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 Reverse-Engineering
EvidenceTest output: 'FUNCTIONS PASSED: 32/37 (partial_score=0.8649) REWARD = 0 (1 iff ALL 37 functions match)'. Agent systematically probed oracle across 20+ steps, implementing all 37 functions. 200-trial comparison of logit functions passed (0 failures), but comprehensive 300-trial testing revealed numerical mismatches in functions like last_token_indices, gather_last_logits, and rollout_last_logits. Task explicitly requires exact numerical agreement (atol=1e-6/1e-5 for primitives, 1e-5/1e-4 for composites) on unseen held-out inputs.
Root causeThe agent successfully reverse-engineered 86.49% of the 37 decoding/collation functions through systematic black-box probing, but 5 functions have subtle numerical behaviors or edge cases that weren't fully captured by the sampled test cases. Reverse-engineering without a written spec inherently cannot guarantee coverage of all hidden test inputs.
RecommendationN/A - task is fine. This is a legitimately hard reverse-engineering task designed to test agent exploration and pattern-discovery under uncertainty. A 32/37 pass rate reflects genuine difficulty, not a task specification problem. The oracle is the ground truth, and the grading requirements are clear.
Trajectory
Tool-by-tool agent trajectory
41 tool calls · 4 tool types · 43 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 probe.py

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.

List reforacle functions

shell
python3 -c "import reforacle; print([x for x in dir(reforacle) if not x.startswith('__')])"
output
['F', 'NEG_INF', 'action_logprobs', 'attention_mask', 'build_position_ids_packed', 'causal_mask', 'completion_loss_labels', 'discount_cumsum', 'entropy_from_logits', 'epsilon_filter', 'first_stop_index', 'frequency_penalty', 'gather_last_logits', 'generation_stop_mask', 'greedy_token', 'last_token_indices', 'logits_warper', 'mask_after_eos', 'masked_argmax', 'min_length_filter', 'min_p_filter', 'pad_to_multiple', 'position_ids', 'prepare_labels', 'presence_penalty', 'renormalize_logits', 'repetition_penalty', 'reward_to_go', 'right_pad_to', 'rollout_last_logits', 'sample_greedy_warped', 'sequence_lengths', 'shift_labels', 'suppress_tokens', 'temperature_scale', 'top_k_filter', 'top_p_filter', 'topk_indices', 'torch', 'truncate_left']

[stdout]
['F', 'NEG_INF', 'action_logprobs', 'attention_mask', 'build_position_ids_packed', 'causal_mask', 'completion_loss_labels', 'discount_cumsum', 'entropy_from_logits', 'epsilon_filter', 'first_stop_index', 'frequency_penalty', 'gather_last_logits', 'generation_stop_mask', 'greedy_token', 'last_token_indices', 'logits_warper', 'mask_after_eos', 'masked_argmax', 'min_length_filter', 'min_p_filter', 'pad_to_multiple', 'position_ids', 'prepare_labels', 'presence_penalty', 'renormalize_logits', 'repetition_penalty', 'reward_to_go', 'right_pad_to', 'rollout_last_logits', 'sample_greedy_warped', 'sequence_lengths', 'shift_labels', 'suppress_tokens', 'temperature_scale', 'top_k_filter', 'top_p_filter', 'topk_indices', 'torch', 'truncate_left']

Probe temperature_scale

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)

# temperature_scale non-zero
logits=torch.randn(2,5,generator=g)
print('logits',logits.tolist())
print('temp2', R.temperature_scale(logits,2.0).tolist())
print('temp0.5', R.temperature_scale(logits,0.5).tolist())
# check: is it logits/T? 
print('check', (logits/2.0).tolist())
"
output
logits [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427], [-1.3985954523086548, 0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914]]
temp2 [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713], [-0.6992977261543274, 0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457]]
temp0.5 [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854], [-2.7971909046173096, 0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828]]
check [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713], [-0.6992977261543274, 0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457]]

[stdout]
logits [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427], [-1.3985954523086548, 0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914]]
temp2 [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713], [-0.6992977261543274, 0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457]]
temp0.5 [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854], [-2.7971909046173096, 0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828]]
check [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713], [-0.6992977261543274, 0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457]]

Probe filters

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(1)
logits=torch.randn(2,6,generator=g)
print('logits',[[round(x,3) for x in r] for r in logits.tolist()])
print('greedy', R.greedy_token(logits).tolist(), R.greedy_token(logits).dtype, R.greedy_token(logits).shape)
# top_k
print('top_k2', R.top_k_filter(logits,2).tolist())
print('topk_indices2', R.topk_indices(logits,2).tolist(), R.topk_indices(logits,2).dtype)
# top_p
print('top_p0.8', R.top_p_filter(logits,0.8).tolist())
# min_p
print('min_p0.3', R.min_p_filter(logits,0.3).tolist())
# epsilon
print('epsilon0.1', R.epsilon_filter(logits,0.1).tolist())
"
output
logits [[0.661, 0.267, 0.062, 0.621, -0.452, -0.166], [-1.523, 0.382, -1.028, -0.563, -0.892, -0.058]]
greedy [0, 1] torch.int64 torch.Size([2])
top_k2 [[0.6613521575927734, -inf, -inf, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -inf, -inf, -0.05825017765164375]]
topk_indices2 [[0, 3], [1, 5]] torch.int64
top_p0.8 [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
min_p0.3 [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -inf, -0.05825017765164375]]
epsilon0.1 [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -inf, -0.16613022983074188], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]

[stdout]
logits [[0.661, 0.267, 0.062, 0.621, -0.452, -0.166], [-1.523, 0.382, -1.028, -0.563, -0.892, -0.058]]
greedy [0, 1] torch.int64 torch.Size([2])
top_k2 [[0.6613521575927734, -inf, -inf, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -inf, -inf, -0.05825017765164375]]
topk_indices2 [[0, 3], [1, 5]] torch.int64
top_p0.8 [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
min_p0.3 [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -inf, -0.05825017765164375]]
epsilon0.1 [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -inf, -0.16613022983074188], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]

Probe penalties

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(2)
logits=torch.randn(2,6,generator=g)
ids=torch.tensor([[1,1,2,3],[0,4,4,4]])
print('logits',[[round(x,3) for x in r] for r in logits.tolist()])
print('rep2', R.repetition_penalty(logits,ids,1.5).tolist())
print('freq', R.frequency_penalty(logits,ids,0.5).tolist())
print('pres', R.presence_penalty(logits,ids,0.5).tolist())
print('suppress', R.suppress_tokens(logits,[1,3]).tolist())
print('minlen', R.min_length_filter(logits,3,5,2).tolist())
print('renorm', R.renormalize_logits(logits).tolist())
print('renorm sum', R.renormalize_logits(logits).exp().sum(-1).tolist())
"
output
logits [[0.392, -0.224, -0.32, -1.205, 1.044, -0.633], [0.573, 0.541, -0.392, -1.043, 1.319, 0.748]]
rep2 [[0.39229682087898254, -0.5030190348625183, -0.47925040125846863, -1.8075556755065918, 1.0444635152816772, -0.6332277059555054], [0.3820711672306061, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 0.39070141315460205, 0.747639000415802]]
freq [[0.39229682087898254, -1.2235640287399292, -0.8195002675056458, -1.7050371170043945, 1.0444635152816772, -0.6332277059555054], [0.07310676574707031, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, -0.1813826560974121, 0.747639000415802]]
pres [[0.39229682087898254, -0.7235640287399292, -0.8195002675056458, -1.7050371170043945, 1.0444635152816772, -0.6332277059555054], [0.07310676574707031, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 0.8186173439025879, 0.747639000415802]]
suppress [[0.39229682087898254, -inf, -0.31950026750564575, -inf, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, -inf, -0.39190584421157837, -inf, 1.318617343902588, 0.747639000415802]]
minlen [[0.39229682087898254, -0.223564013838768, -inf, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -inf, -1.0426788330078125, 1.318617343902588, 0.747639000415802]]
renorm [[-1.5066688060760498, -2.1225297451019287, -2.21846604347229, -3.1040029525756836, -0.8545021414756775, -2.532193422317505], [-1.7658040523529053, -1.7979633808135986, -2.730816602706909, -3.381589651107788, -1.0202934741973877, -1.5912718772888184]]
renorm sum [1.0, 1.0]

[stdout]
logits [[0.392, -0.224, -0.32, -1.205, 1.044, -0.633], [0.573, 0.541, -0.392, -1.043, 1.319, 0.748]]
rep2 [[0.39229682087898254, -0.5030190348625183, -0.47925040125846863, -1.8075556755065918, 1.0444635152816772, -0.6332277059555054], [0.3820711672306061, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 0.39070141315460205, 0.747639000415802]]
freq [[0.39229682087898254, -1.2235640287399292, -0.8195002675056458, -1.7050371170043945, 1.0444635152816772, -0.6332277059555054], [0.07310676574707031, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, -0.1813826560974121, 0.747639000415802]]
pres [[0.39229682087898254, -0.7235640287399292, -0.8195002675056458, -1.7050371170043945, 1.0444635152816772, -0.6332277059555054], [0.07310676574707031, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 0.8186173439025879, 0.747639000415802]]
suppress [[0.39229682087898254, -inf, -0.31950026750564575, -inf, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, -inf, -0.39190584421157837, -inf, 1.318617343902588, 0.747639000415802]]
minlen [[0.39229682087898254, -0.223564013838768, -inf, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -inf, -1.0426788330078125, 1.318617343902588, 0.747639000415802]]
renorm [[-1.5066688060760498, -2.1225297451019287, -2.21846604347229, -3.1040029525756836, -0.8545021414756775, -2.532193422317505], [-1.7658040523529053, -1.7979633808135986, -2.730816602706909, -3.381589651107788, -1.0202934741973877, -1.5912718772888184]]
renorm sum [1.0, 1.0]

Probe masks

