SyncValsverifier → artifact → classifier → verdict
SyncVals · Trajectory

reverse-engineer-decoding

claude-code claude-opus-4-8 ✗ failed GOOD_FAILURE ↑ View task
Solved from the instruction alone, tests/ and solution/ were withheld from the agent's workspace and restored only for grading.
Reward = tests/test.sh exit code (0 → resolved); the classification below is post-hoc and cannot change it.
Classification , post-hoc; cannot change the reward
GOOD_FAILUREHonest miss, the agent ran correctly but couldn't solve it. Expected for a hard task; the task is sound.
SubtypeIncomplete Exploration / Timeout
EvidenceTest output: 'FUNCTIONS PASSED: 35/37 (partial_score=0.9459)'. Agent probed reforacle systematically across 600+ test cases and achieved 100% match on all inputs they tested locally (two stress harnesses confirmed 'ALL PASS'). However, grading on hidden held-out inputs revealed 2/37 functions fail on inputs the agent never probed. The instruction explicitly warns: 'Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on.' Agent's implementation was standalone (no reforacle dependency), complete (all 37 functions), and internally consistent. The failure is due to incomplete coverage of the oracle's behavior space, which is fundamentally a hard reverse-engineering task.
Root causeThe agent successfully reverse-engineered 35/37 functions by systematically probing reforacle, discovering non-obvious conventions (last-max tie-break, per-occurrence repetition penalty, etc.). However, with finite time and a sealed oracle that must be probed blindly, achieving 100% match on unseen held-out test cases is extremely difficult. Two functions failed on edge cases or inputs the agent's probing harnesses didn't encounter.
RecommendationN/A - task is fine. This is expected behavior for hard reverse-engineering tasks. A 94.59% success rate with only 37 functions to match (where each requires discovering exact numerical behavior) is reasonable given the constraints. The task is well-specified (oracle is ground truth, grading is transparent about hidden test cases). Agent's approach was sound; this is agent limitation, not a task problem.
Trajectory
Tool-by-tool agent trajectory
27 tool calls · 3 tool types · 29 steps
Reimplement the post-training decoding-and-collation library in `/app/objectives.py` so that every function reproduces the sealed reference module `reforacle` exactly. `reforacle` is importable and callable (`import reforacle`) but its source cannot be read. It is the only ground truth; there is no written specification of what any function computes or of the conventions it follows. Recover its behavior by calling it. `/app/objectives.py` contains thirty-seven function stubs , logit processors and sampling filters, sequence collation and masking utilities, and rollout helpers used to generate and batch data for RLHF and SFT. Keep every signature exactly as given and implement every function. The names indicate the standard quantity each is meant to produce, but the reference's exact definition is whatever `reforacle` returns , not what you assume. Determine each definition by probing. Grading compares your implementation to `reforacle` on held-out inputs you do not see, function by function, requiring exact numerical agreement (primitives within `atol=1e-6, rtol=1e-5`; composite functions within `atol=1e-5, rtol=1e-4`). Each function passes only if it agrees on every graded input, and the task passes only if all thirty-seven functions pass. Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on. Grading runs your `/app/objectives.py` in a clean sandbox that does not contain `reforacle`, so any submission that imports, wraps, calls, loads, or otherwise depends on the reference at run time fails outright. Your file must be a standalone implementation whose only dependency is `torch`. Do not read, decompile, disassemble, or otherwise extract the compiled reference; recover its behavior only by calling it and comparing. `python3 /app/probe.py` shows the calling pattern; `bash /app/timer.sh` reports remaining budget.

/app/objectives.py

contents
1	from __future__ import annotations
2	
3	import torch
4	
5	
6	def greedy_token(logits):
7	    raise NotImplementedError
8	
9	
10	def temperature_scale(logits, temperature):
11	    raise NotImplementedError
12	
13	
14	def top_k_filter(logits, k):
15	    raise NotImplementedError
16	
17	
18	def top_p_filter(logits, p):
19	    raise NotImplementedError
20	
21	
22	def min_p_filter(logits, p):
23	    raise NotImplementedError
24	
25	
26	def repetition_penalty(logits, input_ids, penalty):
27	    raise NotImplementedError
28	
29	
30	def frequency_penalty(logits, input_ids, penalty):
31	    raise NotImplementedError
32	
33	
34	def presence_penalty(logits, input_ids, penalty):
35	    raise NotImplementedError
36	
37	
38	def suppress_tokens(logits, token_ids):
39	    raise NotImplementedError
40	
41	
42	def min_length_filter(logits, cur_len, min_len, eos_id):
43	    raise NotImplementedError
44	
45	
46	def renormalize_logits(logits):
47	    raise NotImplementedError
48	
49	
50	def attention_mask(input_ids, pad_id):
51	    raise NotImplementedError
52	
53	
54	def position_ids(input_ids, pad_id):
55	    raise NotImplementedError
56	
57	
58	def causal_mask(t):
59	    raise NotImplementedError
60	
61	
62	def shift_labels(input_ids, ignore_index):
63	    raise NotImplementedError
64	
65	
66	def completion_loss_labels(input_ids, prompt_lengths, ignore_index):
67	    raise NotImplementedError
68	
69	
70	def mask_after_eos(input_ids, eos_id):
71	    raise NotImplementedError
72	
73	
74	def sequence_lengths(input_ids, pad_id):
75	    raise NotImplementedError
76	
77	
78	def last_token_indices(attn_mask):
79	    raise NotImplementedError
80	
81	
82	def gather_last_logits(logits, attn_mask):
83	    raise NotImplementedError
84	
85	
86	def truncate_left(input_ids, max_len):
87	    raise NotImplementedError
88	
89	
90	def pad_to_multiple(input_ids, pad_id, multiple):
91	    raise NotImplementedError
92	
93	
94	def action_logprobs(logits, actions):
95	    raise NotImplementedError
96	
97	
98	def entropy_from_logits(logits):
99	    raise NotImplementedError
100	
101	
102	def reward_to_go(rewards, mask):
103	    raise NotImplementedError
104	
105	
106	def discount_cumsum(x, gamma):
107	    raise NotImplementedError
108	
109	
110	def build_position_ids_packed(seq_lengths, total_len):
111	    raise NotImplementedError
112	
113	
114	def logits_warper(logits, temperature, top_k, top_p):
115	    raise NotImplementedError
116	
117	
118	def sample_greedy_warped(logits, temperature, top_k, top_p):
119	    raise NotImplementedError
120	
121	
122	def prepare_labels(input_ids, prompt_lengths, ignore_index):
123	    raise NotImplementedError
124	
125	
126	def rollout_last_logits(logits, input_ids, pad_id):
127	    raise NotImplementedError
128	
129	
130	def epsilon_filter(logits, epsilon):
131	    raise NotImplementedError
132	
133	
134	def first_stop_index(input_ids, stop_ids):
135	    raise NotImplementedError
136	
137	
138	def masked_argmax(logits, mask):
139	    raise NotImplementedError
140	
141	
142	def topk_indices(logits, k):
143	    raise NotImplementedError
144	
145	
146	def right_pad_to(input_ids, pad_id, width):
147	    raise NotImplementedError
148	
149	
150	def generation_stop_mask(input_ids, eos_id, min_len):
151	    raise NotImplementedError
152

1	from __future__ import annotations
2	
3	import torch
4	
5	import reforacle
6	
7	g = torch.Generator().manual_seed(0)
8	logits = torch.randn(2, 6, generator=g)
9	
10	print("logits:", logits.tolist())
11	print("reference greedy_token ->", reforacle.greedy_token(logits).tolist())
12	print("reference temperature_scale(T=0) ->", reforacle.temperature_scale(logits, 0.0).tolist())
13	print()
14	print("Probe any reference function this way to discover its exact behavior,")
15	print("then implement it in /app/objectives.py to match.")
16

Remaining time (hours:minutes):
1:59

[stdout]
Remaining time (hours:minutes):
1:59

/app/probe.py

contents
1	from __future__ import annotations
2	
3	import torch
4	
5	
6	def greedy_token(logits):
7	    raise NotImplementedError
8	
9	
10	def temperature_scale(logits, temperature):
11	    raise NotImplementedError
12	
13	
14	def top_k_filter(logits, k):
15	    raise NotImplementedError
16	
17	
18	def top_p_filter(logits, p):
19	    raise NotImplementedError
20	
21	
22	def min_p_filter(logits, p):
23	    raise NotImplementedError
24	
25	
26	def repetition_penalty(logits, input_ids, penalty):
27	    raise NotImplementedError
28	
29	
30	def frequency_penalty(logits, input_ids, penalty):
31	    raise NotImplementedError
32	
33	
34	def presence_penalty(logits, input_ids, penalty):
35	    raise NotImplementedError
36	
37	
38	def suppress_tokens(logits, token_ids):
39	    raise NotImplementedError
40	
41	
42	def min_length_filter(logits, cur_len, min_len, eos_id):
43	    raise NotImplementedError
44	
45	
46	def renormalize_logits(logits):
47	    raise NotImplementedError
48	
49	
50	def attention_mask(input_ids, pad_id):
51	    raise NotImplementedError
52	
53	
54	def position_ids(input_ids, pad_id):
55	    raise NotImplementedError
56	
57	
58	def causal_mask(t):
59	    raise NotImplementedError
60	
61	
62	def shift_labels(input_ids, ignore_index):
63	    raise NotImplementedError
64	
65	
66	def completion_loss_labels(input_ids, prompt_lengths, ignore_index):
67	    raise NotImplementedError
68	
69	
70	def mask_after_eos(input_ids, eos_id):
71	    raise NotImplementedError
72	
73	
74	def sequence_lengths(input_ids, pad_id):
75	    raise NotImplementedError
76	
77	
78	def last_token_indices(attn_mask):
79	    raise NotImplementedError
80	
81	
82	def gather_last_logits(logits, attn_mask):
83	    raise NotImplementedError
84	
85	
86	def truncate_left(input_ids, max_len):
87	    raise NotImplementedError
88	
89	
90	def pad_to_multiple(input_ids, pad_id, multiple):
91	    raise NotImplementedError
92	
93	
94	def action_logprobs(logits, actions):
95	    raise NotImplementedError
96	
97	
98	def entropy_from_logits(logits):
99	    raise NotImplementedError
100	
101	
102	def reward_to_go(rewards, mask):
103	    raise NotImplementedError
104	
105	
106	def discount_cumsum(x, gamma):
107	    raise NotImplementedError
108	
109	
110	def build_position_ids_packed(seq_lengths, total_len):
111	    raise NotImplementedError
112	
113	
114	def logits_warper(logits, temperature, top_k, top_p):
115	    raise NotImplementedError
116	
117	
118	def sample_greedy_warped(logits, temperature, top_k, top_p):
119	    raise NotImplementedError
120	
121	
122	def prepare_labels(input_ids, prompt_lengths, ignore_index):
123	    raise NotImplementedError
124	
125	
126	def rollout_last_logits(logits, input_ids, pad_id):
127	    raise NotImplementedError
128	
129	
130	def epsilon_filter(logits, epsilon):
131	    raise NotImplementedError
132	
133	
134	def first_stop_index(input_ids, stop_ids):
135	    raise NotImplementedError
136	
137	
138	def masked_argmax(logits, mask):
139	    raise NotImplementedError
140	
141	
142	def topk_indices(logits, k):
143	    raise NotImplementedError
144	
145	
146	def right_pad_to(input_ids, pad_id, width):
147	    raise NotImplementedError
148	
149	
150	def generation_stop_mask(input_ids, eos_id, min_len):
151	    raise NotImplementedError
152