shell
python3 -c "
import torch, reforacle as R
ids=torch.tensor([[5,6,7,0,0],[1,2,0,0,0]])
print('attn', R.attention_mask(ids,0).tolist(), R.attention_mask(ids,0).dtype)
print('posids', R.position_ids(ids,0).tolist(), R.position_ids(ids,0).dtype)
print('causal4', R.causal_mask(4).tolist(), R.causal_mask(4).dtype)
print('seqlen', R.sequence_lengths(ids,0).tolist(), R.sequence_lengths(ids,0).dtype)
am=R.attention_mask(ids,0)
print('last_tok', R.last_token_indices(am).tolist(), R.last_token_indices(am).dtype)
print('shift', R.shift_labels(ids,-100).tolist())
print('mask_after_eos', R.mask_after_eos(torch.tensor([[5,2,6,7],[1,2,3,2]]),2).tolist())
"
output
attn [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] torch.float32
posids [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1]] torch.int64
causal4 [[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]] torch.float32
seqlen [3, 2] torch.int64
last_tok [2, 1] torch.int64
shift [[6, 7, 0, 0, -100], [2, 0, 0, 0, -100]]
mask_after_eos [[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]

[stdout]
attn [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] torch.float32
posids [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1]] torch.int64
causal4 [[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]] torch.float32
seqlen [3, 2] torch.int64
last_tok [2, 1] torch.int64
shift [[6, 7, 0, 0, -100], [2, 0, 0, 0, -100]]
mask_after_eos [[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]

Probe collation

shell
python3 -c "
import torch, reforacle as R
# mask_after_eos no eos
print('noeos', R.mask_after_eos(torch.tensor([[5,6,7]]),2).tolist())
# completion_loss_labels
ids=torch.tensor([[5,6,7,8,9],[1,2,3,4,5]])
print('comp', R.completion_loss_labels(ids,torch.tensor([2,3]),-100).tolist())
print('prepare', R.prepare_labels(ids,torch.tensor([2,3]),-100).tolist())
# gather_last_logits
logits=torch.arange(2*4*3,dtype=torch.float32).reshape(2,4,3)
am=torch.tensor([[1,1,1,0],[1,1,0,0]])
print('gather', R.gather_last_logits(logits,am).tolist())
# truncate_left
print('trunc', R.truncate_left(torch.tensor([[1,2,3,4,5]]),3).tolist())
# pad_to_multiple
print('pad_mult', R.pad_to_multiple(torch.tensor([[1,2,3]]),0,4).tolist())
print('pad_mult5', R.pad_to_multiple(torch.tensor([[1,2,3,4,5]]),0,4).tolist())
# right_pad_to
print('rpad', R.right_pad_to(torch.tensor([[1,2,3]]),0,6).tolist())
"
output
noeos [[1.0, 1.0, 1.0]]
comp [[-100, -100, 7, 8, 9], [-100, -100, -100, 4, 5]]
prepare [[-100, 7, 8, 9, -100], [-100, -100, 4, 5, -100]]
gather [[6.0, 7.0, 8.0], [15.0, 16.0, 17.0]]
trunc [[3, 4, 5]]
pad_mult [[1, 2, 3, 0]]
pad_mult5 [[1, 2, 3, 4, 5, 0, 0, 0]]
rpad [[1, 2, 3, 0, 0, 0]]

[stdout]
noeos [[1.0, 1.0, 1.0]]
comp [[-100, -100, 7, 8, 9], [-100, -100, -100, 4, 5]]
prepare [[-100, 7, 8, 9, -100], [-100, -100, 4, 5, -100]]
gather [[6.0, 7.0, 8.0], [15.0, 16.0, 17.0]]
trunc [[3, 4, 5]]
pad_mult [[1, 2, 3, 0]]
pad_mult5 [[1, 2, 3, 4, 5, 0, 0, 0]]
rpad [[1, 2, 3, 0, 0, 0]]

Probe logprobs/entropy

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(3)
logits=torch.randn(2,4,5,generator=g)
actions=torch.tensor([[1,2,0,3],[4,1,2,0]])
print('alp', R.action_logprobs(logits,actions).tolist())
print('alp shape', R.action_logprobs(logits,actions).shape)
# verify
lsm=torch.log_softmax(logits,-1)
print('verify', lsm.gather(-1,actions.unsqueeze(-1)).squeeze(-1).tolist())
l2=torch.randn(2,5,generator=g)
print('entropy', R.entropy_from_logits(l2).tolist())
p=torch.softmax(l2,-1)
print('ent verify', (-(p*torch.log_softmax(l2,-1)).sum(-1)).tolist())
"
output
alp [[-1.4040248394012451, -3.637298345565796, -1.7765860557556152, -2.329641580581665], [-3.1036596298217773, -1.931173324584961, -0.9780085682868958, -1.1429297924041748]]
alp shape torch.Size([2, 4])
verify [[-1.4040248394012451, -3.637298345565796, -1.7765860557556152, -2.329641580581665], [-3.1036596298217773, -1.931173324584961, -0.9780085682868958, -1.1429297924041748]]
entropy [1.4500930309295654, 1.4904046058654785]
ent verify [1.4500930309295654, 1.490404725074768]

[stdout]
alp [[-1.4040248394012451, -3.637298345565796, -1.7765860557556152, -2.329641580581665], [-3.1036596298217773, -1.931173324584961, -0.9780085682868958, -1.1429297924041748]]
alp shape torch.Size([2, 4])
verify [[-1.4040248394012451, -3.637298345565796, -1.7765860557556152, -2.329641580581665], [-3.1036596298217773, -1.931173324584961, -0.9780085682868958, -1.1429297924041748]]
entropy [1.4500930309295654, 1.4904046058654785]
ent verify [1.4500930309295654, 1.490404725074768]

Probe RL functions

shell
python3 -c "
import torch, reforacle as R
rewards=torch.tensor([[1.,2.,3.,4.],[5.,6.,7.,8.]])
mask=torch.tensor([[1.,1.,1.,0.],[1.,1.,0.,0.]])
print('r2g', R.reward_to_go(rewards,mask).tolist())
print('r2g nomask?', R.reward_to_go(rewards,torch.ones_like(mask)).tolist())
x=torch.tensor([[1.,2.,3.,4.]])
print('disc0.9', R.discount_cumsum(x,0.9).tolist())
print('bpp', R.build_position_ids_packed(torch.tensor([2,3,1]),6).tolist(), R.build_position_ids_packed(torch.tensor([2,3,1]),6).dtype)
"
output
r2g [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]]
r2g nomask? [[10.0, 9.0, 7.0, 4.0], [26.0, 21.0, 15.0, 8.0]]
disc0.9 [[8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]]
bpp [0, 1, 0, 1, 2, 0] torch.int64

[stdout]
r2g [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]]
r2g nomask? [[10.0, 9.0, 7.0, 4.0], [26.0, 21.0, 15.0, 8.0]]
disc0.9 [[8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]]
bpp [0, 1, 0, 1, 2, 0] torch.int64

Probe more

shell
python3 -c "
import torch, reforacle as R
# r2g mask edge
r=torch.tensor([[1.,2.,3.,4.]]); m=torch.tensor([[1.,0.,1.,1.]])
print('r2g edge', R.reward_to_go(r,m).tolist())
# masked_argmax
logits=torch.tensor([[1.,5.,3.,2.],[4.,2.,9.,1.]])
mask=torch.tensor([[1,1,0,1],[0,1,1,1]])
print('margmax', R.masked_argmax(logits,mask).tolist(), R.masked_argmax(logits,mask).dtype)
# first_stop_index
ids=torch.tensor([[5,6,7,8,9],[1,3,2,4,3]])
print('fsi', R.first_stop_index(ids,[7,3]).tolist(), R.first_stop_index(ids,[7,3]).dtype)
print('fsi nostop', R.first_stop_index(torch.tensor([[5,6]]),[9]).tolist())
# generation_stop_mask
ids2=torch.tensor([[5,2,6],[1,2,2]])
print('gsm', R.generation_stop_mask(ids2,2,1).tolist())
"
output
Exit code 1
r2g edge [[8.0, 0.0, 7.0, 4.0]]
margmax [1, 2] torch.int64
Traceback (most recent call last):
  File "<string>", line 12, in <module>
  File "reforacle.py", line 243, in reforacle.first_stop_index
AttributeError: 'list' object has no attribute 'tolist'

[error] tool reported failure

first_stop_index type

shell
python3 -c "
import torch, reforacle as R
ids=torch.tensor([[5,6,7,8,9],[1,3,2,4,3]])
r=R.first_stop_index(ids,[7,3])
print('fsi', type(r), r)
r2=R.first_stop_index(torch.tensor([[5,6]]),[9])
print('fsi nostop', type(r2), r2)
print('fsi tensorarg', R.first_stop_index(ids,torch.tensor([7,3])))
ids2=torch.tensor([[5,2,6],[1,2,2]])
print('gsm', R.generation_stop_mask(ids2,2,1).tolist())
print('gsm minlen3', R.generation_stop_mask(torch.tensor([[5,2,6,2]]),2,3).tolist())
"
output
Exit code 1
Traceback (most recent call last):
  File "<string>", line 4, in <module>
  File "reforacle.py", line 243, in reforacle.first_stop_index
AttributeError: 'list' object has no attribute 'tolist'

[error] tool reported failure

first_stop_index tensor arg

shell
python3 -c "
import torch, reforacle as R
ids=torch.tensor([[5,6,7,8,9],[1,3,2,4,3]])
r=R.first_stop_index(ids,torch.tensor([7,3]))
print('fsi', type(r), r)
r2=R.first_stop_index(torch.tensor([[5,6]]),torch.tensor([9]))
print('fsi nostop', type(r2), r2, r2.dtype)
ids2=torch.tensor([[5,2,6],[1,2,2]])
print('gsm', R.generation_stop_mask(ids2,2,1).tolist(), R.generation_stop_mask(ids2,2,1).dtype)
print('gsm minlen3', R.generation_stop_mask(torch.tensor([[5,2,6,2]]),2,3).tolist())
print('gsm noeos', R.generation_stop_mask(torch.tensor([[5,6,7]]),2,1).tolist())
"
output
fsi <class 'torch.Tensor'> tensor([2, 1])
fsi nostop <class 'torch.Tensor'> tensor([-1]) torch.int64
gsm [[1.0, 1.0, 0.0], [1.0, 1.0, 0.0]] torch.float32
gsm minlen3 [[1.0, 1.0, 1.0, 0.0]]
gsm noeos [[1.0, 1.0, 1.0]]

[stdout]
fsi <class 'torch.Tensor'> tensor([2, 1])
fsi nostop <class 'torch.Tensor'> tensor([-1]) torch.int64
gsm [[1.0, 1.0, 0.0], [1.0, 1.0, 0.0]] torch.float32
gsm minlen3 [[1.0, 1.0, 1.0, 0.0]]
gsm noeos [[1.0, 1.0, 1.0]]

Composite warper

shell
python3 -c "
import torch, reforacle as R
print('gsm test', R.generation_stop_mask(torch.tensor([[2,6,2,7]]),2,2).tolist())
g=torch.Generator().manual_seed(5)
logits=torch.randn(2,8,generator=g)
print('logits',[[round(x,3) for x in r] for r in logits.tolist()])
print('warper', R.logits_warper(logits,2.0,3,0.9).tolist())
print('sgw', R.sample_greedy_warped(logits,2.0,3,0.9).tolist(), R.sample_greedy_warped(logits,2.0,3,0.9).dtype)
"
output
gsm test [[1.0, 1.0, 0.0, 0.0]]
logits [[1.842, 0.519, -1.712, -1.701, 2.019, -0.269, -0.131, -1.437], [0.391, -0.019, -1.353, -0.731, 0.988, -0.419, -0.585, -0.782]]
warper [[0.9211649298667908, 0.25944358110427856, -inf, -inf, 1.0097097158432007, -inf, -inf, -inf], [0.19538411498069763, -0.00948388036340475, -inf, -inf, 0.4939611256122589, -inf, -inf, -inf]]
sgw [4, 4] torch.int64

[stdout]
gsm test [[1.0, 1.0, 0.0, 0.0]]
logits [[1.842, 0.519, -1.712, -1.701, 2.019, -0.269, -0.131, -1.437], [0.391, -0.019, -1.353, -0.731, 0.988, -0.419, -0.585, -0.782]]
warper [[0.9211649298667908, 0.25944358110427856, -inf, -inf, 1.0097097158432007, -inf, -inf, -inf], [0.19538411498069763, -0.00948388036340475, -inf, -inf, 0.4939611256122589, -inf, -inf, -inf]]
sgw [4, 4] torch.int64

rollout_last_logits

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(6)
logits=torch.randn(2,4,5,generator=g)
ids=torch.tensor([[5,6,7,0],[1,2,0,0]])
print('roll', R.rollout_last_logits(logits,ids,0).tolist())
# expect gather last non-pad
am=(ids!=0)
lastidx=am.sum(-1)-1
print('expect', logits[torch.arange(2),lastidx].tolist())
"
output
roll [[2.548506021499634, 0.30063536763191223, -0.543218195438385, -1.0841294527053833, 1.461236596107483], [-0.240716814994812, -2.641396999359131, -1.8220865726470947, 0.9003403782844543, -0.19853952527046204]]
expect [[2.548506021499634, 0.30063536763191223, -0.543218195438385, -1.0841294527053833, 1.461236596107483], [-0.240716814994812, -2.641396999359131, -1.8220865726470947, 0.9003403782844543, -0.19853952527046204]]

[stdout]
roll [[2.548506021499634, 0.30063536763191223, -0.543218195438385, -1.0841294527053833, 1.461236596107483], [-0.240716814994812, -2.641396999359131, -1.8220865726470947, 0.9003403782844543, -0.19853952527046204]]
expect [[2.548506021499634, 0.30063536763191223, -0.543218195438385, -1.0841294527053833, 1.461236596107483], [-0.240716814994812, -2.641396999359131, -1.8220865726470947, 0.9003403782844543, -0.19853952527046204]]

Edge cases

shell
python3 -c "
import torch, reforacle as R
print('trunc short', R.truncate_left(torch.tensor([[1,2]]),5).tolist())
print('rpad short', R.right_pad_to(torch.tensor([[1,2,3,4]]),0,2).tolist())
print('topk big', R.top_k_filter(torch.tensor([[1.,2.,3.]]),5).tolist())
print('topp1', R.top_p_filter(torch.tensor([[1.,2.,3.]]),1.0).tolist())
print('topp0', R.top_p_filter(torch.tensor([[1.,2.,3.]]),0.0).tolist())
# float64 preserve
l64=torch.tensor([[1.,2.,3.]],dtype=torch.float64)
print('temp dtype', R.temperature_scale(l64,2.0).dtype)
print('greedy dtype in float', R.greedy_token(l64).dtype)
# min_length cur>=min
print('minlen noop', R.min_length_filter(torch.tensor([[1.,2.,3.]]),5,3,1).tolist())
# shift dtype
print('shift dtype', R.shift_labels(torch.tensor([[1,2,3]]),-100).dtype)
"
output
trunc short [[1, 2]]
rpad short [[1, 2]]
topk big [[1.0, 2.0, 3.0]]
topp1 [[1.0, 2.0, 3.0]]
topp0 [[-inf, -inf, 3.0]]
temp dtype torch.float64
greedy dtype in float torch.int64
minlen noop [[1.0, 2.0, 3.0]]
shift dtype torch.int64

[stdout]
trunc short [[1, 2]]
rpad short [[1, 2]]
topk big [[1.0, 2.0, 3.0]]
topp1 [[1.0, 2.0, 3.0]]
topp0 [[-inf, -inf, 3.0]]
temp dtype torch.float64
greedy dtype in float torch.int64
minlen noop [[1.0, 2.0, 3.0]]
shift dtype torch.int64

More edges

shell
python3 -c "
import torch, reforacle as R
# epsilon all below
print('eps big', R.epsilon_filter(torch.tensor([[1.,1.,1.,1.]]),0.5).tolist())
# min_p keep
print('minp big', R.min_p_filter(torch.tensor([[1.,2.,3.]]),0.99).tolist())
# top_k ties
print('topk ties', R.top_k_filter(torch.tensor([[1.,3.,3.,2.]]),1).tolist())
print('topk ties2', R.top_k_filter(torch.tensor([[1.,3.,3.,2.]]),2).tolist())
# bpp total>sum
print('bpp big', R.build_position_ids_packed(torch.tensor([2,2]),6).tolist())
"
output
eps big [[-inf, -inf, -inf, 1.0]]
minp big [[-inf, -inf, 3.0]]
topk ties [[-inf, 3.0, 3.0, -inf]]
topk ties2 [[-inf, 3.0, 3.0, -inf]]
bpp big [0, 1, 0, 1, 0, 0]

[stdout]
eps big [[-inf, -inf, -inf, 1.0]]
minp big [[-inf, -inf, 3.0]]
topk ties [[-inf, 3.0, 3.0, -inf]]
topk ties2 [[-inf, 3.0, 3.0, -inf]]
bpp big [0, 1, 0, 1, 0, 0]

epsilon keep-one tie

shell
python3 -c "
import torch, reforacle as R
print('eps a', R.epsilon_filter(torch.tensor([[1.,4.,2.,3.]]),0.9).tolist())
print('eps b', R.epsilon_filter(torch.tensor([[4.,1.,2.,3.]]),0.9).tolist())
print('eps ties', R.epsilon_filter(torch.tensor([[3.,3.,1.,3.]]),0.9).tolist())
# what about min_p all below (can't since max always kept)
# check min_p tie
print('minp tie', R.min_p_filter(torch.tensor([[3.,3.,1.]]),1.0).tolist())
"
output
eps a [[-inf, 4.0, -inf, -inf]]
eps b [[4.0, -inf, -inf, -inf]]
eps ties [[-inf, -inf, -inf, 3.0]]
minp tie [[3.0, 3.0, -inf]]

[stdout]
eps a [[-inf, 4.0, -inf, -inf]]
eps b [[4.0, -inf, -inf, -inf]]
eps ties [[-inf, -inf, -inf, 3.0]]
minp tie [[3.0, 3.0, -inf]]

Confirm last-max

shell
python3 -c "
import torch, reforacle as R
print(R.epsilon_filter(torch.tensor([[3.,1.,3.,3.]]),0.9).tolist())
print(R.epsilon_filter(torch.tensor([[3.,3.,3.,1.]]),0.9).tolist())
print(R.epsilon_filter(torch.tensor([[3.,3.,3.,3.]]),0.9).tolist())
print('topk ties idx', torch.topk(torch.tensor([3.,3.,1.,3.]),1).indices)
print('flip', torch.tensor([3.,3.,1.,3.]).flip(-1).argmax())
"
output
[[-inf, -inf, -inf, 3.0]]
[[-inf, -inf, 3.0, -inf]]
[[-inf, -inf, -inf, 3.0]]
topk ties idx tensor([1])
flip tensor(0)

[stdout]
[[-inf, -inf, -inf, 3.0]]
[[-inf, -inf, 3.0, -inf]]
[[-inf, -inf, -inf, 3.0]]
topk ties idx tensor([1])
flip tensor(0)
write · /app/objectives.py
from __future__ import annotations

import torch

NEG_INF = float("-inf")


def greedy_token(logits):
    return logits.argmax(dim=-1)


def temperature_scale(logits, temperature):
    if temperature == 0:
        out = torch.full_like(logits, NEG_INF)
        idx = logits.argmax(dim=-1, keepdim=True)
        out.scatter_(-1, idx, 0.0)
        return out
    return logits / temperature


def top_k_filter(logits, k):
    k = min(k, logits.size(-1))
    kth = torch.topk(logits, k, dim=-1).values[..., -1, None]
    remove = logits < kth
    return logits.masked_fill(remove, NEG_INF)


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


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


def repetition_penalty(logits, input_ids, penalty):
    counts = torch.zeros_like(logits)
    counts.scatter_add_(-1, input_ids, torch.ones_like(input_ids, dtype=logits.dtype))
    factor = penalty ** counts
    return torch.where(logits > 0, logits / factor, logits * factor)


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


def presence_penalty(logits, input_ids, penalty):
    counts = torch.zeros_like(logits)
    counts.scatter_add_(-1, input_ids, torch.ones_like(input_ids, dtype=logits.dtype))
    present = (counts > 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).float()


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


def causal_mask(t):
    mask = torch.full((t, t), NEG_INF)
    return torch.triu(mask, 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.size(-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):
    is_eos = (input_ids == eos_id)
    before = is_eos.long().cumsum(dim=-1) - is_eos.long()
    return (before == 0).float()


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


def last_token_indices(attn_mask):
    return attn_mask.sum(dim=-1).long() - 1


def gather_last_logits(logits, attn_mask):
    idx = attn_mask.sum(dim=-1).long() - 1
    batch = torch.arange(logits.size(0), device=logits.device)
    return logits[batch, idx]


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


def pad_to_multiple(input_ids, pad_id, multiple):
    seq = input_ids.size(-1)
    rem = seq % multiple
    if rem == 0:
        return input_ids
    pad = multiple - rem
    return torch.nn.functional.pad(input_ids, (0, pad), value=pad_id)


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)
    p = logp.exp()
    return -(p * logp).sum(dim=-1)


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


def discount_cumsum(x, gamma):
    out = torch.zeros_like(x)
    running = torch.zeros(x.shape[:-1], dtype=x.dtype, device=x.device)
    for t in range(x.size(-1) - 1, -1, -1):
        running = x[..., t] + gamma * running
        out[..., t] = running
    return out


def build_position_ids_packed(seq_lengths, total_len):
    parts = [torch.arange(int(l), dtype=torch.long) for l in seq_lengths.tolist()]
    out = torch.cat(parts) if parts else torch.zeros(0, dtype=torch.long)
    if out.numel() < total_len:
        pad = torch.zeros(total_len - out.numel(), dtype=torch.long)
        out = torch.cat([out, pad])
    elif out.numel() > total_len:
        out = out[:total_len]
    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.size(-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):
    idx = (input_ids != pad_id).sum(dim=-1) - 1
    batch = torch.arange(logits.size(0), device=logits.device)
    return logits[batch, idx]


def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    remove = probs < epsilon
    n = logits.size(-1)
    last_max = (n - 1) - logits.flip(-1).argmax(dim=-1)
    remove.scatter_(-1, last_max.unsqueeze(-1), False)
    return logits.masked_fill(remove, NEG_INF)


def first_stop_index(input_ids, stop_ids):
    is_stop = (input_ids.unsqueeze(-1) == stop_ids).any(dim=-1)
    any_match = is_stop.any(dim=-1)
    idx = is_stop.float().argmax(dim=-1)
    return torch.where(any_match, idx, torch.full_like(idx, -1))


def masked_argmax(logits, mask):
    masked = logits.masked_fill(~mask.bool(), NEG_INF)
    return masked.argmax(dim=-1)


def topk_indices(logits, k):
    return torch.topk(logits, k, dim=-1).indices


def right_pad_to(input_ids, pad_id, width):
    seq = input_ids.size(-1)
    if seq >= width:
        return input_ids[:, :width]
    pad = width - seq
    return torch.nn.functional.pad(input_ids, (0, pad), value=pad_id)


def generation_stop_mask(input_ids, eos_id, min_len):
    mask = mask_after_eos(input_ids, eos_id)
    mask[:, :min_len] = 1.0
    return mask

Compare logit functions

shell
python3 -c "
import torch, importlib
import reforacle as R
import objectives as O

torch.manual_seed(0)
fails=[]
def cmp(name, a, b, atol=1e-6, rtol=1e-5):
    try:
        if isinstance(a, torch.Tensor):
            if a.dtype != b.dtype:
                fails.append(f'{name}: dtype {a.dtype} vs {b.dtype}')
                return
            if a.shape != b.shape:
                fails.append(f'{name}: shape {a.shape} vs {b.shape}')
                return
            if not torch.allclose(a.float(), b.float(), atol=atol, rtol=rtol, equal_nan=True):
                fails.append(f'{name}: values differ max {(a.float()-b.float()).abs().max()}')
    except Exception as e:
        fails.append(f'{name}: EXC {e}')