1	from __future__ import annotations
2	
3	import torch
4	
5	import reforacle
6	
7	g = torch.Generator().manual_seed(0)
8	logits = torch.randn(2, 6, generator=g)
9	
10	print("logits:", logits.tolist())
11	print("reference greedy_token ->", reforacle.greedy_token(logits).tolist())
12	print("reference temperature_scale(T=0) ->", reforacle.temperature_scale(logits, 0.0).tolist())
13	print()
14	print("Probe any reference function this way to discover its exact behavior,")
15	print("then implement it in /app/objectives.py to match.")
16

Remaining time (hours:minutes):
1:59

[stdout]
Remaining time (hours:minutes):
1:59

Check remaining budget

shell
bash /app/timer.sh
output
1	from __future__ import annotations
2	
3	import torch
4	
5	
6	def greedy_token(logits):
7	    raise NotImplementedError
8	
9	
10	def temperature_scale(logits, temperature):
11	    raise NotImplementedError
12	
13	
14	def top_k_filter(logits, k):
15	    raise NotImplementedError
16	
17	
18	def top_p_filter(logits, p):
19	    raise NotImplementedError
20	
21	
22	def min_p_filter(logits, p):
23	    raise NotImplementedError
24	
25	
26	def repetition_penalty(logits, input_ids, penalty):
27	    raise NotImplementedError
28	
29	
30	def frequency_penalty(logits, input_ids, penalty):
31	    raise NotImplementedError
32	
33	
34	def presence_penalty(logits, input_ids, penalty):
35	    raise NotImplementedError
36	
37	
38	def suppress_tokens(logits, token_ids):
39	    raise NotImplementedError
40	
41	
42	def min_length_filter(logits, cur_len, min_len, eos_id):
43	    raise NotImplementedError
44	
45	
46	def renormalize_logits(logits):
47	    raise NotImplementedError
48	
49	
50	def attention_mask(input_ids, pad_id):
51	    raise NotImplementedError
52	
53	
54	def position_ids(input_ids, pad_id):
55	    raise NotImplementedError
56	
57	
58	def causal_mask(t):
59	    raise NotImplementedError
60	
61	
62	def shift_labels(input_ids, ignore_index):
63	    raise NotImplementedError
64	
65	
66	def completion_loss_labels(input_ids, prompt_lengths, ignore_index):
67	    raise NotImplementedError
68	
69	
70	def mask_after_eos(input_ids, eos_id):
71	    raise NotImplementedError
72	
73	
74	def sequence_lengths(input_ids, pad_id):
75	    raise NotImplementedError
76	
77	
78	def last_token_indices(attn_mask):
79	    raise NotImplementedError
80	
81	
82	def gather_last_logits(logits, attn_mask):
83	    raise NotImplementedError
84	
85	
86	def truncate_left(input_ids, max_len):
87	    raise NotImplementedError
88	
89	
90	def pad_to_multiple(input_ids, pad_id, multiple):
91	    raise NotImplementedError
92	
93	
94	def action_logprobs(logits, actions):
95	    raise NotImplementedError
96	
97	
98	def entropy_from_logits(logits):
99	    raise NotImplementedError
100	
101	
102	def reward_to_go(rewards, mask):
103	    raise NotImplementedError
104	
105	
106	def discount_cumsum(x, gamma):
107	    raise NotImplementedError
108	
109	
110	def build_position_ids_packed(seq_lengths, total_len):
111	    raise NotImplementedError
112	
113	
114	def logits_warper(logits, temperature, top_k, top_p):
115	    raise NotImplementedError
116	
117	
118	def sample_greedy_warped(logits, temperature, top_k, top_p):
119	    raise NotImplementedError
120	
121	
122	def prepare_labels(input_ids, prompt_lengths, ignore_index):
123	    raise NotImplementedError
124	
125	
126	def rollout_last_logits(logits, input_ids, pad_id):
127	    raise NotImplementedError
128	
129	
130	def epsilon_filter(logits, epsilon):
131	    raise NotImplementedError
132	
133	
134	def first_stop_index(input_ids, stop_ids):
135	    raise NotImplementedError
136	
137	
138	def masked_argmax(logits, mask):
139	    raise NotImplementedError
140	
141	
142	def topk_indices(logits, k):
143	    raise NotImplementedError
144	
145	
146	def right_pad_to(input_ids, pad_id, width):
147	    raise NotImplementedError
148	
149	
150	def generation_stop_mask(input_ids, eos_id, min_len):
151	    raise NotImplementedError
152

1	from __future__ import annotations
2	
3	import torch
4	
5	import reforacle
6	
7	g = torch.Generator().manual_seed(0)
8	logits = torch.randn(2, 6, generator=g)
9	
10	print("logits:", logits.tolist())
11	print("reference greedy_token ->", reforacle.greedy_token(logits).tolist())
12	print("reference temperature_scale(T=0) ->", reforacle.temperature_scale(logits, 0.0).tolist())
13	print()
14	print("Probe any reference function this way to discover its exact behavior,")
15	print("then implement it in /app/objectives.py to match.")
16

Remaining time (hours:minutes):
1:59

[stdout]
Remaining time (hours:minutes):
1:59

Run the probe script

shell
python3 probe.py
output
logits: [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
reference greedy_token -> [0, 1]
reference temperature_scale(T=0) -> [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]]

Probe any reference function this way to discover its exact behavior,
then implement it in /app/objectives.py to match.

[stdout]
logits: [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
reference greedy_token -> [0, 1]
reference temperature_scale(T=0) -> [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]]

Probe any reference function this way to discover its exact behavior,
then implement it in /app/objectives.py to match.

Probe batch 1

shell
cat > /tmp/p1.py << 'EOF'
import torch, reforacle
g = torch.Generator().manual_seed(0)

def show(name, *a, **k):
    try:
        r = getattr(reforacle, name)(*a, **k)
        if torch.is_tensor(r):
            print(name, "->", r.dtype, r.shape, r.tolist())
        else:
            print(name, "->", type(r), r)
    except Exception as e:
        print(name, "ERR", repr(e))

logits = torch.randn(2, 6, generator=g)
print("logits", logits.tolist())

# greedy_token
show("greedy_token", logits)
show("greedy_token", torch.randn(3,4,generator=g))  # 2d
# 1d?
show("greedy_token", torch.tensor([1.0,3.0,2.0]))
# 3d
show("greedy_token", torch.randn(2,3,4,generator=g))

# temperature_scale
show("temperature_scale", logits, 2.0)
show("temperature_scale", logits, 0.5)
show("temperature_scale", logits, 1.0)
# tie in T=0
show("temperature_scale", torch.tensor([[1.0,1.0,0.0]]), 0.0)

# top_k_filter
show("top_k_filter", logits, 2)
show("top_k_filter", logits, 1)
show("top_k_filter", logits, 6)
EOF
python3 /tmp/p1.py
output
Exit code 1
Traceback (most recent call last):
  File "/tmp/p1.py", line 1, in <module>
    import torch, reforacle
ModuleNotFoundError: No module named 'reforacle'

[error] tool reported failure

Probe batch 1 from /app

shell
python3 /tmp/p1.py
output
Exit code 1
Traceback (most recent call last):
  File "/tmp/p1.py", line 1, in <module>
    import torch, reforacle
ModuleNotFoundError: No module named 'reforacle'

[error] tool reported failure

List app dir and check path

shell
ls -la /app; python3 -c "import sys; print(sys.path)"
output
total 224
drwxr-xr-x 1 root root   4096 Jul  1 08:20 .
drwxr-xr-x 1 root root   4096 Jul  1 08:23 ..
drwxr-xr-x 2 root root   4096 Jul  1 08:19 __pycache__
-rw-r--r-- 1 root root   2785 Jul  1 08:19 objectives.py
-rw-r--r-- 1 root root    498 Jul  1 08:19 probe.py
-rwxr-xr-x 1 root root 201968 Jul  1 08:20 reforacle.cpython-312-aarch64-linux-gnu.so
-rwxr-xr-x 1 root root    557 Jul  1 08:19 timer.sh
['', '/usr/lib/python312.zip', '/usr/lib/python3.12', '/usr/lib/python3.12/lib-dynload', '/usr/local/lib/python3.12/dist-packages', '/usr/lib/python3/dist-packages']

[stdout]
total 224
drwxr-xr-x 1 root root   4096 Jul  1 08:20 .
drwxr-xr-x 1 root root   4096 Jul  1 08:23 ..
drwxr-xr-x 2 root root   4096 Jul  1 08:19 __pycache__
-rw-r--r-- 1 root root   2785 Jul  1 08:19 objectives.py
-rw-r--r-- 1 root root    498 Jul  1 08:19 probe.py
-rwxr-xr-x 1 root root 201968 Jul  1 08:20 reforacle.cpython-312-aarch64-linux-gnu.so
-rwxr-xr-x 1 root root    557 Jul  1 08:19 timer.sh
['', '/usr/lib/python312.zip', '/usr/lib/python3.12', '/usr/lib/python3.12/lib-dynload', '/usr/local/lib/python3.12/dist-packages', '/usr/lib/python3/dist-packages']

Probe batch 1 with PYTHONPATH

shell
PYTHONPATH=/app python3 /tmp/p1.py
output
logits [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
greedy_token -> torch.int64 torch.Size([2]) [0, 1]
greedy_token -> torch.int64 torch.Size([3]) [1, 1, 1]
greedy_token -> torch.int64 torch.Size([]) 1
greedy_token -> torch.int64 torch.Size([2, 3]) [[2, 2, 3], [3, 0, 3]]
temperature_scale -> torch.float32 torch.Size([2, 6]) [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
temperature_scale -> torch.float32 torch.Size([2, 6]) [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854, -2.7971909046173096], [0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828, -1.1932706832885742, 0.3640729784965515]]
temperature_scale -> torch.float32 torch.Size([2, 6]) [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
temperature_scale -> torch.float32 torch.Size([1, 3]) [[-inf, 0.0, -inf]]
top_k_filter -> torch.float32 torch.Size([2, 6]) [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k_filter -> torch.float32 torch.Size([2, 6]) [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k_filter -> torch.float32 torch.Size([2, 6]) [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]

[stdout]
logits [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
greedy_token -> torch.int64 torch.Size([2]) [0, 1]
greedy_token -> torch.int64 torch.Size([3]) [1, 1, 1]
greedy_token -> torch.int64 torch.Size([]) 1
greedy_token -> torch.int64 torch.Size([2, 3]) [[2, 2, 3], [3, 0, 3]]
temperature_scale -> torch.float32 torch.Size([2, 6]) [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
temperature_scale -> torch.float32 torch.Size([2, 6]) [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854, -2.7971909046173096], [0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828, -1.1932706832885742, 0.3640729784965515]]
temperature_scale -> torch.float32 torch.Size([2, 6]) [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
temperature_scale -> torch.float32 torch.Size([1, 3]) [[-inf, 0.0, -inf]]
top_k_filter -> torch.float32 torch.Size([2, 6]) [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k_filter -> torch.float32 torch.Size([2, 6]) [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k_filter -> torch.float32 torch.Size([2, 6]) [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]

Check argmax tie behavior

shell
PYTHONPATH=/app python3 -c "
import torch
print(torch.tensor([1.0,1.0,0.0]).argmax(-1))
print(torch.tensor([[1.0,1.0,0.0]]).argmax(-1))
print(torch.tensor([[1.,1.,0.],[0.,2.,2.]]).max(-1))
"
output
tensor(0)
tensor([0])
torch.return_types.max(
values=tensor([1., 2.]),
indices=tensor([0, 1]))

[stdout]
tensor(0)
tensor([0])
torch.return_types.max(
values=tensor([1., 2.]),
indices=tensor([0, 1]))

Probe T=0 tie-break

shell
PYTHONPATH=/app python3 -c "
import torch, reforacle
for v in [[1.,1.,0.],[0.,1.,1.],[1.,1.,1.],[2.,0.,2.,1.,2.]]:
    t=torch.tensor([v])
    print(v, reforacle.temperature_scale(t,0.0).tolist())
"
output
[1.0, 1.0, 0.0] [[-inf, 0.0, -inf]]
[0.0, 1.0, 1.0] [[-inf, -inf, 0.0]]
[1.0, 1.0, 1.0] [[-inf, -inf, 0.0]]
[2.0, 0.0, 2.0, 1.0, 2.0] [[-inf, -inf, -inf, -inf, 0.0]]

[stdout]
[1.0, 1.0, 0.0] [[-inf, 0.0, -inf]]
[0.0, 1.0, 1.0] [[-inf, -inf, 0.0]]
[1.0, 1.0, 1.0] [[-inf, -inf, 0.0]]
[2.0, 0.0, 2.0, 1.0, 2.0] [[-inf, -inf, -inf, -inf, 0.0]]

Probe top_p min_p epsilon

shell
cat > /tmp/p2.py << 'EOF'
import torch, reforacle
g = torch.Generator().manual_seed(1)
def show(name, *a, **k):
    try:
        r = getattr(reforacle, name)(*a, **k)
        if torch.is_tensor(r):
            print(name, "->", r.dtype, list(r.shape), r.tolist())
        else:
            print(name, "->", r)
    except Exception as e:
        print(name, "ERR", repr(e))

logits = torch.tensor([[2.0,1.0,0.5,0.2,-1.0,-3.0],[0.1,0.2,0.3,0.4,0.5,0.6]])
print("logits", logits.tolist())
# softmax probs for reference
print("probs row0", torch.softmax(logits[0],-1).tolist())
print("probs row1", torch.softmax(logits[1],-1).tolist())
show("top_p_filter", logits, 0.9)
show("top_p_filter", logits, 0.5)
show("top_p_filter", logits, 1.0)
show("top_p_filter", logits, 0.0)
show("min_p_filter", logits, 0.1)
show("min_p_filter", logits, 0.5)
show("min_p_filter", logits, 0.0)
show("epsilon_filter", logits, 0.05)
show("epsilon_filter", logits, 0.2)
EOF
PYTHONPATH=/app python3 /tmp/p2.py
output
logits [[2.0, 1.0, 0.5, 0.20000000298023224, -1.0, -3.0], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
probs row0 [0.5516226291656494, 0.20293062925338745, 0.12308365106582642, 0.09118261188268661, 0.027463672682642937, 0.003716804087162018]
probs row1 [0.12792666256427765, 0.14138083159923553, 0.1562499850988388, 0.17268292605876923, 0.19084414839744568, 0.21091541647911072]
top_p_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -inf, -inf], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
top_p_filter -> torch.float32 [2, 6] [[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 0.4000000059604645, 0.5, 0.6000000238418579]]
top_p_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -1.0, -3.0], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
top_p_filter -> torch.float32 [2, 6] [[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 0.6000000238418579]]
min_p_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -inf, -inf], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
min_p_filter -> torch.float32 [2, 6] [[2.0, -inf, -inf, -inf, -inf, -inf], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
min_p_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -1.0, -3.0], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
epsilon_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -inf, -inf], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
epsilon_filter -> torch.float32 [2, 6] [[2.0, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 0.6000000238418579]]

[stdout]
logits [[2.0, 1.0, 0.5, 0.20000000298023224, -1.0, -3.0], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
probs row0 [0.5516226291656494, 0.20293062925338745, 0.12308365106582642, 0.09118261188268661, 0.027463672682642937, 0.003716804087162018]
probs row1 [0.12792666256427765, 0.14138083159923553, 0.1562499850988388, 0.17268292605876923, 0.19084414839744568, 0.21091541647911072]
top_p_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -inf, -inf], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
top_p_filter -> torch.float32 [2, 6] [[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 0.4000000059604645, 0.5, 0.6000000238418579]]
top_p_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -1.0, -3.0], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
top_p_filter -> torch.float32 [2, 6] [[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 0.6000000238418579]]
min_p_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -inf, -inf], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
min_p_filter -> torch.float32 [2, 6] [[2.0, -inf, -inf, -inf, -inf, -inf], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
min_p_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -1.0, -3.0], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
epsilon_filter -> torch.float32 [2, 6] [[2.0, 1.0, 0.5, 0.20000000298023224, -inf, -inf], [0.10000000149011612, 0.20000000298023224, 0.30000001192092896, 0.4000000059604645, 0.5, 0.6000000238418579]]
epsilon_filter -> torch.float32 [2, 6] [[2.0, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 0.6000000238418579]]

Probe penalties, suppress, safety

shell
cat > /tmp/p3.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    try:
        r = getattr(reforacle, name)(*a, **k)
        print(name, "->", r.tolist() if torch.is_tensor(r) else r)
    except Exception as e:
        print(name, "ERR", repr(e))

logits = torch.tensor([[2.0,1.0,0.5,0.2,-1.0,-3.0]])
# safety: all removed cases
show("top_p_filter", logits, 0.0)
show("min_p_filter", logits, 1.5)
show("epsilon_filter", logits, 0.9)
show("epsilon_filter", logits, 0.99)
show("top_k_filter", logits, 0)
# penalties
ll = torch.tensor([[1.0,1.0,1.0,1.0,1.0]])
ids = torch.tensor([[0,0,2]])
show("repetition_penalty", torch.tensor([[1.0,-1.0,2.0,0.5,0.0]]), ids, 2.0)
show("frequency_penalty", torch.tensor([[1.0,1.0,1.0,1.0,1.0]]), ids, 0.5)
show("presence_penalty", torch.tensor([[1.0,1.0,1.0,1.0,1.0]]), ids, 0.5)
show("suppress_tokens", torch.tensor([[1.0,2.0,3.0,4.0]]), [1,3])
show("suppress_tokens", torch.tensor([[1.0,2.0,3.0,4.0]]), torch.tensor([0,2]))
EOF
PYTHONPATH=/app python3 /tmp/p3.py
output
top_p_filter -> [[2.0, -inf, -inf, -inf, -inf, -inf]]
min_p_filter -> [[-inf, -inf, -inf, -inf, -inf, -inf]]
epsilon_filter -> [[2.0, -inf, -inf, -inf, -inf, -inf]]
epsilon_filter -> [[2.0, -inf, -inf, -inf, -inf, -inf]]
top_k_filter ERR RuntimeError('The size of tensor a (6) must match the size of tensor b (0) at non-singleton dimension 1')
repetition_penalty -> [[0.25, -1.0, 1.0, 0.5, 0.0]]
frequency_penalty -> [[0.0, 1.0, 0.5, 1.0, 1.0]]
presence_penalty -> [[0.5, 1.0, 0.5, 1.0, 1.0]]
suppress_tokens -> [[1.0, -inf, 3.0, -inf]]
suppress_tokens -> [[-inf, 2.0, -inf, 4.0]]

[stdout]
top_p_filter -> [[2.0, -inf, -inf, -inf, -inf, -inf]]
min_p_filter -> [[-inf, -inf, -inf, -inf, -inf, -inf]]
epsilon_filter -> [[2.0, -inf, -inf, -inf, -inf, -inf]]
epsilon_filter -> [[2.0, -inf, -inf, -inf, -inf, -inf]]
top_k_filter ERR RuntimeError('The size of tensor a (6) must match the size of tensor b (0) at non-singleton dimension 1')
repetition_penalty -> [[0.25, -1.0, 1.0, 0.5, 0.0]]
frequency_penalty -> [[0.0, 1.0, 0.5, 1.0, 1.0]]
presence_penalty -> [[0.5, 1.0, 0.5, 1.0, 1.0]]
suppress_tokens -> [[1.0, -inf, 3.0, -inf]]
suppress_tokens -> [[-inf, 2.0, -inf, 4.0]]

Probe repetition penalty details

shell
cat > /tmp/p4.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    r = getattr(reforacle, name)(*a, **k)
    print(name, "->", r.tolist() if torch.is_tensor(r) else r)

# rep penalty with negatives and counts
lg = torch.tensor([[2.0,-2.0,1.0,-1.0]])
show("repetition_penalty", lg, torch.tensor([[0,0,1,3,3,3]]), 2.0)
# token0 pos logit 2.0 count2 ; token1 neg -2 count1; token3 neg -1 count3
# single occurrence check
show("repetition_penalty", torch.tensor([[2.0,-2.0]]), torch.tensor([[0,1]]), 2.0)
# batch with different rows
lg2 = torch.tensor([[1.0,2.0,3.0],[3.0,2.0,1.0]])
show("repetition_penalty", lg2, torch.tensor([[0,0],[2,2]]), 3.0)
show("frequency_penalty", lg2, torch.tensor([[0,0],[2,2]]), 1.0)
show("presence_penalty", lg2, torch.tensor([[0,1],[2,2]]), 1.0)
EOF
PYTHONPATH=/app python3 /tmp/p4.py
output
repetition_penalty -> [[0.5, -4.0, 1.0, -8.0]]
repetition_penalty -> [[1.0, -4.0]]
repetition_penalty -> [[0.1111111119389534, 2.0, 3.0], [3.0, 2.0, 0.1111111119389534]]
frequency_penalty -> [[-1.0, 2.0, 3.0], [3.0, 2.0, -1.0]]
presence_penalty -> [[0.0, 1.0, 3.0], [3.0, 2.0, 0.0]]

[stdout]
repetition_penalty -> [[0.5, -4.0, 1.0, -8.0]]
repetition_penalty -> [[1.0, -4.0]]
repetition_penalty -> [[0.1111111119389534, 2.0, 3.0], [3.0, 2.0, 0.1111111119389534]]
frequency_penalty -> [[-1.0, 2.0, 3.0], [3.0, 2.0, -1.0]]
presence_penalty -> [[0.0, 1.0, 3.0], [3.0, 2.0, 0.0]]

Probe masking/collation group

shell
cat > /tmp/p5.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    try:
        r = getattr(reforacle, name)(*a, **k)
        if torch.is_tensor(r): print(name, "->", r.dtype, list(r.shape), r.tolist())
        else: print(name,"->", r)
    except Exception as e:
        print(name, "ERR", repr(e))

lg = torch.tensor([[1.0,2.0,3.0,4.0]])
show("min_length_filter", lg, 3, 5, 2)   # cur_len<min -> suppress eos
show("min_length_filter", lg, 5, 5, 2)   # cur_len==min
show("min_length_filter", lg, 6, 5, 2)
show("renormalize_logits", torch.tensor([[1.0,2.0,3.0],[0.0,0.0,0.0]]))
ids = torch.tensor([[5,6,7,0,0],[1,2,0,0,0]])
show("attention_mask", ids, 0)
show("position_ids", ids, 0)
show("causal_mask", 4)
show("sequence_lengths", ids, 0)
show("last_token_indices", torch.tensor([[1,1,1,0,0],[1,1,0,0,0]]))
lg3 = torch.arange(2*5*3).float().reshape(2,5,3)
show("gather_last_logits", lg3, torch.tensor([[1,1,1,0,0],[1,1,0,0,0]]))
EOF
PYTHONPATH=/app python3 /tmp/p5.py
output
min_length_filter -> torch.float32 [1, 4] [[1.0, 2.0, -inf, 4.0]]
min_length_filter -> torch.float32 [1, 4] [[1.0, 2.0, 3.0, 4.0]]
min_length_filter -> torch.float32 [1, 4] [[1.0, 2.0, 3.0, 4.0]]
renormalize_logits -> torch.float32 [2, 3] [[-2.4076058864593506, -1.4076058864593506, -0.40760594606399536], [-1.0986123085021973, -1.0986123085021973, -1.0986123085021973]]
attention_mask -> torch.float32 [2, 5] [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]]
position_ids -> torch.int64 [2, 5] [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1]]
causal_mask -> torch.float32 [4, 4] [[0.0, -inf, -inf, -inf], [0.0, 0.0, -inf, -inf], [0.0, 0.0, 0.0, -inf], [0.0, 0.0, 0.0, 0.0]]
sequence_lengths -> torch.int64 [2] [3, 2]
last_token_indices -> torch.int64 [2] [2, 1]
gather_last_logits -> torch.float32 [2, 3] [[6.0, 7.0, 8.0], [18.0, 19.0, 20.0]]

[stdout]
min_length_filter -> torch.float32 [1, 4] [[1.0, 2.0, -inf, 4.0]]
min_length_filter -> torch.float32 [1, 4] [[1.0, 2.0, 3.0, 4.0]]
min_length_filter -> torch.float32 [1, 4] [[1.0, 2.0, 3.0, 4.0]]
renormalize_logits -> torch.float32 [2, 3] [[-2.4076058864593506, -1.4076058864593506, -0.40760594606399536], [-1.0986123085021973, -1.0986123085021973, -1.0986123085021973]]
attention_mask -> torch.float32 [2, 5] [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]]
position_ids -> torch.int64 [2, 5] [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1]]
causal_mask -> torch.float32 [4, 4] [[0.0, -inf, -inf, -inf], [0.0, 0.0, -inf, -inf], [0.0, 0.0, 0.0, -inf], [0.0, 0.0, 0.0, 0.0]]
sequence_lengths -> torch.int64 [2] [3, 2]
last_token_indices -> torch.int64 [2] [2, 1]
gather_last_logits -> torch.float32 [2, 3] [[6.0, 7.0, 8.0], [18.0, 19.0, 20.0]]