for trial in range(200):
    g=torch.Generator().manual_seed(trial)
    B=int(torch.randint(1,4,(1,),generator=g))
    V=int(torch.randint(3,10,(1,),generator=g))
    logits=torch.randn(B,V,generator=g)
    ids=torch.randint(0,V,(B,int(torch.randint(2,6,(1,),generator=g))),generator=g)
    cmp('greedy',R.greedy_token(logits),O.greedy_token(logits))
    for T in [0.0,0.5,1.0,2.0]:
        cmp(f'temp{T}',R.temperature_scale(logits,T),O.temperature_scale(logits,T))
    for k in [1,2,V]:
        cmp(f'topk{k}',R.top_k_filter(logits,k),O.top_k_filter(logits,k))
        cmp(f'topki{k}',R.topk_indices(logits,k),O.topk_indices(logits,k))
    for p in [0.0,0.3,0.8,0.95,1.0]:
        cmp(f'topp{p}',R.top_p_filter(logits,p),O.top_p_filter(logits,p))
        cmp(f'minp{p}',R.min_p_filter(logits,p),O.min_p_filter(logits,p))
    for e in [0.05,0.2,0.5,0.9]:
        cmp(f'eps{e}',R.epsilon_filter(logits,e),O.epsilon_filter(logits,e))
    for pen in [1.1,1.5,2.0]:
        cmp(f'rep{pen}',R.repetition_penalty(logits,ids,pen),O.repetition_penalty(logits,ids,pen))
    for pen in [0.1,0.5,1.0]:
        cmp(f'freq{pen}',R.frequency_penalty(logits,ids,pen),O.frequency_penalty(logits,ids,pen))
        cmp(f'pres{pen}',R.presence_penalty(logits,ids,pen),O.presence_penalty(logits,ids,pen))
    tk=[int(x) for x in torch.randint(0,V,(2,),generator=g).tolist()]
    cmp('suppress',R.suppress_tokens(logits,tk),O.suppress_tokens(logits,tk))
    for cl in [0,3,5]:
        cmp(f'minlen{cl}',R.min_length_filter(logits,cl,4,min(2,V-1)),O.min_length_filter(logits,cl,4,min(2,V-1)))
    cmp('renorm',R.renormalize_logits(logits),O.renormalize_logits(logits))
    for T in [0.5,1.0]:
      for k in [1,2,V]:
        for p in [0.5,0.9,1.0]:
            cmp(f'warp',R.logits_warper(logits,T,k,p),O.logits_warper(logits,T,k,p))
            cmp(f'sgw',R.sample_greedy_warped(logits,T,k,p),O.sample_greedy_warped(logits,T,k,p))

print('FILTER/LOGIT trials done, fails:', len(fails))
for f in fails[:20]: print(' ',f)
"
output
FILTER/LOGIT trials done, fails: 0

[stdout]
FILTER/LOGIT trials done, fails: 0

Compare collation functions

shell
python3 -c "
import torch
import reforacle as R
import objectives as O
fails=[]
def cmp(name, a, b, atol=1e-5, rtol=1e-4):
    try:
        if isinstance(a,tuple): a=a[0]
        if a.dtype != b.dtype: fails.append(f'{name}: dtype {a.dtype} vs {b.dtype}'); return
        if a.shape != b.shape: fails.append(f'{name}: shape {a.shape} vs {b.shape}'); return
        if not torch.allclose(a.float(), b.float(), atol=atol, rtol=rtol, equal_nan=True):
            fails.append(f'{name}: diff')
    except Exception as e: fails.append(f'{name}: EXC {e}')

for trial in range(300):
    g=torch.Generator().manual_seed(trial+1000)
    B=int(torch.randint(1,4,(1,),generator=g)); L=int(torch.randint(2,8,(1,),generator=g)); V=int(torch.randint(3,10,(1,),generator=g))
    pad=0
    ids=torch.randint(0,V,(B,L),generator=g)
    # inject some pads at end
    ids2=ids.clone()
    for b in range(B):
        nz=int(torch.randint(1,L+1,(1,),generator=g))
        ids2[b,nz:]=pad
    cmp('attn',R.attention_mask(ids2,pad),O.attention_mask(ids2,pad))
    cmp('posids',R.position_ids(ids2,pad),O.position_ids(ids2,pad))
    cmp('seqlen',R.sequence_lengths(ids2,pad),O.sequence_lengths(ids2,pad))
    am=R.attention_mask(ids2,pad)
    cmp('lasttok',R.last_token_indices(am),O.last_token_indices(am))
    cmp('causal',R.causal_mask(L),O.causal_mask(L))
    cmp('shift',R.shift_labels(ids,-100),O.shift_labels(ids,-100))
    logits3=torch.randn(B,L,V,generator=g)
    cmp('gather',R.gather_last_logits(logits3,am),O.gather_last_logits(logits3,am))
    cmp('rollout',R.rollout_last_logits(logits3,ids2,pad),O.rollout_last_logits(logits3,ids2,pad))
    pl=torch.randint(1,L+1,(B,),generator=g)
    cmp('comp',R.completion_loss_labels(ids,pl,-100),O.completion_loss_labels(ids,pl,-100))
    cmp('prepare',R.prepare_labels(ids,pl,-100),O.prepare_labels(ids,pl,-100))
    eos=1
    cmp('maskeos',R.mask_after_eos(ids,eos),O.mask_after_eos(ids,eos))
    for ml in [1,2,3]:
        cmp(f'gsm{ml}',R.generation_stop_mask(ids,eos,ml),O.generation_stop_mask(ids,eos,ml))
    for m in [1,3,L,L+2]:
        cmp(f'trunc{m}',R.truncate_left(ids,m),O.truncate_left(ids,m))
        cmp(f'rpad{m}',R.right_pad_to(ids,pad,m),O.right_pad_to(ids,pad,m))
    for mu in [2,3,4]:
        cmp(f'padmul{mu}',R.pad_to_multiple(ids,pad,mu),O.pad_to_multiple(ids,pad,mu))
    acts=torch.randint(0,V,(B,L),generator=g)
    cmp('alp',R.action_logprobs(logits3,acts),O.action_logprobs(logits3,acts))
    cmp('ent',R.entropy_from_logits(logits3),O.entropy_from_logits(logits3))
    rew=torch.randn(B,L,generator=g); mask=(torch.rand(B,L,generator=g)>0.3).float()
    cmp('r2g',R.reward_to_go(rew,mask),O.reward_to_go(rew,mask))
    for gam in [0.9,0.99,1.0,0.5]:
        cmp(f'disc{gam}',R.discount_cumsum(rew,gam),O.discount_cumsum(rew,gam))
    stop=torch.randint(0,V,(2,),generator=g)
    cmp('fsi',R.first_stop_index(ids,stop),O.first_stop_index(ids,stop))
    mask2=(torch.rand(B,V,generator=g)>0.3)
    mask2[:,0]=True
    cmp('margmax',R.masked_argmax(logits,mask2) if False else R.masked_argmax(torch.randn(B,V,generator=g),mask2), O.masked_argmax(torch.randn(B,V,generator=torch.Generator().manual_seed(trial+1000)) if False else None,None) if False else None)
    # packed
    sl=torch.randint(1,4,(B,),generator=g); tl=int(sl.sum())
    cmp('bpp',R.build_position_ids_packed(sl,tl),O.build_position_ids_packed(sl,tl))
    cmp('bpp2',R.build_position_ids_packed(sl,tl+2),O.build_position_ids_packed(sl,tl+2))

print('done fails', len(fails))
for f in fails[:30]: print(' ',f)
"
output
done fails 716
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff

[stdout]
done fails 716
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff
  rollout: diff
  margmax: EXC 'NoneType' object has no attribute 'dtype'
  lasttok: diff
  gather: diff

last token with internal zeros

shell
python3 -c "
import torch, reforacle as R
am=torch.tensor([[1.,0.,1.,0.,0.],[1.,1.,1.,0.,1.]])
print('lasttok', R.last_token_indices(am).tolist())
logits=torch.arange(2*5*2,dtype=torch.float32).reshape(2,5,2)
print('gather', R.gather_last_logits(logits,am).tolist())
ids=torch.tensor([[5,0,7,0,0],[1,2,0,3,0]])
print('rollout', R.rollout_last_logits(logits,ids,0).tolist())
"
output
lasttok [2, 4]
gather [[4.0, 5.0], [18.0, 19.0]]
rollout [[4.0, 5.0], [16.0, 17.0]]

[stdout]
lasttok [2, 4]
gather [[4.0, 5.0], [18.0, 19.0]]
rollout [[4.0, 5.0], [16.0, 17.0]]

/app/objectives.py

− removed
def last_token_indices(attn_mask):
    return attn_mask.sum(dim=-1).long() - 1


def gather_last_logits(logits, attn_mask):
    idx = attn_mask.sum(dim=-1).long() - 1
    batch = torch.arange(logits.size(0), device=logits.device)
    return logits[batch, idx]
+ added
def last_token_indices(attn_mask):
    ar = torch.arange(attn_mask.size(-1), device=attn_mask.device)
    return torch.where(attn_mask.bool(), ar, torch.full_like(ar, -1)).max(dim=-1).values


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

/app/objectives.py

− removed
def rollout_last_logits(logits, input_ids, pad_id):
    idx = (input_ids != pad_id).sum(dim=-1) - 1
    batch = torch.arange(logits.size(0), device=logits.device)
    return logits[batch, idx]
+ added
def rollout_last_logits(logits, input_ids, pad_id):
    idx = last_token_indices((input_ids != pad_id))
    batch = torch.arange(logits.size(0), device=logits.device)
    return logits[batch, idx]

Retest fixed functions

shell
python3 -c "
import torch
import reforacle as R
import objectives as O
fails=[]
def cmp(name, a, b, atol=1e-5, rtol=1e-4):
    try:
        if a.dtype != b.dtype: fails.append(f'{name}: dtype {a.dtype} vs {b.dtype}'); return
        if a.shape != b.shape: fails.append(f'{name}: shape {a.shape} vs {b.shape}'); return
        if not torch.allclose(a.float(), b.float(), atol=atol, rtol=rtol, equal_nan=True):
            fails.append(f'{name}: diff {a.tolist()} {b.tolist()}')
    except Exception as e: fails.append(f'{name}: EXC {e}')

for trial in range(300):
    g=torch.Generator().manual_seed(trial+1000)
    B=int(torch.randint(1,4,(1,),generator=g)); L=int(torch.randint(2,8,(1,),generator=g)); V=int(torch.randint(3,10,(1,),generator=g))
    ids=torch.randint(0,V,(B,L),generator=g)
    ids2=ids.clone()
    for b in range(B):
        nz=int(torch.randint(1,L+1,(1,),generator=g)); ids2[b,nz:]=0
    am=R.attention_mask(ids2,0)
    cmp('lasttok',R.last_token_indices(am),O.last_token_indices(am))
    logits3=torch.randn(B,L,V,generator=g)
    cmp('gather',R.gather_last_logits(logits3,am),O.gather_last_logits(logits3,am))
    cmp('rollout',R.rollout_last_logits(logits3,ids2,0),O.rollout_last_logits(logits3,ids2,0))
    cmp('seqlen',R.sequence_lengths(ids2,0),O.sequence_lengths(ids2,0))
    logitsm=torch.randn(B,V,generator=g)
    m=(torch.rand(B,V,generator=g)>0.4); m[:,0]=True
    cmp('margmax',R.masked_argmax(logitsm,m),O.masked_argmax(logitsm,m))
    # also int mask
    mi=m.long()
    cmp('margmax_int',R.masked_argmax(logitsm,mi),O.masked_argmax(logitsm,mi))
    mf=m.float()
    cmp('margmax_f',R.masked_argmax(logitsm,mf),O.masked_argmax(logitsm,mf))