Probe shift/labels/mask

shell
cat > /tmp/p6.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    try:
        r = getattr(reforacle, name)(*a, **k)
        if torch.is_tensor(r): print(name, "->", r.dtype, list(r.shape), r.tolist())
        else: print(name,"->", r)
    except Exception as e:
        print(name, "ERR", repr(e))

# leading pad
show("position_ids", torch.tensor([[0,0,5,6]]), 0)
show("last_token_indices", torch.tensor([[1,0,1,0]]))  # gap
show("sequence_lengths", torch.tensor([[0,5,0,6]]), 0)
show("attention_mask", torch.tensor([[0,5,0,6]]), 0)
# shift_labels
show("shift_labels", torch.tensor([[10,11,12,13]]), -100)
show("shift_labels", torch.tensor([[10,11,12,13],[20,21,22,23]]), -100)
# mask_after_eos
show("mask_after_eos", torch.tensor([[1,2,9,3,4],[9,1,2,3,9]]), 9)
# completion_loss_labels
show("completion_loss_labels", torch.tensor([[1,2,3,4,5]]), torch.tensor([2]), -100)
show("completion_loss_labels", torch.tensor([[1,2,3,4,5],[6,7,8,9,10]]), torch.tensor([2,3]), -100)
# prepare_labels
show("prepare_labels", torch.tensor([[1,2,3,4,5],[6,7,8,9,10]]), torch.tensor([2,3]), -100)
EOF
PYTHONPATH=/app python3 /tmp/p6.py
output
position_ids -> torch.int64 [1, 4] [[0, 0, 0, 1]]
last_token_indices -> torch.int64 [1] [2]
sequence_lengths -> torch.int64 [1] [2]
attention_mask -> torch.float32 [1, 4] [[0.0, 1.0, 0.0, 1.0]]
shift_labels -> torch.int64 [1, 4] [[11, 12, 13, -100]]
shift_labels -> torch.int64 [2, 4] [[11, 12, 13, -100], [21, 22, 23, -100]]
mask_after_eos -> torch.float32 [2, 5] [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0, 0.0]]
completion_loss_labels -> torch.int64 [1, 5] [[-100, -100, 3, 4, 5]]
completion_loss_labels -> torch.int64 [2, 5] [[-100, -100, 3, 4, 5], [-100, -100, -100, 9, 10]]
prepare_labels -> torch.int64 [2, 5] [[-100, 3, 4, 5, -100], [-100, -100, 9, 10, -100]]

[stdout]
position_ids -> torch.int64 [1, 4] [[0, 0, 0, 1]]
last_token_indices -> torch.int64 [1] [2]
sequence_lengths -> torch.int64 [1] [2]
attention_mask -> torch.float32 [1, 4] [[0.0, 1.0, 0.0, 1.0]]
shift_labels -> torch.int64 [1, 4] [[11, 12, 13, -100]]
shift_labels -> torch.int64 [2, 4] [[11, 12, 13, -100], [21, 22, 23, -100]]
mask_after_eos -> torch.float32 [2, 5] [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0, 0.0]]
completion_loss_labels -> torch.int64 [1, 5] [[-100, -100, 3, 4, 5]]
completion_loss_labels -> torch.int64 [2, 5] [[-100, -100, 3, 4, 5], [-100, -100, -100, 9, 10]]
prepare_labels -> torch.int64 [2, 5] [[-100, 3, 4, 5, -100], [-100, -100, 9, 10, -100]]

Probe truncate/pad/logprobs/entropy

shell
cat > /tmp/p7.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    try:
        r = getattr(reforacle, name)(*a, **k)
        if torch.is_tensor(r): print(name, "->", r.dtype, list(r.shape), r.tolist())
        else: print(name,"->", r)
    except Exception as e:
        print(name, "ERR", repr(e))

ids = torch.tensor([[1,2,3,4,5],[6,7,8,9,10]])
show("truncate_left", ids, 3)
show("truncate_left", ids, 7)
show("pad_to_multiple", torch.tensor([[1,2,3],[4,5,6]]), 0, 4)
show("pad_to_multiple", torch.tensor([[1,2,3,4]]), 0, 4)
show("right_pad_to", torch.tensor([[1,2,3]]), 0, 6)
show("right_pad_to", torch.tensor([[1,2,3]]), 0, 2)
# action_logprobs
lg = torch.randn(2,3,4, generator=torch.Generator().manual_seed(3))
acts = torch.tensor([[0,1,2],[3,2,1]])
show("action_logprobs", lg, acts)
print("check", torch.log_softmax(lg,-1).gather(-1, acts.unsqueeze(-1)).squeeze(-1).tolist())
# entropy
show("entropy_from_logits", torch.tensor([[1.0,2.0,3.0],[0.0,0.0,0.0]]))
p=torch.softmax(torch.tensor([1.0,2.0,3.0]),-1); print("ent0", -(p*p.log()).sum().item())
EOF
PYTHONPATH=/app python3 /tmp/p7.py
output
truncate_left -> torch.int64 [2, 3] [[3, 4, 5], [8, 9, 10]]
truncate_left -> torch.int64 [2, 5] [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]
pad_to_multiple -> torch.int64 [2, 4] [[1, 2, 3, 0], [4, 5, 6, 0]]
pad_to_multiple -> torch.int64 [1, 4] [[1, 2, 3, 4]]
right_pad_to -> torch.int64 [1, 6] [[1, 2, 3, 0, 0, 0]]
right_pad_to -> torch.int64 [1, 2] [[1, 2]]
action_logprobs -> torch.float32 [2, 3] [[-1.4353127479553223, -2.360366106033325, -1.1033122539520264], [-3.039785385131836, -1.7562686204910278, -1.6969051361083984]]
check [[-1.4353127479553223, -2.360366106033325, -1.1033122539520264], [-3.039785385131836, -1.7562686204910278, -1.6969051361083984]]
entropy_from_logits -> torch.float32 [2] [0.832395613193512, 1.0986123085021973]
ent0 0.8323955535888672

[stdout]
truncate_left -> torch.int64 [2, 3] [[3, 4, 5], [8, 9, 10]]
truncate_left -> torch.int64 [2, 5] [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]
pad_to_multiple -> torch.int64 [2, 4] [[1, 2, 3, 0], [4, 5, 6, 0]]
pad_to_multiple -> torch.int64 [1, 4] [[1, 2, 3, 4]]
right_pad_to -> torch.int64 [1, 6] [[1, 2, 3, 0, 0, 0]]
right_pad_to -> torch.int64 [1, 2] [[1, 2]]
action_logprobs -> torch.float32 [2, 3] [[-1.4353127479553223, -2.360366106033325, -1.1033122539520264], [-3.039785385131836, -1.7562686204910278, -1.6969051361083984]]
check [[-1.4353127479553223, -2.360366106033325, -1.1033122539520264], [-3.039785385131836, -1.7562686204910278, -1.6969051361083984]]
entropy_from_logits -> torch.float32 [2] [0.832395613193512, 1.0986123085021973]
ent0 0.8323955535888672

Probe RL helpers

shell
cat > /tmp/p8.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    try:
        r = getattr(reforacle, name)(*a, **k)
        if torch.is_tensor(r): print(name, "->", r.dtype, list(r.shape), r.tolist())
        else: print(name,"->", r)
    except Exception as e:
        print(name, "ERR", repr(e))

# discount_cumsum
show("discount_cumsum", torch.tensor([1.0,2.0,3.0]), 0.9)
show("discount_cumsum", torch.tensor([[1.0,2.0,3.0],[1.0,1.0,1.0]]), 0.5)
# reward_to_go with mask
show("reward_to_go", torch.tensor([[1.0,2.0,3.0,4.0]]), torch.tensor([[1,1,1,0]]))
show("reward_to_go", torch.tensor([[1.0,2.0,3.0,4.0]]), torch.tensor([[1,1,1,1]]))
# build_position_ids_packed
show("build_position_ids_packed", torch.tensor([2,3,1]), 6)
# first_stop_index
show("first_stop_index", torch.tensor([[1,2,3,9,5],[1,2,3,4,5]]), [9,3])
# masked_argmax
show("masked_argmax", torch.tensor([[1.0,5.0,3.0,2.0]]), torch.tensor([[1,0,1,1]]))
# topk_indices
show("topk_indices", torch.tensor([[1.0,5.0,3.0,2.0,4.0]]), 3)
# generation_stop_mask
show("generation_stop_mask", torch.tensor([[1,2,9,3],[9,1,2,3]]), 9, 2)
EOF
PYTHONPATH=/app python3 /tmp/p8.py
output
discount_cumsum -> torch.float32 [3] [5.229999542236328, 4.699999809265137, 3.0]
discount_cumsum -> torch.float32 [2, 3] [[2.75, 3.5, 3.0], [1.75, 1.5, 1.0]]
reward_to_go -> torch.float32 [1, 4] [[6.0, 5.0, 3.0, 0.0]]
reward_to_go -> torch.float32 [1, 4] [[10.0, 9.0, 7.0, 4.0]]
build_position_ids_packed -> torch.int64 [6] [0, 1, 0, 1, 2, 0]
first_stop_index ERR AttributeError("'list' object has no attribute 'tolist'")
masked_argmax -> torch.int64 [1] [2]
topk_indices -> torch.int64 [1, 3] [[1, 4, 2]]
generation_stop_mask -> torch.float32 [2, 4] [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]]

[stdout]
discount_cumsum -> torch.float32 [3] [5.229999542236328, 4.699999809265137, 3.0]
discount_cumsum -> torch.float32 [2, 3] [[2.75, 3.5, 3.0], [1.75, 1.5, 1.0]]
reward_to_go -> torch.float32 [1, 4] [[6.0, 5.0, 3.0, 0.0]]
reward_to_go -> torch.float32 [1, 4] [[10.0, 9.0, 7.0, 4.0]]
build_position_ids_packed -> torch.int64 [6] [0, 1, 0, 1, 2, 0]
first_stop_index ERR AttributeError("'list' object has no attribute 'tolist'")
masked_argmax -> torch.int64 [1] [2]
topk_indices -> torch.int64 [1, 3] [[1, 4, 2]]
generation_stop_mask -> torch.float32 [2, 4] [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]]

Probe edge cases

shell
cat > /tmp/p9.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    try:
        r = getattr(reforacle, name)(*a, **k)
        if torch.is_tensor(r): print(name, "->", r.dtype, list(r.shape), r.tolist())
        else: print(name,"->", r)
    except Exception as e:
        print(name, "ERR", repr(e))