print('done fails', len(fails))
for f in fails[:20]: print(' ',f)
"
output
done fails 62
  gather: diff [[-1.5073312520980835, -1.2262433767318726, 0.46977147459983826]] [[-1.7056041955947876, -0.19066767394542694, -1.5646264553070068]]
  rollout: diff [[-1.5073312520980835, -1.2262433767318726, 0.46977147459983826]] [[-1.7056041955947876, -0.19066767394542694, -1.5646264553070068]]
  gather: diff [[-0.25451749563217163, 1.5206695795059204, -1.6270747184753418, -1.0166531801223755, 0.5836646556854248], [0.6953851580619812, -0.4673751890659332, -0.39667850732803345, 0.1963307410478592, -0.6226924657821655], [-0.8233465552330017, -2.463714599609375, 2.0366220474243164, 1.5396276712417603, 0.506711483001709]] [[-0.25451749563217163, 1.5206695795059204, -1.6270747184753418, -1.0166531801223755, 0.5836646556854248], [0.5969623327255249, 1.180668830871582, 0.6582942008972168, 0.28955864906311035, -1.221045732498169], [-0.8233465552330017, -2.463714599609375, 2.0366220474243164, 1.5396276712417603, 0.506711483001709]]
  rollout: diff [[-0.25451749563217163, 1.5206695795059204, -1.6270747184753418, -1.0166531801223755, 0.5836646556854248], [0.6953851580619812, -0.4673751890659332, -0.39667850732803345, 0.1963307410478592, -0.6226924657821655], [-0.8233465552330017, -2.463714599609375, 2.0366220474243164, 1.5396276712417603, 0.506711483001709]] [[-0.25451749563217163, 1.5206695795059204, -1.6270747184753418, -1.0166531801223755, 0.5836646556854248], [0.5969623327255249, 1.180668830871582, 0.6582942008972168, 0.28955864906311035, -1.221045732498169], [-0.8233465552330017, -2.463714599609375, 2.0366220474243164, 1.5396276712417603, 0.506711483001709]]
  gather: diff [[-0.3800773620605469, 0.2723458409309387, 0.22695547342300415, 0.08441971987485886, 0.16896279156208038], [-0.5035187602043152, -1.8280056715011597, -0.12055995315313339, 0.18705636262893677, 0.3587751090526581], [-1.5794363021850586, 1.1280145645141602, 0.5960278511047363, 0.6684494018554688, -1.1003074645996094]] [[-0.3800773620605469, 0.2723458409309387, 0.22695547342300415, 0.08441971987485886, 0.16896279156208038], [-1.067372441291809, 0.0322505421936512, 0.3384777903556824, -0.9966886639595032, 0.5397519469261169], [-1.5794363021850586, 1.1280145645141602, 0.5960278511047363, 0.6684494018554688, -1.1003074645996094]]
  rollout: diff [[-0.3800773620605469, 0.2723458409309387, 0.22695547342300415, 0.08441971987485886, 0.16896279156208038], [-0.5035187602043152, -1.8280056715011597, -0.12055995315313339, 0.18705636262893677, 0.3587751090526581], [-1.5794363021850586, 1.1280145645141602, 0.5960278511047363, 0.6684494018554688, -1.1003074645996094]] [[-0.3800773620605469, 0.2723458409309387, 0.22695547342300415, 0.08441971987485886, 0.16896279156208038], [-1.067372441291809, 0.0322505421936512, 0.3384777903556824, -0.9966886639595032, 0.5397519469261169], [-1.5794363021850586, 1.1280145645141602, 0.5960278511047363, 0.6684494018554688, -1.1003074645996094]]
  gather: diff [[-0.7397835850715637, -0.09017974883317947, -0.5498442649841309, -0.33297863602638245], [1.0400713682174683, 0.8422277569770813, -0.904862105846405, -0.5839735269546509]] [[-0.7397835850715637, -0.09017974883317947, -0.5498442649841309, -0.33297863602638245], [-1.8087124824523926, -0.259879469871521, 0.6529757380485535, -0.31304994225502014]]
  rollout: diff [[-0.7397835850715637, -0.09017974883317947, -0.5498442649841309, -0.33297863602638245], [1.0400713682174683, 0.8422277569770813, -0.904862105846405, -0.5839735269546509]] [[-0.7397835850715637, -0.09017974883317947, -0.5498442649841309, -0.33297863602638245], [-1.8087124824523926, -0.259879469871521, 0.6529757380485535, -0.31304994225502014]]
  gather: diff [[-0.15988218784332275, -0.42961007356643677, 0.4164445400238037, 1.3355345726013184]] [[0.5760172009468079, 0.6705266237258911, -1.5619317293167114, 0.8643422722816467]]
  rollout: diff [[-0.15988218784332275, -0.42961007356643677, 0.4164445400238037, 1.3355345726013184]] [[0.5760172009468079, 0.6705266237258911, -1.5619317293167114, 0.8643422722816467]]
  gather: diff [[0.8644434213638306, 2.3115234375, -0.23842452466487885, 0.1389683485031128, 2.6364939212799072, -0.4712688624858856, 2.4827094078063965, 0.9293696284294128], [-2.044372797012329, -0.5340074896812439, 0.661652147769928, 0.4482860267162323, -1.9976638555526733, 0.3227129280567169, 0.09909894317388535, 0.5903045535087585]] [[0.8644434213638306, 2.3115234375, -0.23842452466487885, 0.1389683485031128, 2.6364939212799072, -0.4712688624858856, 2.4827094078063965, 0.9293696284294128], [-0.27697309851646423, 0.9316501617431641, 1.154234766960144, 0.3634006679058075, 1.0140272378921509, -0.7374575138092041, 0.13643352687358856, -1.3142989873886108]]
  rollout: diff [[0.8644434213638306, 2.3115234375, -0.23842452466487885, 0.1389683485031128, 2.6364939212799072, -0.4712688624858856, 2.4827094078063965, 0.9293696284294128], [-2.044372797012329, -0.5340074896812439, 0.661652147769928, 0.4482860267162323, -1.9976638555526733, 0.3227129280567169, 0.09909894317388535, 0.5903045535087585]] [[0.8644434213638306, 2.3115234375, -0.23842452466487885, 0.1389683485031128, 2.6364939212799072, -0.4712688624858856, 2.4827094078063965, 0.9293696284294128], [-0.27697309851646423, 0.9316501617431641, 1.154234766960144, 0.3634006679058075, 1.0140272378921509, -0.7374575138092041, 0.13643352687358856, -1.3142989873886108]]
  gather: diff [[-1.2892212867736816, -0.13683712482452393, -0.6789252758026123, 0.3674507141113281, 2.4607529640197754, 0.31124797463417053], [0.1555379033088684, 0.7151826024055481, -0.4527703523635864, 0.6383377313613892, 0.22653184831142426, 0.10496827960014343]] [[-1.2892212867736816, -0.13683712482452393, -0.6789252758026123, 0.3674507141113281, 2.4607529640197754, 0.31124797463417053], [-0.04339570179581642, 1.2061368227005005, -0.027847718447446823, 0.2285156100988388, -1.349202036857605, 0.6043699979782104]]
  rollout: diff [[-1.2892212867736816, -0.13683712482452393, -0.6789252758026123, 0.3674507141113281, 2.4607529640197754, 0.31124797463417053], [0.1555379033088684, 0.7151826024055481, -0.4527703523635864, 0.6383377313613892, 0.22653184831142426, 0.10496827960014343]] [[-1.2892212867736816, -0.13683712482452393, -0.6789252758026123, 0.3674507141113281, 2.4607529640197754, 0.31124797463417053], [-0.04339570179581642, 1.2061368227005005, -0.027847718447446823, 0.2285156100988388, -1.349202036857605, 0.6043699979782104]]
  gather: diff [[-1.0534507036209106, 2.7689197063446045, -0.6309090852737427], [1.1679619550704956, -1.5942761898040771, 0.4347749650478363], [0.01339460164308548, -1.4778693914413452, 0.31811830401420593]] [[1.1154446601867676, -0.4749634563922882, -0.03460348770022392], [1.1679619550704956, -1.5942761898040771, 0.4347749650478363], [0.01339460164308548, -1.4778693914413452, 0.31811830401420593]]
  rollout: diff [[-1.0534507036209106, 2.7689197063446045, -0.6309090852737427], [1.1679619550704956, -1.5942761898040771, 0.4347749650478363], [0.01339460164308548, -1.4778693914413452, 0.31811830401420593]] [[1.1154446601867676, -0.4749634563922882, -0.03460348770022392], [1.1679619550704956, -1.5942761898040771, 0.4347749650478363], [0.01339460164308548, -1.4778693914413452, 0.31811830401420593]]
  gather: diff [[-0.9157993197441101, -0.08487596362829208, -0.9139502048492432, 0.65134596824646], [-0.7970607876777649, 2.76434588432312, 0.21474957466125488, 0.9075921177864075], [1.050678014755249, 0.7621136903762817, -1.8566186428070068, 0.7022791504859924]] [[-0.9157993197441101, -0.08487596362829208, -0.9139502048492432, 0.65134596824646], [-0.7970607876777649, 2.76434588432312, 0.21474957466125488, 0.9075921177864075], [-0.045340586453676224, 0.6977095007896423, -0.21585851907730103, -0.6049559712409973]]
  rollout: diff [[-0.9157993197441101, -0.08487596362829208, -0.9139502048492432, 0.65134596824646], [-0.7970607876777649, 2.76434588432312, 0.21474957466125488, 0.9075921177864075], [1.050678014755249, 0.7621136903762817, -1.8566186428070068, 0.7022791504859924]] [[-0.9157993197441101, -0.08487596362829208, -0.9139502048492432, 0.65134596824646], [-0.7970607876777649, 2.76434588432312, 0.21474957466125488, 0.9075921177864075], [-0.045340586453676224, 0.6977095007896423, -0.21585851907730103, -0.6049559712409973]]
  gather: diff [[-1.8066365718841553, -0.1278151273727417, 0.8762345910072327, -0.7333471775054932, 0.46251946687698364, -0.42460358142852783, -1.1208643913269043, -1.6713074445724487, 0.45846855640411377], [0.4334733784198761, -0.28486311435699463, 1.4628337621688843, -1.0664767026901245, 0.6954727172851562, -0.786920964717865, -1.1808123588562012, -0.7880973219871521, -0.6179247498512268], [-0.736928403377533, 2.0226821899414062, 1.4570287466049194, 0.28779369592666626, -0.9521523714065552, -0.5665543079376221, 1.3453789949417114, -0.27101340889930725, 0.6359113454818726]] [[1.1064019203186035, 0.08809513598680496, -0.2578117251396179, 0.3963477313518524, -0.08793409913778305, 0.443652480840683, -1.1620314121246338, -0.35910001397132874, 0.2007512003183365], [-1.8392595052719116, -0.5155808925628662, 0.5831047892570496, 0.08208966255187988, -0.8998916149139404, 0.35301101207733154, -1.9619250297546387, -0.7747390866279602, -0.08775407820940018], [-0.736928403377533, 2.0226821899414062, 1.4570287466049194, 0.28779369592666626, -0.9521523714065552, -0.5665543079376221, 1.3453789949417114, -0.27101340889930725, 0.6359113454818726]]
  rollout: diff [[-1.8066365718841553, -0.1278151273727417, 0.8762345910072327, -0.7333471775054932, 0.46251946687698364, -0.42460358142852783, -1.1208643913269043, -1.6713074445724487, 0.45846855640411377], [0.4334733784198761, -0.28486311435699463, 1.4628337621688843, -1.0664767026901245, 0.6954727172851562, -0.786920964717865, -1.1808123588562012, -0.7880973219871521, -0.6179247498512268], [-0.736928403377533, 2.0226821899414062, 1.4570287466049194, 0.28779369592666626, -0.9521523714065552, -0.5665543079376221, 1.3453789949417114, -0.27101340889930725, 0.6359113454818726]] [[1.1064019203186035, 0.08809513598680496, -0.2578117251396179, 0.3963477313518524, -0.08793409913778305, 0.443652480840683, -1.1620314121246338, -0.35910001397132874, 0.2007512003183365], [-1.8392595052719116, -0.5155808925628662, 0.5831047892570496, 0.08208966255187988, -0.8998916149139404, 0.35301101207733154, -1.9619250297546387, -0.7747390866279602, -0.08775407820940018], [-0.736928403377533, 2.0226821899414062, 1.4570287466049194, 0.28779369592666626, -0.9521523714065552, -0.5665543079376221, 1.3453789949417114, -0.27101340889930725, 0.6359113454818726]]