# no eos cases
show("mask_after_eos", torch.tensor([[1,2,3,4]]), 9)
show("generation_stop_mask", torch.tensor([[1,2,3,4],[1,2,3,9]]), 9, 2)
# eos as last / multiple eos
show("mask_after_eos", torch.tensor([[9,9,9]]), 9)
# first_stop_index with tensor stop_ids
show("first_stop_index", torch.tensor([[1,2,3,9,5],[1,2,3,4,5]]), torch.tensor([9,3]))
show("first_stop_index", torch.tensor([[5,3,9,3]]), torch.tensor([9,3]))  # earliest of any stop
show("first_stop_index", torch.tensor([[1,2,3]]), torch.tensor([9]))  # none
# reward_to_go interior mask zero
show("reward_to_go", torch.tensor([[1.0,2.0,3.0]]), torch.tensor([[1,0,1]]))
# masked_argmax tie
show("masked_argmax", torch.tensor([[5.0,5.0,1.0]]), torch.tensor([[1,1,1]]))
EOF
PYTHONPATH=/app python3 /tmp/p9.py
output
mask_after_eos -> torch.float32 [1, 4] [[1.0, 1.0, 1.0, 1.0]]
generation_stop_mask -> torch.float32 [2, 4] [[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]
mask_after_eos -> torch.float32 [1, 3] [[1.0, 0.0, 0.0]]
first_stop_index -> torch.int64 [2] [2, 2]
first_stop_index -> torch.int64 [1] [1]
first_stop_index -> torch.int64 [1] [-1]
reward_to_go -> torch.float32 [1, 3] [[4.0, 0.0, 3.0]]
masked_argmax -> torch.int64 [1] [1]

[stdout]
mask_after_eos -> torch.float32 [1, 4] [[1.0, 1.0, 1.0, 1.0]]
generation_stop_mask -> torch.float32 [2, 4] [[1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0]]
mask_after_eos -> torch.float32 [1, 3] [[1.0, 0.0, 0.0]]
first_stop_index -> torch.int64 [2] [2, 2]
first_stop_index -> torch.int64 [1] [1]
first_stop_index -> torch.int64 [1] [-1]
reward_to_go -> torch.float32 [1, 3] [[4.0, 0.0, 3.0]]
masked_argmax -> torch.int64 [1] [1]

Probe tie-break and warpers

shell
cat > /tmp/p10.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    r = getattr(reforacle, name)(*a, **k)
    print(name, "->", r.tolist() if torch.is_tensor(r) else r)
t=torch.tensor([[5.0,5.0,1.0,5.0]])
show("greedy_token", t)
show("masked_argmax", t, torch.tensor([[1,1,1,1]]))
show("topk_indices", t, 2)
show("topk_indices", torch.tensor([[1.0,3.0,3.0,3.0,2.0]]), 2)
# torch reference
print("torch argmax", t.argmax(-1).tolist())
print("torch topk", torch.topk(torch.tensor([1.0,3.0,3.0,3.0,2.0]),2))
# sample_greedy_warped & logits_warper
lg=torch.tensor([[2.0,1.0,0.5,0.2,-1.0,-3.0]])
show("logits_warper", lg, 1.0, 3, 0.9)
show("logits_warper", lg, 2.0, 0, 1.0)
show("sample_greedy_warped", lg, 1.0, 3, 0.9)
show("sample_greedy_warped", torch.tensor([[2.0,1.0,0.5,0.2,-1.0,-3.0],[0.1,0.2,0.3,5.0,0.5,0.6]]), 0.7, 2, 0.95)
EOF
PYTHONPATH=/app python3 /tmp/p10.py
output
Exit code 1
greedy_token -> [3]
masked_argmax -> [3]
topk_indices -> [[0, 1]]
topk_indices -> [[1, 2]]
torch argmax [0]
torch topk torch.return_types.topk(
values=tensor([3., 3.]),
indices=tensor([1, 3]))
logits_warper -> [[2.0, 1.0, 0.5, -inf, -inf, -inf]]
Traceback (most recent call last):
  File "/tmp/p10.py", line 16, in <module>
    show("logits_warper", lg, 2.0, 0, 1.0)
  File "/tmp/p10.py", line 3, in show
    r = getattr(reforacle, name)(*a, **k)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "reforacle.py", line 211, in reforacle.logits_warper
  File "reforacle.py", line 30, in reforacle.top_k_filter
RuntimeError: The size of tensor a (6) must match the size of tensor b (0) at non-singleton dimension 1

[error] tool reported failure

Confirm warper order and rollout

shell
cat > /tmp/p11.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    r = getattr(reforacle, name)(*a, **k)
    print(name, "->", r.tolist() if torch.is_tensor(r) else r)
lg=torch.tensor([[2.0,1.0,0.5,0.2,-1.0,-3.0]])
show("logits_warper", lg, 2.0, 4, 1.0)   # temp then topk4
print("expect temp:", (lg/2.0).tolist())
show("sample_greedy_warped", lg, 1.0, 3, 0.9)
show("sample_greedy_warped", torch.tensor([[2.0,1.0,0.5,0.2,-1.0,-3.0],[0.1,0.2,0.3,5.0,0.5,0.6]]), 0.7, 2, 0.95)
# order: does top_p come before top_k? test
show("logits_warper", lg, 1.0, 5, 0.5)
# rollout_last_logits
lg3=torch.arange(2*4*3).float().reshape(2,4,3)
show("rollout_last_logits", lg3, torch.tensor([[1,2,0,0],[5,6,7,0]]), 0)
show("last_token_indices", torch.tensor([[1,1,0,0],[1,1,1,0]]))
EOF
PYTHONPATH=/app python3 /tmp/p11.py
output
logits_warper -> [[1.0, 0.5, 0.25, 0.10000000149011612, -inf, -inf]]
expect temp: [[1.0, 0.5, 0.25, 0.10000000149011612, -0.5, -1.5]]
sample_greedy_warped -> [0]
sample_greedy_warped -> [0, 3]
logits_warper -> [[2.0, -inf, -inf, -inf, -inf, -inf]]
rollout_last_logits -> [[3.0, 4.0, 5.0], [18.0, 19.0, 20.0]]
last_token_indices -> [1, 2]

[stdout]
logits_warper -> [[1.0, 0.5, 0.25, 0.10000000149011612, -inf, -inf]]
expect temp: [[1.0, 0.5, 0.25, 0.10000000149011612, -0.5, -1.5]]
sample_greedy_warped -> [0]
sample_greedy_warped -> [0, 3]
logits_warper -> [[2.0, -inf, -inf, -inf, -inf, -inf]]
rollout_last_logits -> [[3.0, 4.0, 5.0], [18.0, 19.0, 20.0]]
last_token_indices -> [1, 2]

More edge probes

shell
cat > /tmp/p12.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    try:
        r = getattr(reforacle, name)(*a, **k)
        if torch.is_tensor(r): print(name, "->", r.dtype, list(r.shape), r.tolist())
        else: print(name,"->", r)
    except Exception as e:
        print(name, "ERR", repr(e))
show("build_position_ids_packed", torch.tensor([2,3]), 8)  # total>sum
show("completion_loss_labels", torch.tensor([[1,2,3]]), torch.tensor([5]), -100) # prompt>seq
show("completion_loss_labels", torch.tensor([[1,2,3]]), torch.tensor([0]), -100)
show("prepare_labels", torch.tensor([[1,2,3]]), torch.tensor([0]), -100)
show("temperature_scale", torch.tensor([[1.0,2.0,3.0]]), 0.0)
# discount_cumsum negative/2d gamma float
show("discount_cumsum", torch.tensor([1.0,0.0,0.0,0.0]), 0.5)
# entropy 3d?
show("entropy_from_logits", torch.randn(2,3,4,generator=torch.Generator().manual_seed(5)))
# action_logprobs 2d (no batch)?
show("action_logprobs", torch.randn(3,4,generator=torch.Generator().manual_seed(6)), torch.tensor([0,1,2]))
# min_length cur==min already tested. min_length with eos_id list?
# pad_to_multiple already exact multiple
show("pad_to_multiple", torch.tensor([[1,2,3,4,5]]), 7, 4)  # 5-> pad to 8
EOF
PYTHONPATH=/app python3 /tmp/p12.py
output
build_position_ids_packed -> torch.int64 [8] [0, 1, 0, 1, 2, 0, 0, 0]
completion_loss_labels -> torch.int64 [1, 3] [[-100, -100, -100]]
completion_loss_labels -> torch.int64 [1, 3] [[1, 2, 3]]
prepare_labels -> torch.int64 [1, 3] [[2, 3, -100]]
temperature_scale -> torch.float32 [1, 3] [[-inf, -inf, 0.0]]
discount_cumsum -> torch.float32 [4] [1.0, 0.0, 0.0, 0.0]
entropy_from_logits -> torch.float32 [2, 3] [[0.7007933855056763, 0.6961748600006104, 0.6611921191215515], [1.283085584640503, 1.0535820722579956, 1.3640565872192383]]
action_logprobs -> torch.float32 [3] [-2.9968271255493164, -1.8515634536743164, -2.4247515201568604]
pad_to_multiple -> torch.int64 [1, 8] [[1, 2, 3, 4, 5, 7, 7, 7]]

[stdout]
build_position_ids_packed -> torch.int64 [8] [0, 1, 0, 1, 2, 0, 0, 0]
completion_loss_labels -> torch.int64 [1, 3] [[-100, -100, -100]]
completion_loss_labels -> torch.int64 [1, 3] [[1, 2, 3]]
prepare_labels -> torch.int64 [1, 3] [[2, 3, -100]]
temperature_scale -> torch.float32 [1, 3] [[-inf, -inf, 0.0]]
discount_cumsum -> torch.float32 [4] [1.0, 0.0, 0.0, 0.0]
entropy_from_logits -> torch.float32 [2, 3] [[0.7007933855056763, 0.6961748600006104, 0.6611921191215515], [1.283085584640503, 1.0535820722579956, 1.3640565872192383]]
action_logprobs -> torch.float32 [3] [-2.9968271255493164, -1.8515634536743164, -2.4247515201568604]
pad_to_multiple -> torch.int64 [1, 8] [[1, 2, 3, 4, 5, 7, 7, 7]]

Verify top_k tie, epsilon tie

shell
cat > /tmp/p13.py << 'EOF'
import torch, reforacle
def show(name, *a, **k):
    r = getattr(reforacle, name)(*a, **k)
    print(name, "->", r.tolist() if torch.is_tensor(r) else r)
# top_k tie threshold
show("top_k_filter", torch.tensor([[5.0,5.0,5.0,1.0]]), 2)  # 3 tied for top2
# epsilon keep-top1 with ties in max
show("epsilon_filter", torch.tensor([[3.0,3.0,0.0]]), 0.9)
# presence with penalty on 3d? skip. frequency scatter multiid
show("frequency_penalty", torch.tensor([[0.0,0.0,0.0,0.0]]), torch.tensor([[1,1,3,3,3]]), 1.0)
EOF
PYTHONPATH=/app python3 /tmp/p13.py
output
top_k_filter -> [[5.0, 5.0, 5.0, -inf]]
epsilon_filter -> [[-inf, 3.0, -inf]]
frequency_penalty -> [[0.0, -2.0, 0.0, -3.0]]

[stdout]
top_k_filter -> [[5.0, 5.0, 5.0, -inf]]
epsilon_filter -> [[-inf, 3.0, -inf]]
frequency_penalty -> [[0.0, -2.0, 0.0, -3.0]]
write · /app/objectives.py
from __future__ import annotations

import torch

NEG_INF = float("-inf")


def _last_argmax(x):
    # Index of the LAST maximum along the last dim (reforacle's tie-break).
    n = x.shape[-1]
    return (n - 1) - x.flip(-1).argmax(dim=-1)


def _counts(logits, input_ids):
    counts = torch.zeros_like(logits)
    counts.scatter_add_(-1, input_ids, torch.ones_like(input_ids, dtype=logits.dtype))
    return counts


def _first_index(hits, fill):
    # For a boolean tensor `hits` [..., L]: first True index along -1, else `fill`.
    has = hits.any(dim=-1)
    idx = hits.to(torch.int64).argmax(dim=-1)
    return torch.where(has, idx, torch.full_like(idx, fill))


def greedy_token(logits):
    return _last_argmax(logits)


def temperature_scale(logits, temperature):
    if temperature == 0:
        idx = _last_argmax(logits)
        out = torch.full_like(logits, NEG_INF)
        out.scatter_(-1, idx.unsqueeze(-1), 0.0)
        return out
    return logits / temperature


def top_k_filter(logits, k):
    k = int(k)
    thresh = logits.topk(k, dim=-1).values[..., -1:]
    return logits.masked_fill(logits < thresh, NEG_INF)


def top_p_filter(logits, p):
    sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1)
    cum = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
    remove_sorted = cum > p
    remove_sorted[..., 1:] = remove_sorted[..., :-1].clone()
    remove_sorted[..., 0] = False
    remove = torch.zeros_like(remove_sorted).scatter_(-1, sorted_idx, remove_sorted)
    return logits.masked_fill(remove, NEG_INF)


def min_p_filter(logits, p):
    probs = logits.softmax(dim=-1)
    max_p = probs.max(dim=-1, keepdim=True).values
    return logits.masked_fill(probs < p * max_p, NEG_INF)


def repetition_penalty(logits, input_ids, penalty):
    factor = penalty ** _counts(logits, input_ids)
    return torch.where(logits > 0, logits / factor, logits * factor)


def frequency_penalty(logits, input_ids, penalty):
    return logits - penalty * _counts(logits, input_ids)


def presence_penalty(logits, input_ids, penalty):
    present = (_counts(logits, input_ids) > 0).to(logits.dtype)
    return logits - penalty * present


def suppress_tokens(logits, token_ids):
    out = logits.clone()
    out[..., token_ids] = NEG_INF
    return out


def min_length_filter(logits, cur_len, min_len, eos_id):
    out = logits.clone()
    if cur_len < min_len:
        out[..., eos_id] = NEG_INF
    return out


def renormalize_logits(logits):
    return torch.log_softmax(logits, dim=-1)


def attention_mask(input_ids, pad_id):
    return (input_ids != pad_id).to(torch.float32)


def position_ids(input_ids, pad_id):
    mask = (input_ids != pad_id).to(torch.int64)
    return (mask.cumsum(dim=-1) - 1).clamp(min=0)


def causal_mask(t):
    m = torch.full((t, t), NEG_INF)
    return torch.triu(m, diagonal=1)


def shift_labels(input_ids, ignore_index):
    labels = torch.full_like(input_ids, ignore_index)
    labels[..., :-1] = input_ids[..., 1:]
    return labels


def completion_loss_labels(input_ids, prompt_lengths, ignore_index):
    labels = input_ids.clone()
    ar = torch.arange(input_ids.shape[-1], device=input_ids.device)
    mask = ar.unsqueeze(0) < prompt_lengths.unsqueeze(-1)
    labels[mask] = ignore_index
    return labels


def mask_after_eos(input_ids, eos_id):
    L = input_ids.shape[-1]
    first_eos = _first_index(input_ids == eos_id, L)
    ar = torch.arange(L, device=input_ids.device)
    return (ar.unsqueeze(0) <= first_eos.unsqueeze(-1)).to(torch.float32)


def sequence_lengths(input_ids, pad_id):
    return (input_ids != pad_id).sum(dim=-1)


def last_token_indices(attn_mask):
    return _last_argmax(attn_mask)


def gather_last_logits(logits, attn_mask):
    idx = last_token_indices(attn_mask)
    return logits[torch.arange(logits.shape[0], device=logits.device), idx]


def truncate_left(input_ids, max_len):
    return input_ids[:, -max_len:]


def pad_to_multiple(input_ids, pad_id, multiple):
    S = input_ids.shape[1]
    rem = S % multiple
    if rem == 0:
        return input_ids
    pad = multiple - rem
    padding = input_ids.new_full((input_ids.shape[0], pad), pad_id)
    return torch.cat([input_ids, padding], dim=1)


def action_logprobs(logits, actions):
    logp = torch.log_softmax(logits, dim=-1)
    return logp.gather(-1, actions.unsqueeze(-1)).squeeze(-1)


def entropy_from_logits(logits):
    logp = torch.log_softmax(logits, dim=-1)
    return -(logp.exp() * logp).sum(dim=-1)


def reward_to_go(rewards, mask):
    rc = (rewards * mask).flip(-1).cumsum(dim=-1).flip(-1)
    return rc * mask


def discount_cumsum(x, gamma):
    out = torch.zeros_like(x)
    running = None
    for i in reversed(range(x.shape[-1])):
        xi = x[..., i]
        running = xi if running is None else xi + gamma * running
        out[..., i] = running
    return out


def build_position_ids_packed(seq_lengths, total_len):
    parts = [torch.arange(int(l)) for l in seq_lengths]
    out = torch.zeros(total_len, dtype=torch.int64)
    if parts:
        pos = torch.cat(parts)
        out[: pos.numel()] = pos
    return out


def logits_warper(logits, temperature, top_k, top_p):
    out = temperature_scale(logits, temperature)
    out = top_k_filter(out, top_k)
    out = top_p_filter(out, top_p)
    return out


def sample_greedy_warped(logits, temperature, top_k, top_p):
    return greedy_token(logits_warper(logits, temperature, top_k, top_p))


def prepare_labels(input_ids, prompt_lengths, ignore_index):
    labels = shift_labels(input_ids, ignore_index)
    ar = torch.arange(input_ids.shape[-1], device=input_ids.device)
    mask = ar.unsqueeze(0) < (prompt_lengths.unsqueeze(-1) - 1)
    labels[mask] = ignore_index
    return labels


def rollout_last_logits(logits, input_ids, pad_id):
    mask = (input_ids != pad_id)
    idx = _last_argmax(mask.to(torch.int64))
    return logits[torch.arange(logits.shape[0], device=logits.device), idx]


def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    remove = probs < epsilon
    keep_idx = _last_argmax(logits)
    remove.scatter_(-1, keep_idx.unsqueeze(-1), False)
    return logits.masked_fill(remove, NEG_INF)


def first_stop_index(input_ids, stop_ids):
    hits = torch.isin(input_ids, stop_ids)
    return _first_index(hits, -1)


def masked_argmax(logits, mask):
    masked = logits.masked_fill(mask == 0, NEG_INF)
    return _last_argmax(masked)


def topk_indices(logits, k):
    k = int(k)
    return torch.sort(logits, dim=-1, descending=True, stable=True).indices[..., :k]


def right_pad_to(input_ids, pad_id, width):
    S = input_ids.shape[1]
    if S >= width:
        return input_ids[:, :width]
    padding = input_ids.new_full((input_ids.shape[0], width - S), pad_id)
    return torch.cat([input_ids, padding], dim=1)


def generation_stop_mask(input_ids, eos_id, min_len):
    L = input_ids.shape[-1]
    first_eos = _first_index(input_ids == eos_id, L)
    effective = first_eos.clamp(min=min_len - 1)
    ar = torch.arange(L, device=input_ids.device)
    return (ar.unsqueeze(0) <= effective.unsqueeze(-1)).to(torch.float32)

Full verification harness

shell
cat > /tmp/verify.py << 'EOF'
import torch, importlib.util, sys
import reforacle as ref
spec = importlib.util.spec_from_file_location("obj", "/app/objectives.py")
obj = importlib.util.module_from_spec(spec); spec.loader.exec_module(obj)

g = torch.Generator().manual_seed(42)
def rl(*s): return torch.randn(*s, generator=g)
def rids(b,s,v): return torch.randint(0,v,(b,s),generator=g)

fails=[]
def chk(name, args, atol=1e-6, rtol=1e-5):
    fa=getattr(ref,name); fo=getattr(obj,name)
    try: r=fa(*args)
    except Exception as e:
        try: o=fo(*args)
        except Exception as e2:
            return
        fails.append((name,"ref-err-only",repr(e))); return
    try: o=fo(*args)
    except Exception as e:
        fails.append((name,"obj-err",repr(e))); return
    if r.dtype!=o.dtype: fails.append((name,"dtype",r.dtype,o.dtype)); return
    if r.shape!=o.shape: fails.append((name,"shape",r.shape,o.shape)); return
    rr=r.float(); oo=o.float()
    rr=torch.nan_to_num(rr,neginf=-1e30,posinf=1e30); oo=torch.nan_to_num(oo,neginf=-1e30,posinf=1e30)
    if not torch.allclose(rr,oo,atol=atol,rtol=rtol):
        d=(rr-oo).abs().max().item(); fails.append((name,"mismatch",d))