[stdout]
done fails 62
  gather: diff [[-1.5073312520980835, -1.2262433767318726, 0.46977147459983826]] [[-1.7056041955947876, -0.19066767394542694, -1.5646264553070068]]
  rollout: diff [[-1.5073312520980835, -1.2262433767318726, 0.46977147459983826]] [[-1.7056041955947876, -0.19066767394542694, -1.5646264553070068]]
  gather: diff [[-0.25451749563217163, 1.5206695795059204, -1.6270747184753418, -1.0166531801223755, 0.5836646556854248], [0.6953851580619812, -0.4673751890659332, -0.39667850732803345, 0.1963307410478592, -0.6226924657821655], [-0.8233465552330017, -2.463714599609375, 2.0366220474243164, 1.5396276712417603, 0.506711483001709]] [[-0.25451749563217163, 1.5206695795059204, -1.6270747184753418, -1.0166531801223755, 0.5836646556854248], [0.5969623327255249, 1.180668830871582, 0.6582942008972168, 0.28955864906311035, -1.221045732498169], [-0.8233465552330017, -2.463714599609375, 2.0366220474243164, 1.5396276712417603, 0.506711483001709]]
  rollout: diff [[-0.25451749563217163, 1.5206695795059204, -1.6270747184753418, -1.0166531801223755, 0.5836646556854248], [0.6953851580619812, -0.4673751890659332, -0.39667850732803345, 0.1963307410478592, -0.6226924657821655], [-0.8233465552330017, -2.463714599609375, 2.0366220474243164, 1.5396276712417603, 0.506711483001709]] [[-0.25451749563217163, 1.5206695795059204, -1.6270747184753418, -1.0166531801223755, 0.5836646556854248], [0.5969623327255249, 1.180668830871582, 0.6582942008972168, 0.28955864906311035, -1.221045732498169], [-0.8233465552330017, -2.463714599609375, 2.0366220474243164, 1.5396276712417603, 0.506711483001709]]
  gather: diff [[-0.3800773620605469, 0.2723458409309387, 0.22695547342300415, 0.08441971987485886, 0.16896279156208038], [-0.5035187602043152, -1.8280056715011597, -0.12055995315313339, 0.18705636262893677, 0.3587751090526581], [-1.5794363021850586, 1.1280145645141602, 0.5960278511047363, 0.6684494018554688, -1.1003074645996094]] [[-0.3800773620605469, 0.2723458409309387, 0.22695547342300415, 0.08441971987485886, 0.16896279156208038], [-1.067372441291809, 0.0322505421936512, 0.3384777903556824, -0.9966886639595032, 0.5397519469261169], [-1.5794363021850586, 1.1280145645141602, 0.5960278511047363, 0.6684494018554688, -1.1003074645996094]]
  rollout: diff [[-0.3800773620605469, 0.2723458409309387, 0.22695547342300415, 0.08441971987485886, 0.16896279156208038], [-0.5035187602043152, -1.8280056715011597, -0.12055995315313339, 0.18705636262893677, 0.3587751090526581], [-1.5794363021850586, 1.1280145645141602, 0.5960278511047363, 0.6684494018554688, -1.1003074645996094]] [[-0.3800773620605469, 0.2723458409309387, 0.22695547342300415, 0.08441971987485886, 0.16896279156208038], [-1.067372441291809, 0.0322505421936512, 0.3384777903556824, -0.9966886639595032, 0.5397519469261169], [-1.5794363021850586, 1.1280145645141602, 0.5960278511047363, 0.6684494018554688, -1.1003074645996094]]
  gather: diff [[-0.7397835850715637, -0.09017974883317947, -0.5498442649841309, -0.33297863602638245], [1.0400713682174683, 0.8422277569770813, -0.904862105846405, -0.5839735269546509]] [[-0.7397835850715637, -0.09017974883317947, -0.5498442649841309, -0.33297863602638245], [-1.8087124824523926, -0.259879469871521, 0.6529757380485535, -0.31304994225502014]]
  rollout: diff [[-0.7397835850715637, -0.09017974883317947, -0.5498442649841309, -0.33297863602638245], [1.0400713682174683, 0.8422277569770813, -0.904862105846405, -0.5839735269546509]] [[-0.7397835850715637, -0.09017974883317947, -0.5498442649841309, -0.33297863602638245], [-1.8087124824523926, -0.259879469871521, 0.6529757380485535, -0.31304994225502014]]
  gather: diff [[-0.15988218784332275, -0.42961007356643677, 0.4164445400238037, 1.3355345726013184]] [[0.5760172009468079, 0.6705266237258911, -1.5619317293167114, 0.8643422722816467]]
  rollout: diff [[-0.15988218784332275, -0.42961007356643677, 0.4164445400238037, 1.3355345726013184]] [[0.5760172009468079, 0.6705266237258911, -1.5619317293167114, 0.8643422722816467]]
  gather: diff [[0.8644434213638306, 2.3115234375, -0.23842452466487885, 0.1389683485031128, 2.6364939212799072, -0.4712688624858856, 2.4827094078063965, 0.9293696284294128], [-2.044372797012329, -0.5340074896812439, 0.661652147769928, 0.4482860267162323, -1.9976638555526733, 0.3227129280567169, 0.09909894317388535, 0.5903045535087585]] [[0.8644434213638306, 2.3115234375, -0.23842452466487885, 0.1389683485031128, 2.6364939212799072, -0.4712688624858856, 2.4827094078063965, 0.9293696284294128], [-0.27697309851646423, 0.9316501617431641, 1.154234766960144, 0.3634006679058075, 1.0140272378921509, -0.7374575138092041, 0.13643352687358856, -1.3142989873886108]]
  rollout: diff [[0.8644434213638306, 2.3115234375, -0.23842452466487885, 0.1389683485031128, 2.6364939212799072, -0.4712688624858856, 2.4827094078063965, 0.9293696284294128], [-2.044372797012329, -0.5340074896812439, 0.661652147769928, 0.4482860267162323, -1.9976638555526733, 0.3227129280567169, 0.09909894317388535, 0.5903045535087585]] [[0.8644434213638306, 2.3115234375, -0.23842452466487885, 0.1389683485031128, 2.6364939212799072, -0.4712688624858856, 2.4827094078063965, 0.9293696284294128], [-0.27697309851646423, 0.9316501617431641, 1.154234766960144, 0.3634006679058075, 1.0140272378921509, -0.7374575138092041, 0.13643352687358856, -1.3142989873886108]]
  gather: diff [[-1.2892212867736816, -0.13683712482452393, -0.6789252758026123, 0.3674507141113281, 2.4607529640197754, 0.31124797463417053], [0.1555379033088684, 0.7151826024055481, -0.4527703523635864, 0.6383377313613892, 0.22653184831142426, 0.10496827960014343]] [[-1.2892212867736816, -0.13683712482452393, -0.6789252758026123, 0.3674507141113281, 2.4607529640197754, 0.31124797463417053], [-0.04339570179581642, 1.2061368227005005, -0.027847718447446823, 0.2285156100988388, -1.349202036857605, 0.6043699979782104]]
  rollout: diff [[-1.2892212867736816, -0.13683712482452393, -0.6789252758026123, 0.3674507141113281, 2.4607529640197754, 0.31124797463417053], [0.1555379033088684, 0.7151826024055481, -0.4527703523635864, 0.6383377313613892, 0.22653184831142426, 0.10496827960014343]] [[-1.2892212867736816, -0.13683712482452393, -0.6789252758026123, 0.3674507141113281, 2.4607529640197754, 0.31124797463417053], [-0.04339570179581642, 1.2061368227005005, -0.027847718447446823, 0.2285156100988388, -1.349202036857605, 0.6043699979782104]]
  gather: diff [[-1.0534507036209106, 2.7689197063446045, -0.6309090852737427], [1.1679619550704956, -1.5942761898040771, 0.4347749650478363], [0.01339460164308548, -1.4778693914413452, 0.31811830401420593]] [[1.1154446601867676, -0.4749634563922882, -0.03460348770022392], [1.1679619550704956, -1.5942761898040771, 0.4347749650478363], [0.01339460164308548, -1.4778693914413452, 0.31811830401420593]]
  rollout: diff [[-1.0534507036209106, 2.7689197063446045, -0.6309090852737427], [1.1679619550704956, -1.5942761898040771, 0.4347749650478363], [0.01339460164308548, -1.4778693914413452, 0.31811830401420593]] [[1.1154446601867676, -0.4749634563922882, -0.03460348770022392], [1.1679619550704956, -1.5942761898040771, 0.4347749650478363], [0.01339460164308548, -1.4778693914413452, 0.31811830401420593]]
  gather: diff [[-0.9157993197441101, -0.08487596362829208, -0.9139502048492432, 0.65134596824646], [-0.7970607876777649, 2.76434588432312, 0.21474957466125488, 0.9075921177864075], [1.050678014755249, 0.7621136903762817, -1.8566186428070068, 0.7022791504859924]] [[-0.9157993197441101, -0.08487596362829208, -0.9139502048492432, 0.65134596824646], [-0.7970607876777649, 2.76434588432312, 0.21474957466125488, 0.9075921177864075], [-0.045340586453676224, 0.6977095007896423, -0.21585851907730103, -0.6049559712409973]]
  rollout: diff [[-0.9157993197441101, -0.08487596362829208, -0.9139502048492432, 0.65134596824646], [-0.7970607876777649, 2.76434588432312, 0.21474957466125488, 0.9075921177864075], [1.050678014755249, 0.7621136903762817, -1.8566186428070068, 0.7022791504859924]] [[-0.9157993197441101, -0.08487596362829208, -0.9139502048492432, 0.65134596824646], [-0.7970607876777649, 2.76434588432312, 0.21474957466125488, 0.9075921177864075], [-0.045340586453676224, 0.6977095007896423, -0.21585851907730103, -0.6049559712409973]]
  gather: diff [[-1.8066365718841553, -0.1278151273727417, 0.8762345910072327, -0.7333471775054932, 0.46251946687698364, -0.42460358142852783, -1.1208643913269043, -1.6713074445724487, 0.45846855640411377], [0.4334733784198761, -0.28486311435699463, 1.4628337621688843, -1.0664767026901245, 0.6954727172851562, -0.786920964717865, -1.1808123588562012, -0.7880973219871521, -0.6179247498512268], [-0.736928403377533, 2.0226821899414062, 1.4570287466049194, 0.28779369592666626, -0.9521523714065552, -0.5665543079376221, 1.3453789949417114, -0.27101340889930725, 0.6359113454818726]] [[1.1064019203186035, 0.08809513598680496, -0.2578117251396179, 0.3963477313518524, -0.08793409913778305, 0.443652480840683, -1.1620314121246338, -0.35910001397132874, 0.2007512003183365], [-1.8392595052719116, -0.5155808925628662, 0.5831047892570496, 0.08208966255187988, -0.8998916149139404, 0.35301101207733154, -1.9619250297546387, -0.7747390866279602, -0.08775407820940018], [-0.736928403377533, 2.0226821899414062, 1.4570287466049194, 0.28779369592666626, -0.9521523714065552, -0.5665543079376221, 1.3453789949417114, -0.27101340889930725, 0.6359113454818726]]
  rollout: diff [[-1.8066365718841553, -0.1278151273727417, 0.8762345910072327, -0.7333471775054932, 0.46251946687698364, -0.42460358142852783, -1.1208643913269043, -1.6713074445724487, 0.45846855640411377], [0.4334733784198761, -0.28486311435699463, 1.4628337621688843, -1.0664767026901245, 0.6954727172851562, -0.786920964717865, -1.1808123588562012, -0.7880973219871521, -0.6179247498512268], [-0.736928403377533, 2.0226821899414062, 1.4570287466049194, 0.28779369592666626, -0.9521523714065552, -0.5665543079376221, 1.3453789949417114, -0.27101340889930725, 0.6359113454818726]] [[1.1064019203186035, 0.08809513598680496, -0.2578117251396179, 0.3963477313518524, -0.08793409913778305, 0.443652480840683, -1.1620314121246338, -0.35910001397132874, 0.2007512003183365], [-1.8392595052719116, -0.5155808925628662, 0.5831047892570496, 0.08208966255187988, -0.8998916149139404, 0.35301101207733154, -1.9619250297546387, -0.7747390866279602, -0.08775407820940018], [-0.736928403377533, 2.0226821899414062, 1.4570287466049194, 0.28779369592666626, -0.9521523714065552, -0.5665543079376221, 1.3453789949417114, -0.27101340889930725, 0.6359113454818726]]

Debug gather index

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(1000)
B=int(torch.randint(1,4,(1,),generator=g)); L=int(torch.randint(2,8,(1,),generator=g)); V=int(torch.randint(3,10,(1,),generator=g))
ids=torch.randint(0,V,(B,L),generator=g)
ids2=ids.clone()
for b in range(B):
    nz=int(torch.randint(1,L+1,(1,),generator=g)); ids2[b,nz:]=0
am=R.attention_mask(ids2,0)
print('ids2',ids2.tolist())
print('am',am.tolist())
print('R lasttok', R.last_token_indices(am).tolist())
print('O lasttok', O.last_token_indices(am).tolist())
logits3=torch.randn(B,L,V,generator=g)
print('R gather idx?')
rg=R.gather_last_logits(logits3,am)
og=O.gather_last_logits(logits3,am)
for b in range(B):
    for i in range(L):
        if torch.allclose(logits3[b,i],rg[b]): print('row',b,'R uses idx',i)
        if torch.allclose(logits3[b,i],og[b]): print('row',b,'O uses idx',i)
"
output
ids2 [[3, 1, 4, 0, 0, 0, 0], [0, 0, 4, 0, 0, 0, 0]]
am [[1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]]
R lasttok [2, 2]
O lasttok [2, 2]
R gather idx?
row 0 R uses idx 2
row 0 O uses idx 2
row 1 R uses idx 2
row 1 O uses idx 2

[stdout]
ids2 [[3, 1, 4, 0, 0, 0, 0], [0, 0, 4, 0, 0, 0, 0]]
am [[1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0]]
R lasttok [2, 2]
O lasttok [2, 2]
R gather idx?
row 0 R uses idx 2
row 0 O uses idx 2
row 1 R uses idx 2
row 1 O uses idx 2

Recount gather fails

shell
python3 -c "
import torch
import reforacle as R
import objectives as O
fails={}
def cmp(name, a, b, atol=1e-5, rtol=1e-4):
    if a.dtype != b.dtype or a.shape!=b.shape or not torch.allclose(a.float(),b.float(),atol=atol,rtol=rtol,equal_nan=True):
        fails[name]=fails.get(name,0)+1

for trial in range(300):
    g=torch.Generator().manual_seed(trial+1000)
    B=int(torch.randint(1,4,(1,),generator=g)); L=int(torch.randint(2,8,(1,),generator=g)); V=int(torch.randint(3,10,(1,),generator=g))
    ids=torch.randint(0,V,(B,L),generator=g)
    ids2=ids.clone()
    for b in range(B):
        nz=int(torch.randint(1,L+1,(1,),generator=g)); ids2[b,nz:]=0
    am=R.attention_mask(ids2,0)
    logits3=torch.randn(B,L,V,generator=g)
    cmp('lasttok',R.last_token_indices(am),O.last_token_indices(am))
    cmp('gather',R.gather_last_logits(logits3,am),O.gather_last_logits(logits3,am))
    cmp('rollout',R.rollout_last_logits(logits3,ids2,0),O.rollout_last_logits(logits3,ids2,0))
print(fails)
"
output
{'gather': 31, 'rollout': 31}

[stdout]
{'gather': 31, 'rollout': 31}

Find gather failing trial

shell
python3 -c "
import torch, reforacle as R, objectives as O
for trial in range(300):
    g=torch.Generator().manual_seed(trial+1000)
    B=int(torch.randint(1,4,(1,),generator=g)); L=int(torch.randint(2,8,(1,),generator=g)); V=int(torch.randint(3,10,(1,),generator=g))
    ids=torch.randint(0,V,(B,L),generator=g)
    ids2=ids.clone()
    for b in range(B):
        nz=int(torch.randint(1,L+1,(1,),generator=g)); ids2[b,nz:]=0
    am=R.attention_mask(ids2,0)
    logits3=torch.randn(B,L,V,generator=g)
    rg=R.gather_last_logits(logits3,am)
    og=O.gather_last_logits(logits3,am)
    if not torch.allclose(rg,og):
        print('trial',trial,'ids2',ids2.tolist())
        print('am',am.tolist())
        print('R lasttok', R.last_token_indices(am).tolist())
        for b in range(B):
            ri=oi=None
            for i in range(L):
                if torch.allclose(logits3[b,i],rg[b]): ri=i
                if torch.allclose(logits3[b,i],og[b]): oi=i
            print('row',b,'R idx',ri,'O idx',oi)
        break
"
output
trial 8 ids2 [[0, 0, 0, 0, 0]]
am [[0.0, 0.0, 0.0, 0.0, 0.0]]
R lasttok [-1]
row 0 R idx 0 O idx 4

[stdout]
trial 8 ids2 [[0, 0, 0, 0, 0]]
am [[0.0, 0.0, 0.0, 0.0, 0.0]]
R lasttok [-1]
row 0 R idx 0 O idx 4

/app/objectives.py

− removed
def gather_last_logits(logits, attn_mask):
    idx = last_token_indices(attn_mask)
    batch = torch.arange(logits.size(0), device=logits.device)
    return logits[batch, idx]
+ added
def gather_last_logits(logits, attn_mask):
    idx = last_token_indices(attn_mask).clamp(min=0)
    batch = torch.arange(logits.size(0), device=logits.device)
    return logits[batch, idx]

/app/objectives.py

− removed
    idx = last_token_indices((input_ids != pad_id))
    batch = torch.arange(logits.size(0), device=logits.device)
    return logits[batch, idx]
+ added
    idx = last_token_indices((input_ids != pad_id)).clamp(min=0)
    batch = torch.arange(logits.size(0), device=logits.device)
    return logits[batch, idx]

/app/objectives.py

− removed
    labels = input_ids.clone()
    ar = torch.arange(input_ids.size(-1), device=input_ids.device)
    mask = ar.unsqueeze(0) < prompt_lengths.unsqueeze(-1)
    labels[mask] = ignore_index
    return labels
+ added
    labels = input_ids.clone()
    prompt_lengths = torch.as_tensor(prompt_lengths, device=input_ids.device)
    ar = torch.arange(input_ids.size(-1), device=input_ids.device)
    mask = ar.unsqueeze(0) < prompt_lengths.unsqueeze(-1)
    labels[mask] = ignore_index
    return labels

/app/objectives.py

− removed
    labels = shift_labels(input_ids, ignore_index)
    ar = torch.arange(input_ids.size(-1), device=input_ids.device)
    mask = ar.unsqueeze(0) < (prompt_lengths.unsqueeze(-1) - 1)
+ added
    labels = shift_labels(input_ids, ignore_index)
    prompt_lengths = torch.as_tensor(prompt_lengths, device=input_ids.device)
    ar = torch.arange(input_ids.size(-1), device=input_ids.device)
    mask = ar.unsqueeze(0) < (prompt_lengths.unsqueeze(-1) - 1)

/app/objectives.py

− removed
    is_stop = (input_ids.unsqueeze(-1) == stop_ids).any(dim=-1)
+ added
    stop_ids = torch.as_tensor(stop_ids, device=input_ids.device)
    is_stop = (input_ids.unsqueeze(-1) == stop_ids).any(dim=-1)

/app/objectives.py

− removed
    parts = [torch.arange(int(l), dtype=torch.long) for l in seq_lengths.tolist()]
+ added
    parts = [torch.arange(int(l), dtype=torch.long) for l in torch.as_tensor(seq_lengths).tolist()]

Full 37-function verification

shell
python3 -c "
import torch
import reforacle as R
import objectives as O
fails={}
def cmp(name, ta, tb, atol=1e-5, rtol=1e-4):
    try:
        a=ta(); b=tb()
        if isinstance(a,torch.Tensor):
            if a.dtype!=b.dtype: fails.setdefault(name,[]).append(('dtype',str(a.dtype),str(b.dtype))); return
            if a.shape!=b.shape: fails.setdefault(name,[]).append(('shape',a.shape,b.shape)); return
            if not torch.allclose(a.float(),b.float(),atol=atol,rtol=rtol,equal_nan=True):
                fails.setdefault(name,[]).append('val')
    except Exception as e:
        fails.setdefault(name,[]).append(('EXC',str(e)))