for trial in range(20):
    b,s,v = 3,7,9
    logits=rl(b,v)
    ids=rids(b,s,v)
    mask=(rl(b,s)>0).int()
    chk("greedy_token",(rl(b,v),))
    chk("temperature_scale",(logits,0.0))
    chk("temperature_scale",(logits,0.7))
    chk("top_k_filter",(logits,3))
    chk("top_p_filter",(logits,0.8))
    chk("top_p_filter",(logits,0.3))
    chk("min_p_filter",(logits,0.2))
    chk("min_p_filter",(logits,1.3))
    chk("repetition_penalty",(logits,ids,1.5))
    chk("frequency_penalty",(logits,ids,0.3))
    chk("presence_penalty",(logits,ids,0.4))
    chk("suppress_tokens",(logits,[1,4,7]))
    chk("min_length_filter",(logits,3,6,2))
    chk("min_length_filter",(logits,8,6,2))
    chk("renormalize_logits",(logits,))
    chk("attention_mask",(ids*(rl(b,s)>0).long(),0))
    idp=rids(b,s,v); idp[:, -2:]=0
    chk("position_ids",(idp,0))
    chk("causal_mask",(6,))
    chk("shift_labels",(ids,-100))
    chk("completion_loss_labels",(ids,torch.tensor([2,3,5]),-100))
    chk("prepare_labels",(ids,torch.tensor([2,3,5]),-100))
    chk("mask_after_eos",(rids(b,s,v),4))
    chk("sequence_lengths",(idp,0))
    am=(rl(b,s)>0).long(); am[:,0]=1
    chk("last_token_indices",(am,))
    chk("gather_last_logits",(rl(b,s,v),am))
    chk("truncate_left",(ids,4))
    chk("truncate_left",(ids,20))
    chk("pad_to_multiple",(ids,0,4))
    chk("action_logprobs",(rl(b,s,v),rids(b,s,v)))
    chk("entropy_from_logits",(rl(b,s,v),))
    chk("reward_to_go",(rl(b,s),mask))
    chk("discount_cumsum",(rl(b,s),0.95))
    chk("discount_cumsum",(rl(s),0.9))
    chk("build_position_ids_packed",(torch.tensor([2,3,1]),6))
    chk("build_position_ids_packed",(torch.tensor([2,3,1]),9))
    chk("logits_warper",(logits,0.8,4,0.9))
    chk("sample_greedy_warped",(logits,0.8,4,0.9))
    chk("rollout_last_logits",(rl(b,s,v),idp,0))
    chk("epsilon_filter",(logits,0.1))
    chk("epsilon_filter",(logits,0.9))
    chk("first_stop_index",(rids(b,s,v),torch.tensor([4,7])))
    chk("masked_argmax",(logits,(rl(b,v)>0).int()))
    chk("topk_indices",(logits,3))
    chk("right_pad_to",(ids,0,10))
    chk("right_pad_to",(ids,0,4))
    chk("generation_stop_mask",(rids(b,s,v),4,3))

from collections import Counter
c=Counter(f[0] for f in fails)
print("total fail records:", len(fails))
for k,v2 in c.items(): print("FAIL",k,v2)
for f in fails[:20]: print(f)
if not fails: print("ALL PASS")
EOF
PYTHONPATH=/app python3 /tmp/verify.py
output
total fail records: 0
ALL PASS

[stdout]
total fail records: 0
ALL PASS

Stress test ties and boundaries

shell
cat > /tmp/verify2.py << 'EOF'
import torch, importlib.util
import reforacle as ref
spec = importlib.util.spec_from_file_location("obj", "/app/objectives.py")
obj = importlib.util.module_from_spec(spec); spec.loader.exec_module(obj)
g = torch.Generator().manual_seed(7)
fails=[]
def chk(name,args,atol=1e-6,rtol=1e-5):
    try: r=getattr(ref,name)(*args)
    except Exception as e: return
    o=getattr(obj,name)(*args)
    if r.dtype!=o.dtype: fails.append((name,"dtype",r.dtype,o.dtype)); return
    if r.shape!=o.shape: fails.append((name,"shape",tuple(r.shape),tuple(o.shape))); return
    rr=torch.nan_to_num(r.float(),neginf=-1e30,posinf=1e30); oo=torch.nan_to_num(o.float(),neginf=-1e30,posinf=1e30)
    if not torch.allclose(rr,oo,atol=atol,rtol=rtol): fails.append((name,"mismatch",(rr-oo).abs().max().item(),r.tolist(),o.tolist()))

# integer logits => lots of ties
for _ in range(30):
    L=torch.randint(-3,4,(4,8),generator=g).float()
    chk("greedy_token",(L,)); chk("temperature_scale",(L,0.0)); chk("temperature_scale",(L,1.5))
    chk("top_k_filter",(L,3)); chk("top_k_filter",(L,1)); chk("top_k_filter",(L,8))
    chk("top_p_filter",(L,0.0)); chk("top_p_filter",(L,0.5)); chk("top_p_filter",(L,1.0))
    chk("min_p_filter",(L,0.0)); chk("min_p_filter",(L,0.3))
    chk("epsilon_filter",(L,0.05)); chk("epsilon_filter",(L,0.5)); chk("epsilon_filter",(L,0.99))
    chk("masked_argmax",(L,torch.randint(0,2,(4,8),generator=g)))
    chk("topk_indices",(L,4)); chk("topk_indices",(L,1))
    chk("renormalize_logits",(L,)); chk("entropy_from_logits",(L,))
    ids=torch.randint(0,8,(4,6),generator=g)
    chk("repetition_penalty",(L,ids,2.0)); chk("frequency_penalty",(L,ids,0.5)); chk("presence_penalty",(L,ids,0.5))
    chk("logits_warper",(L,1.0,4,0.9),1e-5,1e-4); chk("sample_greedy_warped",(L,1.0,4,0.9))
    m=torch.randint(0,2,(4,6),generator=g); m[:,0]=1
    chk("last_token_indices",(m,)); chk("reward_to_go",(torch.randn(4,6,generator=g),m))
    chk("mask_after_eos",(ids,3)); chk("generation_stop_mask",(ids,3,2)); chk("generation_stop_mask",(ids,3,0))
    chk("first_stop_index",(ids,torch.tensor([3,5])))
    chk("gather_last_logits",(torch.randn(4,6,5,generator=g),m))
    chk("rollout_last_logits",(torch.randn(4,6,5,generator=g),ids,0))

# 3D warper inputs and batch=1
for _ in range(10):
    L=torch.randn(1,12,generator=g)
    chk("top_p_filter",(L,0.7)); chk("min_p_filter",(L,0.15)); chk("epsilon_filter",(L,0.08))
    chk("action_logprobs",(torch.randn(2,3,7,generator=g),torch.randint(0,7,(2,3),generator=g)))
    chk("discount_cumsum",(torch.randn(5,generator=g),0.99))
    chk("build_position_ids_packed",(torch.tensor([1,1,4]),10))

print("fails",len(fails))
for f in fails[:15]: print(f)
if not fails: print("ALL PASS")
EOF
PYTHONPATH=/app python3 /tmp/verify2.py
output
fails 0
ALL PASS

[stdout]
fails 0
ALL PASS

Confirm standalone and complete

shell
grep -nE "reforacle|NotImplementedError|^import|^from" /app/objectives.py; echo "---defs:"; grep -c "^def " /app/objectives.py
output
1:from __future__ import annotations
3:import torch
9:    # Index of the LAST maximum along the last dim (reforacle's tie-break).
---defs:
40

[stdout]
1:from __future__ import annotations
3:import torch
9:    # Index of the LAST maximum along the last dim (reforacle's tie-break).
---defs:
40

Verify standalone load without reforacle

shell
cd /tmp && python3 -c "
import importlib.util
spec=importlib.util.spec_from_file_location('objectives','/app/objectives.py')
m=importlib.util.module_from_spec(spec); spec.loader.exec_module(m)
import torch
print('loads standalone OK')
print('greedy', m.greedy_token(torch.tensor([[1.,3.,2.]])).tolist())
funcs=[n for n in dir(m) if not n.startswith('_') and callable(getattr(m,n)) and n!='torch']
print('callable functions:', len([n for n in funcs if n not in ('annotations',)]))
"
output
loads standalone OK
greedy [1]
callable functions: 37
Shell cwd was reset to /app

[stdout]
loads standalone OK
greedy [1]
callable functions: 37
[stderr]

Shell cwd was reset to /app
All 37 functions are implemented and verified. Summary of what I recovered by probing `reforacle`: ## Key behaviors discovered **Non-obvious conventions** (would have been wrong if assumed from names): - **Last-maximum tie-break** , `greedy_token`, `masked_argmax`, and `temperature_scale(T=0)` all select the *last* argmax, not torch's default first. Implemented via a `_last_argmax` helper (`flip(-1).argmax`). - **`repetition_penalty` applies per-occurrence** , a token repeated *c* times is penalized `penalty**c` (divide if logit>0, multiply if <0), not the standard once-only. - **`epsilon_filter` safety keeps exactly one token** , the single last-argmax token, even when the max is tied (it dropped a tied-max sibling). `min_p_filter` has *no* safety (can zero everything). - **`prepare_labels` shifts then masks prompt** (`-100` on first `prompt_len-1`), while **`completion_loss_labels` masks without shifting**. - **`last_token_indices` finds the last *non-zero* index** (handles gaps), not `sum-1`. - **`generation_stop_mask`** stop point is `max(first_eos, min_len-1)`; **`reward_to_go`** is `reverse_cumsum(rewards*mask)*mask`; **`right_pad_to`** truncates when width < seq; **`topk_indices`** uses a stable descending sort (differs from `torch.topk` tie order). - **`position_ids`** is `(mask.cumsum-1).clamp(min=0)` (pad positions carry the running value, not HF's fill-with-1). ## Verification Ran two comparison harnesses against `reforacle` (600+ calls total), including integer-valued logits to force heavy ties and boundary values (p∈{0,1}, over-threshold min_p/epsilon, no-eos rows, leading/interior padding). **Zero mismatches** on values, dtypes, and shapes at the required tolerances. The file imports only `torch`, contains no reference to `reforacle`, and loads cleanly in a sandbox where `reforacle` is absent.
Final diff , what the agent changed

No file changes, the agent left the workspace untouched.

Verifier output , the proof behind the reward
tests/test.sh
=== SyncVals verifier (post-training decoding/collation, sealed oracle) ===

FUNCTIONS PASSED: 35/37  (partial_score=0.9459)
REWARD = 0   (1 iff ALL 37 functions match the sealed reference exactly)

The agent probes a sealed compiled oracle (reforacle) to recover 37 decoding/sampling/collation conventions (logit processors, sampling filters, sequence collation, RLHF/SFT rollout helpers) and reimplements them in torch.

Reproduce this trial: git checkout 2f94510 && PYTHONPATH=src python3 scripts/build_site.py , then open trial/trial_0caf178aa1bf40be. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.

Trial trial_0caf178aa1bf40be · verifier authoritative; classifier explanatory.