import random
for trial in range(500):
    g=torch.Generator().manual_seed(trial)
    B=int(torch.randint(1,5,(1,),generator=g)); V=int(torch.randint(2,12,(1,),generator=g)); L=int(torch.randint(1,9,(1,),generator=g))
    dt = torch.float64 if trial%3==0 else torch.float32
    logits=torch.randn(B,V,generator=g).to(dt)
    logits3=torch.randn(B,L,V,generator=g).to(dt)
    pad=int(torch.randint(0,V,(1,),generator=g))
    ids=torch.randint(0,V,(B,L),generator=g)
    ids2=ids.clone()
    for b in range(B):
        nz=int(torch.randint(0,L+1,(1,),generator=g)); ids2[b,nz:]=pad
    cmp('greedy_token',lambda:R.greedy_token(logits),lambda:O.greedy_token(logits))
    for T in [0.0,0.3,0.7,1.0,2.5]:
        cmp('temperature_scale',lambda T=T:R.temperature_scale(logits,T),lambda T=T:O.temperature_scale(logits,T))
    for k in [1,2,max(1,V-1),V,V+3]:
        cmp('top_k_filter',lambda k=k:R.top_k_filter(logits,k),lambda k=k:O.top_k_filter(logits,k))
        cmp('topk_indices',lambda k=k:R.topk_indices(logits,min(k,V)),lambda k=k:O.topk_indices(logits,min(k,V)))
    for p in [0.0,0.1,0.5,0.9,1.0]:
        cmp('top_p_filter',lambda p=p:R.top_p_filter(logits,p),lambda p=p:O.top_p_filter(logits,p))
        cmp('min_p_filter',lambda p=p:R.min_p_filter(logits,p),lambda p=p:O.min_p_filter(logits,p))
        cmp('epsilon_filter',lambda p=p:R.epsilon_filter(logits,p*0.3+0.01),lambda p=p:O.epsilon_filter(logits,p*0.3+0.01))
    for pen in [1.0,1.2,1.8]:
        cmp('repetition_penalty',lambda pen=pen:R.repetition_penalty(logits,ids,pen),lambda pen=pen:O.repetition_penalty(logits,ids,pen))
    for pen in [0.0,0.3,0.9]:
        cmp('frequency_penalty',lambda pen=pen:R.frequency_penalty(logits,ids,pen),lambda pen=pen:O.frequency_penalty(logits,ids,pen))
        cmp('presence_penalty',lambda pen=pen:R.presence_penalty(logits,ids,pen),lambda pen=pen:O.presence_penalty(logits,ids,pen))
    tk=list({int(x) for x in torch.randint(0,V,(3,),generator=g).tolist()})
    cmp('suppress_tokens',lambda:R.suppress_tokens(logits,tk),lambda:O.suppress_tokens(logits,tk))
    for cl in [0,2,4,7]:
        cmp('min_length_filter',lambda cl=cl:R.min_length_filter(logits,cl,4,pad),lambda cl=cl:O.min_length_filter(logits,cl,4,pad))
    cmp('renormalize_logits',lambda:R.renormalize_logits(logits),lambda:O.renormalize_logits(logits))
    cmp('attention_mask',lambda:R.attention_mask(ids2,pad),lambda:O.attention_mask(ids2,pad))
    cmp('position_ids',lambda:R.position_ids(ids2,pad),lambda:O.position_ids(ids2,pad))
    cmp('causal_mask',lambda:R.causal_mask(L),lambda:O.causal_mask(L))
    cmp('shift_labels',lambda:R.shift_labels(ids2,-100),lambda:O.shift_labels(ids2,-100))
    pl=torch.randint(0,L+1,(B,),generator=g)
    cmp('completion_loss_labels',lambda:R.completion_loss_labels(ids2,pl,-100),lambda:O.completion_loss_labels(ids2,pl,-100))
    cmp('prepare_labels',lambda:R.prepare_labels(ids2,pl,-100),lambda:O.prepare_labels(ids2,pl,-100))
    eos=int(torch.randint(0,V,(1,),generator=g))
    cmp('mask_after_eos',lambda:R.mask_after_eos(ids2,eos),lambda:O.mask_after_eos(ids2,eos))
    for ml in [0,1,2,L]:
        cmp('generation_stop_mask',lambda ml=ml:R.generation_stop_mask(ids2,eos,ml),lambda ml=ml:O.generation_stop_mask(ids2,eos,ml))
    cmp('sequence_lengths',lambda:R.sequence_lengths(ids2,pad),lambda:O.sequence_lengths(ids2,pad))
    am=R.attention_mask(ids2,pad)
    cmp('last_token_indices',lambda:R.last_token_indices(am),lambda:O.last_token_indices(am))
    cmp('gather_last_logits',lambda:R.gather_last_logits(logits3,am),lambda:O.gather_last_logits(logits3,am))
    cmp('rollout_last_logits',lambda:R.rollout_last_logits(logits3,ids2,pad),lambda:O.rollout_last_logits(logits3,ids2,pad))
    for m in [1,3,L,L+2]:
        cmp('truncate_left',lambda m=m:R.truncate_left(ids2,m),lambda m=m:O.truncate_left(ids2,m))
        cmp('right_pad_to',lambda m=m:R.right_pad_to(ids2,pad,m),lambda m=m:O.right_pad_to(ids2,pad,m))
    for mu in [1,2,3,5]:
        cmp('pad_to_multiple',lambda mu=mu:R.pad_to_multiple(ids2,pad,mu),lambda mu=mu:O.pad_to_multiple(ids2,pad,mu))
    acts=torch.randint(0,V,(B,L),generator=g)
    cmp('action_logprobs',lambda:R.action_logprobs(logits3,acts),lambda:O.action_logprobs(logits3,acts))
    cmp('entropy_from_logits',lambda:R.entropy_from_logits(logits3),lambda:O.entropy_from_logits(logits3))
    rew=torch.randn(B,L,generator=g).to(dt); mask=(torch.rand(B,L,generator=g)>0.3).to(dt)
    cmp('reward_to_go',lambda:R.reward_to_go(rew,mask),lambda:O.reward_to_go(rew,mask))
    for gam in [0.0,0.5,0.9,1.0]:
        cmp('discount_cumsum',lambda gam=gam:R.discount_cumsum(rew,gam),lambda gam=gam:O.discount_cumsum(rew,gam))
    sl=torch.randint(0,4,(B,),generator=g); tl=int(sl.sum())+int(torch.randint(0,3,(1,),generator=g))
    cmp('build_position_ids_packed',lambda:R.build_position_ids_packed(sl,tl),lambda:O.build_position_ids_packed(sl,tl))
    for T in [0.5,1.0]:
      for k in [1,2,V]:
        for p in [0.5,0.9,1.0]:
            cmp('logits_warper',lambda T=T,k=k,p=p:R.logits_warper(logits,T,k,p),lambda T=T,k=k,p=p:O.logits_warper(logits,T,k,p))
            cmp('sample_greedy_warped',lambda T=T,k=k,p=p:R.sample_greedy_warped(logits,T,k,p),lambda T=T,k=k,p=p:O.sample_greedy_warped(logits,T,k,p))
    stop=torch.randint(0,V,(2,),generator=g)
    cmp('first_stop_index',lambda:R.first_stop_index(ids2,stop),lambda:O.first_stop_index(ids2,stop))
    mm=(torch.rand(B,V,generator=g)>0.4); mm[:,0]=True
    cmp('masked_argmax',lambda:R.masked_argmax(logits,mm),lambda:O.masked_argmax(logits,mm))

print('TOTAL fail categories:', len(fails))
for k,v in fails.items(): print(' ',k, len(v), v[:3])
"
output
TOTAL fail categories: 0

[stdout]
TOTAL fail categories: 0

Edge dimensionality/dtype tests

shell
python3 -c "
import torch, reforacle as R, objectives as O
def chk(name,a,b):
    ok = (a.dtype==b.dtype and a.shape==b.shape and torch.allclose(a.float(),b.float(),atol=1e-5,rtol=1e-4,equal_nan=True))
    print(('OK ' if ok else 'FAIL ')+name, '' if ok else (a.tolist(),b.tolist()))
# 1D logits
l1=torch.randn(6)
chk('greedy1d',R.greedy_token(l1),O.greedy_token(l1))
chk('temp1d',R.temperature_scale(l1,0.0),O.temperature_scale(l1,0.0))
chk('topk1d',R.top_k_filter(l1,2),O.top_k_filter(l1,2))
chk('topp1d',R.top_p_filter(l1,0.8),O.top_p_filter(l1,0.8))
chk('minp1d',R.min_p_filter(l1,0.3),O.min_p_filter(l1,0.3))
chk('eps1d',R.epsilon_filter(l1,0.1),O.epsilon_filter(l1,0.1))
chk('renorm1d',R.renormalize_logits(l1),O.renormalize_logits(l1))
chk('ent1d',R.entropy_from_logits(l1),O.entropy_from_logits(l1))
# 1D reward / discount
r1=torch.randn(5); m1=(torch.rand(5)>0.3).float()
chk('r2g1d',R.reward_to_go(r1,m1),O.reward_to_go(r1,m1))
chk('disc1d',R.discount_cumsum(r1,0.9),O.discount_cumsum(r1,0.9))
# int32 ids
ids=torch.randint(0,5,(2,4),dtype=torch.int32)
chk('shift_i32',R.shift_labels(ids,-100),O.shift_labels(ids,-100))
chk('attn_i32',R.attention_mask(ids,0),O.attention_mask(ids,0))
chk('pad_i32',R.pad_to_multiple(ids,0,3),O.pad_to_multiple(ids,0,3))
# neg temperature
chk('tempneg',R.temperature_scale(torch.randn(2,5),-1.5),O.temperature_scale(torch.randn(2,5),-1.5) if False else O.temperature_scale(torch.tensor([[1.,2.],[3.,4.]]),-1.5)) if False else None
tl=torch.tensor([[1.,2.,3.],[4.,5.,6.]])
chk('tempneg',R.temperature_scale(tl,-2.0),O.temperature_scale(tl,-2.0))
# 3D logits for filters?
l3=torch.randn(2,3,6)
chk('topk3d',R.top_k_filter(l3,2),O.top_k_filter(l3,2))
chk('greedy3d',R.greedy_token(l3),O.greedy_token(l3))
chk('actlp3d',R.action_logprobs(l3,torch.randint(0,6,(2,3))),O.action_logprobs(l3,torch.randint(0,6,(2,3))) if False else O.action_logprobs(l3,torch.zeros(2,3,dtype=torch.long)))
" 2>&1 | grep -v '^OK.*None'
output
OK greedy1d 
OK temp1d 
OK topk1d 
OK topp1d 
OK minp1d 
OK eps1d 
OK renorm1d 
OK ent1d 
OK r2g1d 
OK disc1d 
OK shift_i32 
OK attn_i32 
OK pad_i32 
OK tempneg 
OK topk3d 
OK greedy3d 
FAIL actlp3d ([[-1.415848731994629, -3.0165419578552246, -2.5094218254089355], [-2.671379327774048, -2.896944999694824, -2.5633933544158936]], [[-3.987278938293457, -2.131946563720703, -1.762450098991394], [-2.671379327774048, -0.36069318652153015, -2.5237371921539307]])

[stdout]
OK greedy1d 
OK temp1d 
OK topk1d 
OK topp1d 
OK minp1d 
OK eps1d 
OK renorm1d 
OK ent1d 
OK r2g1d 
OK disc1d 
OK shift_i32 
OK attn_i32 
OK pad_i32 
OK tempneg 
OK topk3d 
OK greedy3d 
FAIL actlp3d ([[-1.415848731994629, -3.0165419578552246, -2.5094218254089355], [-2.671379327774048, -2.896944999694824, -2.5633933544158936]], [[-3.987278938293457, -2.131946563720703, -1.762450098991394], [-2.671379327774048, -0.36069318652153015, -2.5237371921539307]])

Final sanity check

shell
python3 -c "
import torch, reforacle as R, objectives as O
l3=torch.randn(2,3,6); a=torch.randint(0,6,(2,3))
print('actlp3d match', torch.allclose(R.action_logprobs(l3,a),O.action_logprobs(l3,a)))
# ensure no stub left
import inspect
src=inspect.getsource(O)
print('NotImplemented present:', 'NotImplementedError' in src)
print('funcs implemented:', sum(1 for n in dir(O) if callable(getattr(O,n)) and not n.startswith('_') and n not in ('torch',)))
"
output
actlp3d match True
NotImplemented present: False
funcs implemented: 37

[stdout]
actlp3d match True
NotImplemented present: False
funcs implemented: 37

Confirm no reforacle dependency

shell
grep -c "reforacle" /app/objectives.py; echo "---imports---"; grep -E "^import|^from" /app/objectives.py
output
0
---imports---
from __future__ import annotations
import torch

[stdout]
0
---imports---
from __future__ import annotations
import torch
The reimplementation is complete. `/app/objectives.py` is a standalone module whose only dependency is `torch`, with all 37 functions implemented to reproduce `reforacle` exactly. ## Summary I recovered each function's behavior purely by probing `reforacle` and comparing, then verified my implementation matched across 500 randomized trials (varied batch sizes, vocab sizes, sequence lengths, `float32`/`float64`, `int32` ids, and edge cases) plus targeted 1D/3D and degenerate-input checks , **zero mismatches**. Key non-obvious conventions I discovered by probing (not assuming the "standard" definition): - **`temperature_scale(T=0)`** → one-hot in log-space: `0` at argmax, `-inf` elsewhere. - **`repetition_penalty`** is *count-based* (`penalty**count` per token), not applied once. - **`epsilon_filter`** keeps at least one token when all fall below ε , specifically the **last** occurrence of the max logit (via flip-argmax). - **`position_ids`** uses `cumsum(mask)-1` clamped at 0 (pad positions repeat the last real index), not the HF `masked_fill(...,1)` convention. - **`last_token_indices`** returns the *actual last non-masked index* (returns `-1` for all-pad rows), while **`gather_last_logits`/`rollout_last_logits`** additionally clamp that index to `0`. - **`generation_stop_mask`** = `mask_after_eos` (includes the eos token) with the first `min_len` positions forced to 1. - **`prepare_labels`** = shifted labels with the first `prompt_length-1` positions masked; **`completion_loss_labels`** masks the first `prompt_length` with no shift. - **`reward_to_go`** = reverse-cumsum of `rewards*mask`, then re-masked; **`logits_warper`** applies temperature → top_k → top_p in that order.
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: 32/37  (partial_score=0.8649)
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_20cbe0b6c87e49a4. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.

Trial trial_20cbe0b6c87e49a4 · verifier authoritative; classifier explanatory.