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.
SubtypeInsufficient Exploration
EvidenceTest output: "FUNCTIONS PASSED: 35/37 (partial_score=0.9459) REWARD = 0 (1 iff ALL 37 functions match the sealed reference exactly)". Agent systematically probed reforacle on ~80 randomized trials plus targeted edge cases (ties, all-pad rows, degenerate filters) and verified locally that all functions matched across tested inputs. However, 2/37 functions failed on held-out grading inputs that the agent never encountered during development. The instruction explicitly warns: "Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on." Agent logged: "All 37 functions implemented, standalone (torch-only)" after testing revealed no local mismatches, but grading found failures on unseen input combinations.
Root causeThe agent systematically reverse-engineered a sealed 37-function library through probing and extensive local testing but failed to capture all corner cases for 2 functions on held-out test inputs with strict tolerances (atol=1e-6/1e-5 for primitives). This is expected behavior for a hard reverse-engineering task where held-out inputs may expose subtle edge cases the agent did not probe.
RecommendationN/A - task is fine. This is a well-designed reverse-engineering challenge. The 94.6% pass rate (35/37) demonstrates the agent's systematic approach was sound; the 2 failures on unseen inputs reflect the fundamental difficulty of the problem, not a task specification issue. The instruction correctly acknowledges that agreement on tried inputs doesn't guarantee agreement on graded inputs.
Trajectory
Tool-by-tool agent trajectory
44 tool calls · 4 tool types · 46 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

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 greedy and temperature

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

# greedy_token
logits = torch.randn(2,6, generator=g)
print('greedy_token 2d:', R.greedy_token(logits).shape, R.greedy_token(logits).dtype, R.greedy_token(logits).tolist())
l1 = torch.randn(6, generator=g)
print('greedy_token 1d:', R.greedy_token(l1).shape, R.greedy_token(l1).tolist())
l3 = torch.randn(2,3,6, generator=g)
print('greedy_token 3d:', R.greedy_token(l3).shape, R.greedy_token(l3).tolist())

# temperature_scale
print('temp T=2:', R.temperature_scale(logits, 2.0).tolist())
print('temp T=0:', R.temperature_scale(logits, 0.0).tolist())
print('temp T=0.5:', R.temperature_scale(logits, 0.5).tolist())
"
output
greedy_token 2d: torch.Size([2]) torch.int64 [0, 1]
greedy_token 1d: torch.Size([]) 1
greedy_token 3d: torch.Size([2, 3]) [[2, 3, 0], [5, 4, 5]]
temp T=2: [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
temp T=0: [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]]
temp T=0.5: [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854, -2.7971909046173096], [0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828, -1.1932706832885742, 0.3640729784965515]]

[stdout]
greedy_token 2d: torch.Size([2]) torch.int64 [0, 1]
greedy_token 1d: torch.Size([]) 1
greedy_token 3d: torch.Size([2, 3]) [[2, 3, 0], [5, 4, 5]]
temp T=2: [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
temp T=0: [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]]
temp T=0.5: [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854, -2.7971909046173096], [0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828, -1.1932706832885742, 0.3640729784965515]]

Probe top_k and constants

shell
python3 -c "
import torch, reforacle as R
print('NEG_INF:', R.NEG_INF)
print('F:', R.F)
g = torch.Generator().manual_seed(1)
# top_k_filter
logits = torch.randn(2,6, generator=g)
print('logits:', logits.tolist())
print('top_k k=3:', R.top_k_filter(logits, 3).tolist())
print('top_k k=1:', R.top_k_filter(logits, 1).tolist())
print('top_k k=10:', R.top_k_filter(logits, 10).tolist())
"
output
NEG_INF: -inf
F: <module 'torch.nn.functional' from '/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py'>
logits: [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
top_k k=3: [[0.6613521575927734, 0.266924113035202, -inf, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -inf, -0.05825017765164375]]
top_k k=1: [[0.6613521575927734, -inf, -inf, -inf, -inf, -inf], [-inf, 0.38168391585350037, -inf, -inf, -inf, -inf]]
top_k k=10: [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]

[stdout]
NEG_INF: -inf
F: <module 'torch.nn.functional' from '/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py'>
logits: [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
top_k k=3: [[0.6613521575927734, 0.266924113035202, -inf, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -inf, -0.05825017765164375]]
top_k k=1: [[0.6613521575927734, -inf, -inf, -inf, -inf, -inf], [-inf, 0.38168391585350037, -inf, -inf, -inf, -inf]]
top_k k=10: [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]

Probe top_p

shell
python3 -c "
import torch, reforacle as R
g = torch.Generator().manual_seed(2)
logits = torch.randn(2,6, generator=g)
print('logits:', [[round(x,4) for x in r] for r in logits.tolist()])
print('softmax:', [[round(x,4) for x in r] for r in torch.softmax(logits,-1).tolist()])
print('top_p 0.9:', R.top_p_filter(logits, 0.9).tolist())
print('top_p 0.5:', R.top_p_filter(logits, 0.5).tolist())
print('top_p 0.0:', R.top_p_filter(logits, 0.0).tolist())
print('top_p 1.0:', R.top_p_filter(logits, 1.0).tolist())
"
output
logits: [[0.3923, -0.2236, -0.3195, -1.205, 1.0445, -0.6332], [0.5731, 0.5409, -0.3919, -1.0427, 1.3186, 0.7476]]
softmax: [[0.2216, 0.1197, 0.1088, 0.0449, 0.4255, 0.0795], [0.171, 0.1656, 0.0652, 0.034, 0.3605, 0.2037]]
top_p 0.9: [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -inf, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -inf, -inf, 1.318617343902588, 0.747639000415802]]
top_p 0.5: [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, 0.747639000415802]]
top_p 0.0: [[-inf, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, -inf]]
top_p 1.0: [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 1.318617343902588, 0.747639000415802]]

[stdout]
logits: [[0.3923, -0.2236, -0.3195, -1.205, 1.0445, -0.6332], [0.5731, 0.5409, -0.3919, -1.0427, 1.3186, 0.7476]]
softmax: [[0.2216, 0.1197, 0.1088, 0.0449, 0.4255, 0.0795], [0.171, 0.1656, 0.0652, 0.034, 0.3605, 0.2037]]
top_p 0.9: [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -inf, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -inf, -inf, 1.318617343902588, 0.747639000415802]]
top_p 0.5: [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, 0.747639000415802]]
top_p 0.0: [[-inf, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, -inf]]
top_p 1.0: [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 1.318617343902588, 0.747639000415802]]

Probe min_p, epsilon, renormalize

shell
python3 -c "
import torch, reforacle as R
g = torch.Generator().manual_seed(3)
logits = torch.randn(2,6, generator=g)
print('softmax:', [[round(x,4) for x in r] for r in torch.softmax(logits,-1).tolist()])
print('max prob row:', torch.softmax(logits,-1).max(-1).values.tolist())
print('min_p 0.1:', R.min_p_filter(logits, 0.1).tolist())
print('min_p 0.5:', R.min_p_filter(logits, 0.5).tolist())
print()
print('epsilon 0.1:', R.epsilon_filter(logits, 0.1).tolist())
print('epsilon 0.15:', R.epsilon_filter(logits, 0.15).tolist())
print()
print('renorm:', R.renormalize_logits(logits).tolist())
print('logsoftmax:', torch.log_softmax(logits,-1).tolist())
"
output
softmax: [[0.3511, 0.1873, 0.1719, 0.0851, 0.1647, 0.04], [0.1151, 0.2257, 0.0195, 0.2182, 0.1384, 0.2831]]
max prob row: [0.35107070207595825, 0.2831287384033203]
min_p 0.1: [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -0.6137180328369141, 0.04618244990706444, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -inf, 0.9774317741394043, 0.5220499634742737, 1.2379167079925537]]
min_p 0.5: [[0.8032760620117188, 0.17483338713645935, -inf, -inf, -inf, -inf], [-inf, 1.0111159086227417, -inf, 0.9774317741394043, -inf, 1.2379167079925537]]

epsilon 0.1: [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -inf, 0.04618244990706444, -inf], [0.3374950885772705, 1.0111159086227417, -inf, 0.9774317741394043, 0.5220499634742737, 1.2379167079925537]]
epsilon 0.15: [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -inf, 0.04618244990706444, -inf], [-inf, 1.0111159086227417, -inf, 0.9774317741394043, -inf, 1.2379167079925537]]

renorm: [[-1.0467677116394043, -1.6752103567123413, -1.7610657215118408, -2.463761806488037, -1.803861379623413, -3.2183029651641846], [-2.1622753143310547, -1.488654375076294, -3.9349498748779297, -1.5223385095596313, -1.9777202606201172, -1.261853575706482]]
logsoftmax: [[-1.0467677116394043, -1.6752103567123413, -1.7610657215118408, -2.463761806488037, -1.803861379623413, -3.2183029651641846], [-2.1622753143310547, -1.488654375076294, -3.9349498748779297, -1.5223385095596313, -1.9777202606201172, -1.261853575706482]]

[stdout]
softmax: [[0.3511, 0.1873, 0.1719, 0.0851, 0.1647, 0.04], [0.1151, 0.2257, 0.0195, 0.2182, 0.1384, 0.2831]]
max prob row: [0.35107070207595825, 0.2831287384033203]
min_p 0.1: [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -0.6137180328369141, 0.04618244990706444, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -inf, 0.9774317741394043, 0.5220499634742737, 1.2379167079925537]]
min_p 0.5: [[0.8032760620117188, 0.17483338713645935, -inf, -inf, -inf, -inf], [-inf, 1.0111159086227417, -inf, 0.9774317741394043, -inf, 1.2379167079925537]]

epsilon 0.1: [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -inf, 0.04618244990706444, -inf], [0.3374950885772705, 1.0111159086227417, -inf, 0.9774317741394043, 0.5220499634742737, 1.2379167079925537]]
epsilon 0.15: [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -inf, 0.04618244990706444, -inf], [-inf, 1.0111159086227417, -inf, 0.9774317741394043, -inf, 1.2379167079925537]]

renorm: [[-1.0467677116394043, -1.6752103567123413, -1.7610657215118408, -2.463761806488037, -1.803861379623413, -3.2183029651641846], [-2.1622753143310547, -1.488654375076294, -3.9349498748779297, -1.5223385095596313, -1.9777202606201172, -1.261853575706482]]
logsoftmax: [[-1.0467677116394043, -1.6752103567123413, -1.7610657215118408, -2.463761806488037, -1.803861379623413, -3.2183029651641846], [-2.1622753143310547, -1.488654375076294, -3.9349498748779297, -1.5223385095596313, -1.9777202606201172, -1.261853575706482]]

Edge cases for filters

shell
python3 -c "
import torch, reforacle as R
# all logits equal -> top_p removes? min_tokens behavior
logits = torch.zeros(1,5)
print('top_p 0.0 equal:', R.top_p_filter(logits, 0.0).tolist())
print('min_p 1.5 equal:', R.min_p_filter(logits, 1.5).tolist())
print('epsilon 0.5 equal:', R.epsilon_filter(logits, 0.5).tolist())  # all prob 0.2 < 0.5 -> all removed?
print('top_k 0:', R.top_k_filter(logits, 0).tolist())
# high epsilon removing all
l2 = torch.tensor([[10.0, 0.0, 0.0, 0.0, 0.0]])
print('epsilon 0.9 peaked:', R.epsilon_filter(l2, 0.9).tolist())
print('min_p 0.99 peaked:', R.min_p_filter(l2, 0.99).tolist())
"
output
Exit code 1
top_p 0.0 equal: [[0.0, -inf, -inf, -inf, -inf]]
min_p 1.5 equal: [[-inf, -inf, -inf, -inf, -inf]]
epsilon 0.5 equal: [[-inf, -inf, -inf, -inf, 0.0]]
Traceback (most recent call last):
  File "<string>", line 8, in <module>
  File "reforacle.py", line 30, in reforacle.top_k_filter
RuntimeError: The size of tensor a (5) must match the size of tensor b (0) at non-singleton dimension 1

[error] tool reported failure

Which token force-kept

shell
python3 -c "
import torch, reforacle as R
# which token force-kept in epsilon when all removed? distinct values
l = torch.tensor([[0.5, 0.1, 0.9, 0.2, 0.3]])
print('probs:', torch.softmax(l,-1).tolist())
print('epsilon 0.9 (all<0.9):', R.epsilon_filter(l, 0.9).tolist())
l2 = torch.tensor([[10.0, 0.0, 0.0, 0.0, 0.0]])
print('epsilon 0.9 peaked:', R.epsilon_filter(l2, 0.9).tolist())
print('min_p 0.99 peaked:', R.min_p_filter(l2, 0.99).tolist())
print('min_p 1.01 peaked:', R.min_p_filter(l2, 1.01).tolist())
# top_p all removed check with min_tokens - which kept
print('top_p 0.0 distinct:', R.top_p_filter(l, 0.0).tolist())
"
output
probs: [[0.21178840100765228, 0.14196601510047913, 0.3159511685371399, 0.156896710395813, 0.17339769005775452]]
epsilon 0.9 (all<0.9): [[-inf, -inf, 0.8999999761581421, -inf, -inf]]
epsilon 0.9 peaked: [[10.0, -inf, -inf, -inf, -inf]]
min_p 0.99 peaked: [[10.0, -inf, -inf, -inf, -inf]]
min_p 1.01 peaked: [[-inf, -inf, -inf, -inf, -inf]]
top_p 0.0 distinct: [[-inf, -inf, 0.8999999761581421, -inf, -inf]]

[stdout]
probs: [[0.21178840100765228, 0.14196601510047913, 0.3159511685371399, 0.156896710395813, 0.17339769005775452]]
epsilon 0.9 (all<0.9): [[-inf, -inf, 0.8999999761581421, -inf, -inf]]
epsilon 0.9 peaked: [[10.0, -inf, -inf, -inf, -inf]]
min_p 0.99 peaked: [[10.0, -inf, -inf, -inf, -inf]]
min_p 1.01 peaked: [[-inf, -inf, -inf, -inf, -inf]]
top_p 0.0 distinct: [[-inf, -inf, 0.8999999761581421, -inf, -inf]]

torch topk tie behavior

shell
python3 -c "
import torch
p = torch.full((1,5), 0.2)
print('topk(1) equal:', torch.topk(p,1).indices.tolist())
print('argmax equal:', torch.argmax(p,-1).tolist())
print('sort desc equal indices:', torch.sort(p, descending=True).indices.tolist())
"
output
topk(1) equal: [[2]]
argmax equal: [0]
sort desc equal indices: [[0, 1, 2, 3, 4]]

[stdout]
topk(1) equal: [[2]]
argmax equal: [0]
sort desc equal indices: [[0, 1, 2, 3, 4]]

Epsilon tie keep pattern by size

shell
python3 -c "
import torch, reforacle as R
for n in [3,4,5,6]:
    l = torch.zeros(1,n)
    r = R.epsilon_filter(l, 0.9)
    kept = [i for i,v in enumerate(r[0].tolist()) if v!=float('-inf')]
    print(f'zeros n={n} epsilon kept:', kept)
# min_p all-equal keep behavior
for n in [3,5]:
    l = torch.zeros(1,n)
    print(f'min_p 1.5 zeros n={n}:', [i for i,v in enumerate(R.min_p_filter(l,1.5)[0].tolist()) if v!=float('-inf')])
"
output
zeros n=3 epsilon kept: [2]
zeros n=4 epsilon kept: [3]
zeros n=5 epsilon kept: [4]
zeros n=6 epsilon kept: [5]
min_p 1.5 zeros n=3: []
min_p 1.5 zeros n=5: []

[stdout]
zeros n=3 epsilon kept: [2]
zeros n=4 epsilon kept: [3]
zeros n=5 epsilon kept: [4]
zeros n=6 epsilon kept: [5]
min_p 1.5 zeros n=3: []
min_p 1.5 zeros n=5: []

Tied maxima keep index

shell
python3 -c "
import torch, reforacle as R
# tied maxima
l = torch.tensor([[1.0, 0.0, 1.0]])
print('probs:', torch.softmax(l,-1).tolist())
print('epsilon 0.5 tied[1,0,1]:', R.epsilon_filter(l,0.5).tolist())
l2 = torch.tensor([[1.0, 2.0, 2.0, 0.0]])
print('epsilon big tied max at 1,2:', R.epsilon_filter(l2, 0.9).tolist())
# does epsilon guarantee keep at all? try top_k min tokens
# Also test top_p min_tokens tie
print('top_p 0.0 tied[1,0,1]:', R.top_p_filter(l,0.0).tolist())
"
output
probs: [[0.42231881618499756, 0.15536241233348846, 0.42231881618499756]]
epsilon 0.5 tied[1,0,1]: [[-inf, -inf, 1.0]]
epsilon big tied max at 1,2: [[-inf, -inf, 2.0, -inf]]
top_p 0.0 tied[1,0,1]: [[1.0, -inf, -inf]]

[stdout]
probs: [[0.42231881618499756, 0.15536241233348846, 0.42231881618499756]]
epsilon 0.5 tied[1,0,1]: [[-inf, -inf, 1.0]]
epsilon big tied max at 1,2: [[-inf, -inf, 2.0, -inf]]
top_p 0.0 tied[1,0,1]: [[1.0, -inf, -inf]]

Probe penalties

shell
python3 -c "
import torch, reforacle as R
g = torch.Generator().manual_seed(4)
logits = torch.randn(2,6, generator=g)
ids = torch.tensor([[0,0,1,4],[2,2,2,5]])
print('logits:', [[round(x,4) for x in r] for r in logits.tolist()])
print('rep pen 1.5:', [[round(x,4) for x in r] for r in R.repetition_penalty(logits, ids, 1.5).tolist()])
print('freq pen 0.5:', [[round(x,4) for x in r] for r in R.frequency_penalty(logits, ids, 0.5).tolist()])
print('pres pen 0.5:', [[round(x,4) for x in r] for r in R.presence_penalty(logits, ids, 0.5).tolist()])
"
output
logits: [[-1.6053, 0.2325, 2.2399, 0.8473, 1.2006, -0.4016], [-1.426, 0.9039, 0.8557, 0.6889, 0.885, 1.7706]]
rep pen 1.5: [[-3.6119, 0.155, 2.2399, 0.8473, 0.8004, -0.4016], [-1.426, 0.9039, 0.2535, 0.6889, 0.885, 1.1804]]
freq pen 0.5: [[-2.6053, -0.2675, 2.2399, 0.8473, 0.7006, -0.4016], [-1.426, 0.9039, -0.6443, 0.6889, 0.885, 1.2706]]
pres pen 0.5: [[-2.1053, -0.2675, 2.2399, 0.8473, 0.7006, -0.4016], [-1.426, 0.9039, 0.3557, 0.6889, 0.885, 1.2706]]

[stdout]
logits: [[-1.6053, 0.2325, 2.2399, 0.8473, 1.2006, -0.4016], [-1.426, 0.9039, 0.8557, 0.6889, 0.885, 1.7706]]
rep pen 1.5: [[-3.6119, 0.155, 2.2399, 0.8473, 0.8004, -0.4016], [-1.426, 0.9039, 0.2535, 0.6889, 0.885, 1.1804]]
freq pen 0.5: [[-2.6053, -0.2675, 2.2399, 0.8473, 0.7006, -0.4016], [-1.426, 0.9039, -0.6443, 0.6889, 0.885, 1.2706]]
pres pen 0.5: [[-2.1053, -0.2675, 2.2399, 0.8473, 0.7006, -0.4016], [-1.426, 0.9039, 0.3557, 0.6889, 0.885, 1.2706]]

Probe suppress and min_length

shell
python3 -c "
import torch, reforacle as R
g = torch.Generator().manual_seed(5)
logits = torch.randn(2,6, generator=g)
print('logits:', [[round(x,4) for x in r] for r in logits.tolist()])
print('suppress [1,3]:', R.suppress_tokens(logits, [1,3]).tolist())
print('suppress tensor:', R.suppress_tokens(logits, torch.tensor([0,5])).tolist())
# min_length_filter
print('minlen cur=2 min=5 eos=3:', R.min_length_filter(logits, 2, 5, 3).tolist())
print('minlen cur=5 min=5 eos=3:', R.min_length_filter(logits, 5, 5, 3).tolist())
print('minlen cur=6 min=5 eos=3:', R.min_length_filter(logits, 6, 5, 3).tolist())
"
output
logits: [[-0.4868, -0.6038, -0.5581, 0.6675, -0.1974, 1.9428], [-1.4017, -0.7626, 0.6312, -0.8991, -0.5578, 0.6907]]
suppress [1,3]: [[-0.4867807626724243, -inf, -0.5580962300300598, -inf, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -inf, 0.6312134265899658, -inf, -0.5577934384346008, 0.6907187700271606]]
suppress tensor: [[-inf, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, -inf], [-inf, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, -inf]]
minlen cur=2 min=5 eos=3: [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, -inf, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -inf, -0.5577934384346008, 0.6907187700271606]]
minlen cur=5 min=5 eos=3: [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
minlen cur=6 min=5 eos=3: [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]

[stdout]
logits: [[-0.4868, -0.6038, -0.5581, 0.6675, -0.1974, 1.9428], [-1.4017, -0.7626, 0.6312, -0.8991, -0.5578, 0.6907]]
suppress [1,3]: [[-0.4867807626724243, -inf, -0.5580962300300598, -inf, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -inf, 0.6312134265899658, -inf, -0.5577934384346008, 0.6907187700271606]]
suppress tensor: [[-inf, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, -inf], [-inf, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, -inf]]
minlen cur=2 min=5 eos=3: [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, -inf, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -inf, -0.5577934384346008, 0.6907187700271606]]
minlen cur=5 min=5 eos=3: [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
minlen cur=6 min=5 eos=3: [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]

Probe mask/position functions

shell
python3 -c "
import torch, reforacle as R
ids = torch.tensor([[5,6,7,0,0],[1,2,0,0,0]])
print('attention_mask pad0:', R.attention_mask(ids, 0).tolist(), R.attention_mask(ids,0).dtype)
print('position_ids pad0:', R.position_ids(ids, 0).tolist(), R.position_ids(ids,0).dtype)
print('sequence_lengths pad0:', R.sequence_lengths(ids, 0).tolist(), R.sequence_lengths(ids,0).dtype)
print('causal_mask(4):', R.causal_mask(4).tolist(), R.causal_mask(4).dtype)
"
output
attention_mask pad0: [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] torch.float32
position_ids pad0: [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1]] torch.int64
sequence_lengths pad0: [3, 2] torch.int64
causal_mask(4): [[0.0, -inf, -inf, -inf], [0.0, 0.0, -inf, -inf], [0.0, 0.0, 0.0, -inf], [0.0, 0.0, 0.0, 0.0]] torch.float32

[stdout]
attention_mask pad0: [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] torch.float32
position_ids pad0: [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1]] torch.int64
sequence_lengths pad0: [3, 2] torch.int64
causal_mask(4): [[0.0, -inf, -inf, -inf], [0.0, 0.0, -inf, -inf], [0.0, 0.0, 0.0, -inf], [0.0, 0.0, 0.0, 0.0]] torch.float32

Position edge cases and 1d

shell
python3 -c "
import torch, reforacle as R
# leading pad / interior pad
ids = torch.tensor([[0,5,6,0,7]])
print('pos leading/interior pad:', R.position_ids(ids,0).tolist())
print('cumsum-1:', ((ids!=0).long().cumsum(-1)-1).tolist())
print('attn:', R.attention_mask(ids,0).tolist())
# 1d input
ids1 = torch.tensor([5,6,0,0])
print('attn 1d:', R.attention_mask(ids1,0).tolist())
print('pos 1d:', R.position_ids(ids1,0).tolist())
print('seqlen 1d:', R.sequence_lengths(ids1,0).tolist())
"
output
pos leading/interior pad: [[0, 0, 1, 1, 2]]
cumsum-1: [[-1, 0, 1, 1, 2]]
attn: [[0.0, 1.0, 1.0, 0.0, 1.0]]
attn 1d: [1.0, 1.0, 0.0, 0.0]
pos 1d: [0, 1, 1, 1]
seqlen 1d: 2

[stdout]
pos leading/interior pad: [[0, 0, 1, 1, 2]]
cumsum-1: [[-1, 0, 1, 1, 2]]
attn: [[0.0, 1.0, 1.0, 0.0, 1.0]]
attn 1d: [1.0, 1.0, 0.0, 0.0]
pos 1d: [0, 1, 1, 1]
seqlen 1d: 2

Probe label functions

shell
python3 -c "
import torch, reforacle as R
ids = torch.tensor([[5,6,7,8,9],[1,2,3,4,0]])
print('shift_labels ign=-100:', R.shift_labels(ids, -100).tolist())
print('mask_after_eos eos=7:', R.mask_after_eos(ids, 7).tolist())
print('mask_after_eos eos=3:', R.mask_after_eos(ids, 3).tolist())
pl = torch.tensor([2,3])
print('completion_loss_labels pl=[2,3] ign=-100:', R.completion_loss_labels(ids, pl, -100).tolist())
print('prepare_labels pl=[2,3] ign=-100:', R.prepare_labels(ids, pl, -100).tolist())
"
output
shift_labels ign=-100: [[6, 7, 8, 9, -100], [2, 3, 4, 0, -100]]
mask_after_eos eos=7: [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
mask_after_eos eos=3: [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 0.0]]
completion_loss_labels pl=[2,3] ign=-100: [[-100, -100, 7, 8, 9], [-100, -100, -100, 4, 0]]
prepare_labels pl=[2,3] ign=-100: [[-100, 7, 8, 9, -100], [-100, -100, 4, 0, -100]]

[stdout]
shift_labels ign=-100: [[6, 7, 8, 9, -100], [2, 3, 4, 0, -100]]
mask_after_eos eos=7: [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
mask_after_eos eos=3: [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 0.0, 0.0]]
completion_loss_labels pl=[2,3] ign=-100: [[-100, -100, 7, 8, 9], [-100, -100, -100, 4, 0]]
prepare_labels pl=[2,3] ign=-100: [[-100, 7, 8, 9, -100], [-100, -100, 4, 0, -100]]

Probe last_token, gather, truncate

shell
python3 -c "
import torch, reforacle as R
am = torch.tensor([[1.,1.,1.,0.,0.],[1.,1.,0.,0.,0.]])
print('last_token_indices:', R.last_token_indices(am).tolist(), R.last_token_indices(am).dtype)
g=torch.Generator().manual_seed(6)
logits = torch.randn(2,5,4, generator=g)
gl = R.gather_last_logits(logits, am)
print('gather_last_logits shape:', gl.shape, gl.tolist())
print('  check row0 idx2:', logits[0,2].tolist())
print('  check row1 idx1:', logits[1,1].tolist())
ids = torch.tensor([[1,2,3,4,5],[6,7,8,9,10]])
print('truncate_left max3:', R.truncate_left(ids,3).tolist())
print('truncate_left max10:', R.truncate_left(ids,10).tolist())
"
output
last_token_indices: [2, 1] torch.int64
gather_last_logits shape: torch.Size([2, 4]) [[-0.48055049777030945, 1.1031601428985596, 2.548506021499634, 0.30063536763191223], [-2.344620704650879, -0.240716814994812, -2.641396999359131, -1.8220865726470947]]
  check row0 idx2: [-0.48055049777030945, 1.1031601428985596, 2.548506021499634, 0.30063536763191223]
  check row1 idx1: [-2.344620704650879, -0.240716814994812, -2.641396999359131, -1.8220865726470947]
truncate_left max3: [[3, 4, 5], [8, 9, 10]]
truncate_left max10: [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]

[stdout]
last_token_indices: [2, 1] torch.int64
gather_last_logits shape: torch.Size([2, 4]) [[-0.48055049777030945, 1.1031601428985596, 2.548506021499634, 0.30063536763191223], [-2.344620704650879, -0.240716814994812, -2.641396999359131, -1.8220865726470947]]
  check row0 idx2: [-0.48055049777030945, 1.1031601428985596, 2.548506021499634, 0.30063536763191223]
  check row1 idx1: [-2.344620704650879, -0.240716814994812, -2.641396999359131, -1.8220865726470947]
truncate_left max3: [[3, 4, 5], [8, 9, 10]]
truncate_left max10: [[1, 2, 3, 4, 5], [6, 7, 8, 9, 10]]

Probe padding functions

shell
python3 -c "
import torch, reforacle as R
ids = torch.tensor([[1,2,3],[4,5,6]])
print('pad_to_multiple pad0 mult4:', R.pad_to_multiple(ids,0,4).tolist())
print('pad_to_multiple pad0 mult3:', R.pad_to_multiple(ids,0,3).tolist())
print('pad_to_multiple pad0 mult2:', R.pad_to_multiple(ids,0,2).tolist())
print('right_pad_to pad0 w5:', R.right_pad_to(ids,0,5).tolist())
print('right_pad_to pad0 w2:', R.right_pad_to(ids,0,2).tolist())
"
output
pad_to_multiple pad0 mult4: [[1, 2, 3, 0], [4, 5, 6, 0]]
pad_to_multiple pad0 mult3: [[1, 2, 3], [4, 5, 6]]
pad_to_multiple pad0 mult2: [[1, 2, 3, 0], [4, 5, 6, 0]]
right_pad_to pad0 w5: [[1, 2, 3, 0, 0], [4, 5, 6, 0, 0]]
right_pad_to pad0 w2: [[1, 2], [4, 5]]

[stdout]
pad_to_multiple pad0 mult4: [[1, 2, 3, 0], [4, 5, 6, 0]]
pad_to_multiple pad0 mult3: [[1, 2, 3], [4, 5, 6]]
pad_to_multiple pad0 mult2: [[1, 2, 3, 0], [4, 5, 6, 0]]
right_pad_to pad0 w5: [[1, 2, 3, 0, 0], [4, 5, 6, 0, 0]]
right_pad_to pad0 w2: [[1, 2], [4, 5]]

Probe action_logprobs, entropy

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(7)
logits = torch.randn(2,3,4, generator=g)
actions = torch.tensor([[0,1,2],[3,2,1]])
al = R.action_logprobs(logits, actions)
print('action_logprobs shape:', al.shape, al.tolist())
lp = torch.log_softmax(logits,-1)
print('  manual:', [[lp[0,i,actions[0,i]].item() for i in range(3)],[lp[1,i,actions[1,i]].item() for i in range(3)]])
ent = R.entropy_from_logits(logits)
print('entropy shape:', ent.shape, ent.tolist())
p = torch.softmax(logits,-1)
print('  manual ent:', (-(p*lp).sum(-1)).tolist())
"
output
action_logprobs shape: torch.Size([2, 3]) [[-2.3531417846679688, -0.850128173828125, -1.7944910526275635], [-0.6806192994117737, -2.6537225246429443, -2.062441349029541]]
  manual: [[-2.3531417846679688, -0.850128173828125, -1.7944910526275635], [-0.6806192994117737, -2.6537225246429443, -2.062441349029541]]
entropy shape: torch.Size([2, 3]) [[1.0820982456207275, 1.2848243713378906, 0.8998391628265381], [1.114432454109192, 1.2564908266067505, 1.3314017057418823]]
  manual ent: [[1.0820982456207275, 1.2848243713378906, 0.8998391628265381], [1.114432454109192, 1.256490707397461, 1.3314017057418823]]

[stdout]
action_logprobs shape: torch.Size([2, 3]) [[-2.3531417846679688, -0.850128173828125, -1.7944910526275635], [-0.6806192994117737, -2.6537225246429443, -2.062441349029541]]
  manual: [[-2.3531417846679688, -0.850128173828125, -1.7944910526275635], [-0.6806192994117737, -2.6537225246429443, -2.062441349029541]]
entropy shape: torch.Size([2, 3]) [[1.0820982456207275, 1.2848243713378906, 0.8998391628265381], [1.114432454109192, 1.2564908266067505, 1.3314017057418823]]
  manual ent: [[1.0820982456207275, 1.2848243713378906, 0.8998391628265381], [1.114432454109192, 1.256490707397461, 1.3314017057418823]]

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('reward_to_go:', R.reward_to_go(rewards, mask).tolist())
x = torch.tensor([[1.,2.,3.,4.]])
print('discount_cumsum g=0.9:', R.discount_cumsum(x, 0.9).tolist())
print('discount_cumsum g=1.0:', R.discount_cumsum(x, 1.0).tolist())
sl = torch.tensor([2,3])
print('build_position_ids_packed total5:', R.build_position_ids_packed(sl,5).tolist())
sl2 = torch.tensor([1,2,2])
print('build_position_ids_packed [1,2,2] total5:', R.build_position_ids_packed(sl2,5).tolist())
"
output
reward_to_go: [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]]
discount_cumsum g=0.9: [[8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]]
discount_cumsum g=1.0: [[10.0, 9.0, 7.0, 4.0]]
build_position_ids_packed total5: [0, 1, 0, 1, 2]
build_position_ids_packed [1,2,2] total5: [0, 0, 1, 0, 1]

[stdout]
reward_to_go: [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]]
discount_cumsum g=0.9: [[8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]]
discount_cumsum g=1.0: [[10.0, 9.0, 7.0, 4.0]]
build_position_ids_packed total5: [0, 1, 0, 1, 2]
build_position_ids_packed [1,2,2] total5: [0, 0, 1, 0, 1]

reward_to_go mask, packed edge, first_stop

shell
python3 -c "
import torch, reforacle as R
# reward_to_go interior mask
rewards = torch.tensor([[1.,2.,3.]])
mask = torch.tensor([[1.,0.,1.]])
print('rtg interior mask [1,0,1]:', R.reward_to_go(rewards,mask).tolist())
# build packed total>sum
print('bpp [2,2] total5:', R.build_position_ids_packed(torch.tensor([2,2]),5).tolist())
# first_stop_index
ids = torch.tensor([[1,2,3,4,5],[6,3,2,3,1]])
print('first_stop_index stop=[3,5]:', R.first_stop_index(ids,[3,5]).tolist())
print('first_stop_index stop=[9]:', R.first_stop_index(ids,[9]).tolist())
print('first_stop_index stop=[3]:', R.first_stop_index(ids,[3]).tolist())
"
output
Exit code 1
rtg interior mask [1,0,1]: [[4.0, 0.0, 3.0]]
bpp [2,2] total5: [0, 1, 0, 1, 0]
Traceback (most recent call last):
  File "<string>", line 11, 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 behavior

shell
python3 -c "
import torch, reforacle as R
ids = torch.tensor([[1,2,3,4,5],[6,3,2,3,1]])
print('first_stop stop=[3,5]:', R.first_stop_index(ids,[3,5]), type(R.first_stop_index(ids,[3,5])))
print('first_stop stop=[9]:', R.first_stop_index(ids,[9]))
print('first_stop stop=[3]:', R.first_stop_index(ids,[3]))
print('first_stop stop=[1]:', R.first_stop_index(ids,[1]))
"
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 with tensor stops

shell
python3 -c "
import torch, reforacle as R
ids = torch.tensor([[1,2,3,4,5],[6,3,2,3,1]])
print('stop=[3,5]:', R.first_stop_index(ids,torch.tensor([3,5])))
print('stop=[9]:', R.first_stop_index(ids,torch.tensor([9])))
print('stop=[3]:', R.first_stop_index(ids,torch.tensor([3])))
print('stop=[1]:', R.first_stop_index(ids,torch.tensor([1])))
r = R.first_stop_index(ids,torch.tensor([3,5]))
print('type/dtype:', type(r), r.dtype if hasattr(r,'dtype') else 'na')
"
output
stop=[3,5]: tensor([2, 1])
stop=[9]: tensor([-1, -1])
stop=[3]: tensor([2, 1])
stop=[1]: tensor([0, 4])
type/dtype: <class 'torch.Tensor'> torch.int64

[stdout]
stop=[3,5]: tensor([2, 1])
stop=[9]: tensor([-1, -1])
stop=[3]: tensor([2, 1])
stop=[1]: tensor([0, 4])
type/dtype: <class 'torch.Tensor'> torch.int64

masked_argmax, topk_indices

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(8)
logits = torch.randn(2,6, generator=g)
mask = torch.tensor([[1.,1.,0.,1.,0.,0.],[0.,1.,1.,0.,1.,1.]])
print('logits:', [[round(x,3) for x in r] for r in logits.tolist()])
print('masked_argmax:', R.masked_argmax(logits, mask).tolist(), R.masked_argmax(logits,mask).dtype)
print('topk_indices k=3:', R.topk_indices(logits, 3).tolist(), R.topk_indices(logits,3).dtype)
print('topk_indices k=1:', R.topk_indices(logits, 1).tolist())
"
output
logits: [[0.271, -1.273, 0.503, 0.418, -0.639, -0.661], [-0.143, -0.104, -1.531, 0.632, -1.345, 1.431]]
masked_argmax: [3, 5] torch.int64
topk_indices k=3: [[2, 3, 0], [5, 3, 1]] torch.int64
topk_indices k=1: [[2], [5]]

[stdout]
logits: [[0.271, -1.273, 0.503, 0.418, -0.639, -0.661], [-0.143, -0.104, -1.531, 0.632, -1.345, 1.431]]
masked_argmax: [3, 5] torch.int64
topk_indices k=3: [[2, 3, 0], [5, 3, 1]] torch.int64
topk_indices k=1: [[2], [5]]

generation_stop_mask

shell
python3 -c "
import torch, reforacle as R
ids = torch.tensor([[5,6,3,8],[3,2,1,9],[5,6,7,3]])
print('gen_stop_mask eos=3 min=0:', R.generation_stop_mask(ids,3,0).tolist(), R.generation_stop_mask(ids,3,0).dtype)
print('gen_stop_mask eos=3 min=2:', R.generation_stop_mask(ids,3,2).tolist())
print('gen_stop_mask eos=3 min=3:', R.generation_stop_mask(ids,3,3).tolist())
"
output
gen_stop_mask eos=3 min=0: [[1.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]] torch.float32
gen_stop_mask eos=3 min=2: [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]]
gen_stop_mask eos=3 min=3: [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0]]

[stdout]
gen_stop_mask eos=3 min=0: [[1.0, 1.0, 1.0, 0.0], [1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]] torch.float32
gen_stop_mask eos=3 min=2: [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0]]
gen_stop_mask eos=3 min=3: [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0]]

Disambiguate generation_stop_mask

shell
python3 -c "
import torch, reforacle as R
ids = torch.tensor([[3,5,3,7],[3,5,6,7],[5,3,6,3]])
for m in [0,2,3]:
    print(f'min={m}:', R.generation_stop_mask(ids,3,m).tolist())
"
output
min=0: [[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]
min=2: [[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]
min=3: [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 0.0]]

[stdout]
min=0: [[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]
min=2: [[1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0]]
min=3: [[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 0.0]]

logits_warper composite

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(9)
logits = torch.randn(2,6, generator=g)
print('logits:', [[round(x,3) for x in r] for r in logits.tolist()])
w = R.logits_warper(logits, 2.0, 3, 0.9)
print('warper T2 k3 p0.9:', w.tolist())
sg = R.sample_greedy_warped(logits, 2.0, 3, 0.9)
print('sample_greedy_warped:', sg.tolist(), sg.dtype)
"
output
logits: [[0.045, 1.911, -0.231, 0.346, 1.318, 0.37], [0.384, 0.297, 0.747, 1.087, -0.772, -0.294]]
warper T2 k3 p0.9: [[-inf, 0.9556195139884949, -inf, -inf, 0.659014105796814, 0.18481850624084473], [0.19205577671527863, -inf, 0.37364327907562256, 0.5436873435974121, -inf, -inf]]
sample_greedy_warped: [1, 3] torch.int64

[stdout]
logits: [[0.045, 1.911, -0.231, 0.346, 1.318, 0.37], [0.384, 0.297, 0.747, 1.087, -0.772, -0.294]]
warper T2 k3 p0.9: [[-inf, 0.9556195139884949, -inf, -inf, 0.659014105796814, 0.18481850624084473], [0.19205577671527863, -inf, 0.37364327907562256, 0.5436873435974121, -inf, -inf]]
sample_greedy_warped: [1, 3] torch.int64

rollout_last_logits

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(10)
logits = torch.randn(2,5,4, generator=g)  # batch, seq, vocab
ids = torch.tensor([[5,6,7,0,0],[1,2,3,4,0]])
r = R.rollout_last_logits(logits, ids, 0)
print('rollout_last_logits shape:', r.shape)
print(r.tolist())
print('logits[0,2]:', logits[0,2].tolist())
print('logits[1,3]:', logits[1,3].tolist())
"
output
rollout_last_logits shape: torch.Size([2, 4])
[[-0.7464958429336548, 1.0050938129425049, -0.25683921575546265, 0.4764931797981262], [1.3358182907104492, 0.3075266480445862, 0.6277765035629272, -0.10053547471761703]]
logits[0,2]: [-0.7464958429336548, 1.0050938129425049, -0.25683921575546265, 0.4764931797981262]
logits[1,3]: [1.3358182907104492, 0.3075266480445862, 0.6277765035629272, -0.10053547471761703]

[stdout]
rollout_last_logits shape: torch.Size([2, 4])
[[-0.7464958429336548, 1.0050938129425049, -0.25683921575546265, 0.4764931797981262], [1.3358182907104492, 0.3075266480445862, 0.6277765035629272, -0.10053547471761703]]
logits[0,2]: [-0.7464958429336548, 1.0050938129425049, -0.25683921575546265, 0.4764931797981262]
logits[1,3]: [1.3358182907104492, 0.3075266480445862, 0.6277765035629272, -0.10053547471761703]

Edge cases 1d, dtypes

shell
python3 -c "
import torch, reforacle as R
# 1d penalties?
try:
    l = torch.randn(6); ids=torch.tensor([0,0,1])
    print('rep 1d:', R.repetition_penalty(l, ids, 1.5).shape)
except Exception as e: print('rep 1d err:', type(e).__name__, e)
# temperature T=0 tie
l = torch.tensor([[1.0,1.0,0.0]])
print('temp T0 tie:', R.temperature_scale(l,0.0).tolist())
# min_length when cur<min but eos already -inf? just check dtype preserve
# causal_mask type
print('causal_mask(2):', R.causal_mask(2).tolist())
# suppress with empty
print('suppress empty:', R.suppress_tokens(torch.randn(1,4), []).shape)
# shift_labels dtype
print('shift dtype:', R.shift_labels(torch.tensor([[1,2,3]]), -100).dtype)
# completion dtype
print('comp dtype:', R.completion_loss_labels(torch.tensor([[1,2,3]]), torch.tensor([1]), -100).dtype)
"
output
rep 1d err: ValueError need more than 1 value to unpack
temp T0 tie: [[-inf, 0.0, -inf]]
causal_mask(2): [[0.0, -inf], [0.0, 0.0]]
suppress empty: torch.Size([1, 4])
shift dtype: torch.int64
comp dtype: torch.int64

[stdout]
rep 1d err: ValueError need more than 1 value to unpack
temp T0 tie: [[-inf, 0.0, -inf]]
causal_mask(2): [[0.0, -inf], [0.0, 0.0]]
suppress empty: torch.Size([1, 4])
shift dtype: torch.int64
comp dtype: torch.int64

Compare greedy vs T0 tie handling

shell
python3 -c "
import torch, reforacle as R
l = torch.tensor([[1.0,1.0,0.0],[0.0,2.0,2.0]])
print('greedy tie:', R.greedy_token(l).tolist())
print('temp T0 tie:', R.temperature_scale(l,0.0).tolist())
print('argmax:', l.argmax(-1).tolist())
print('argsort asc last:', l.argsort(-1)[...,-1].tolist())
"
output
greedy tie: [1, 2]
temp T0 tie: [[-inf, 0.0, -inf], [-inf, -inf, 0.0]]
argmax: [0, 1]
argsort asc last: [1, 2]

[stdout]
greedy tie: [1, 2]
temp T0 tie: [[-inf, 0.0, -inf], [-inf, -inf, 0.0]]
argmax: [0, 1]
argsort asc last: [1, 2]

Tie behavior for argmax functions

shell
python3 -c "
import torch, reforacle as R
l = torch.tensor([[2.0,1.0,2.0,0.0]])
print('greedy:', R.greedy_token(l).tolist())  # expect 2 (last)
mask = torch.tensor([[1.,1.,1.,1.]])
print('masked_argmax full:', R.masked_argmax(l,mask).tolist())
mask2 = torch.tensor([[1.,1.,0.,1.]])  # exclude idx2, tie now idx0 only
print('masked_argmax excl idx2:', R.masked_argmax(l,mask2).tolist())
print('topk_indices k=3 tie:', R.topk_indices(l,3).tolist())
# sample_greedy_warped tie: T=1,large k,p=1 -> greedy of logits
print('sgw tie:', R.sample_greedy_warped(l,1.0,10,1.0).tolist())
# first_stop already fine
"
output
greedy: [2]
masked_argmax full: [2]
masked_argmax excl idx2: [0]
topk_indices k=3 tie: [[0, 2, 1]]
sgw tie: [2]

[stdout]
greedy: [2]
masked_argmax full: [2]
masked_argmax excl idx2: [0]
topk_indices k=3 tie: [[0, 2, 1]]
sgw tie: [2]
write · /app/objectives.py
from __future__ import annotations

import torch
import torch.nn.functional as F

NEG_INF = float("-inf")


def _last_argmax(x):
    # argmax that returns the LAST index on ties (matches reforacle convention)
    flipped = torch.flip(x, [-1])
    idx = torch.argmax(flipped, dim=-1)
    return (x.shape[-1] - 1) - idx


def greedy_token(logits):
    return _last_argmax(logits)


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


def top_k_filter(logits, k):
    k = min(int(k), logits.shape[-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_idx = torch.sort(logits, descending=True, dim=-1)
    cum = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
    sorted_remove = cum > p
    # keep at least the top token: shift removal mask right by one
    sorted_remove[..., 1:] = sorted_remove[..., :-1].clone()
    sorted_remove[..., 0] = False
    remove = torch.zeros_like(sorted_remove).scatter(-1, sorted_idx, 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))
    return logits - penalty * (counts > 0).to(logits.dtype)


def suppress_tokens(logits, token_ids):
    idx = torch.as_tensor(token_ids, dtype=torch.long, device=logits.device)
    return logits.index_fill(-1, idx, NEG_INF)


def min_length_filter(logits, cur_len, min_len, eos_id):
    if cur_len < min_len:
        idx = torch.tensor([eos_id], dtype=torch.long, device=logits.device)
        return logits.index_fill(-1, idx, NEG_INF)
    return logits


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()
    return (mask.cumsum(dim=-1) - 1).clamp(min=0)


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


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


def completion_loss_labels(input_ids, prompt_lengths, ignore_index):
    labels = input_ids.clone()
    s = input_ids.shape[-1]
    ar = torch.arange(s, 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).long()
    before = is_eos.cumsum(dim=-1) - is_eos
    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) - 1).long()
    b = logits.shape[0]
    return logits[torch.arange(b, device=logits.device), idx]


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


def pad_to_multiple(input_ids, pad_id, multiple):
    s = input_ids.shape[-1]
    new = ((s + multiple - 1) // multiple) * multiple
    if new > s:
        return F.pad(input_ids, (0, new - s), value=pad_id)
    return input_ids


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


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


def reward_to_go(rewards, mask):
    masked = rewards * mask
    rev = torch.flip(torch.cumsum(torch.flip(masked, [-1]), dim=-1), [-1])
    return rev * mask


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


def build_position_ids_packed(seq_lengths, total_len):
    out = torch.zeros(total_len, dtype=torch.long)
    pos = 0
    for L in seq_lengths.tolist():
        out[pos:pos + L] = torch.arange(L)
        pos += L
    return out


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


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)
    s = input_ids.shape[-1]
    ar = torch.arange(s, device=input_ids.device)
    mask = ar.unsqueeze(0) < (prompt_lengths - 1).unsqueeze(1)
    labels[mask] = ignore_index
    return labels


def rollout_last_logits(logits, input_ids, pad_id):
    mask = (input_ids != pad_id).float()
    return gather_last_logits(logits, mask)


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


def first_stop_index(input_ids, stop_ids):
    isin = torch.isin(input_ids, stop_ids)
    any_stop = isin.any(dim=-1)
    idx = isin.long().argmax(dim=-1)
    return torch.where(any_stop, idx, torch.full_like(idx, -1))


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


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


def right_pad_to(input_ids, pad_id, width):
    s = input_ids.shape[-1]
    if width <= s:
        return input_ids[..., :width]
    return F.pad(input_ids, (0, width - s), value=pad_id)


def generation_stop_mask(input_ids, eos_id, min_len):
    base = mask_after_eos(input_ids, eos_id)
    s = input_ids.shape[-1]
    ar = torch.arange(s, device=input_ids.device)
    minmask = (ar < min_len).to(base.dtype)
    return torch.maximum(base, minmask)

Comprehensive comparison test

shell
python3 -c "
import torch, importlib
import reforacle as R
import objectives as O
importlib.reload(O)

g = torch.Generator().manual_seed(123)
def rl(*s): return torch.randn(*s, generator=g)
def ri(hi,*s): return torch.randint(0,hi,s,generator=g)

fails=[]
def chk(name, a, b, atol=1e-6, rtol=1e-5):
    try:
        a=a if torch.is_tensor(a) else torch.tensor(a)
        b=b if torch.is_tensor(b) else torch.tensor(b)
        if a.shape!=b.shape: fails.append((name,'shape',a.shape,b.shape)); return
        if a.dtype!=b.dtype: fails.append((name,'dtype',a.dtype,b.dtype)); return
        af=a.float(); bf=b.float()
        # treat -inf equal
        m=torch.isinf(af)&torch.isinf(bf)&(torch.sign(af)==torch.sign(bf))
        af=af.masked_fill(m,0); bf=bf.masked_fill(m,0)
        if not torch.allclose(af,bf,atol=atol,rtol=rtol,equal_nan=True):
            d=(af-bf).abs().max().item(); fails.append((name,'val',d))
    except Exception as e:
        fails.append((name,'exc',repr(e)))

for trial in range(50):
    L=rl(3,8); ids=ri(8,3,5); pl=torch.randint(1,5,(3,),generator=g)
    chk('greedy',O.greedy_token(L),R.greedy_token(L))
    for T in [0.0,0.5,1.0,2.0,0.7]:
        chk(f'temp{T}',O.temperature_scale(L,T),R.temperature_scale(L,T))
    for k in [1,2,3,5,8,10]:
        chk(f'topk{k}',O.top_k_filter(L,k),R.top_k_filter(L,k))
    for p in [0.0,0.3,0.5,0.9,1.0]:
        chk(f'topp{p}',O.top_p_filter(L,p),R.top_p_filter(L,p))
        chk(f'minp{p}',O.min_p_filter(L,p),R.min_p_filter(L,p))
    for e in [0.0,0.05,0.1,0.2,0.5]:
        chk(f'eps{e}',O.epsilon_filter(L,e),R.epsilon_filter(L,e))
    for pen in [0.5,1.0,1.2,1.5,2.0]:
        chk(f'rep{pen}',O.repetition_penalty(L,ids,pen),R.repetition_penalty(L,ids,pen))
        chk(f'freq{pen}',O.frequency_penalty(L,ids,pen),R.frequency_penalty(L,ids,pen))
        chk(f'pres{pen}',O.presence_penalty(L,ids,pen),R.presence_penalty(L,ids,pen))
    chk('suppress',O.suppress_tokens(L,[1,3,5]),R.suppress_tokens(L,[1,3,5]))
    for cl in [1,3,5]:
        chk(f'minlen{cl}',O.min_length_filter(L,cl,4,2),R.min_length_filter(L,cl,4,2))
    chk('renorm',O.renormalize_logits(L),R.renormalize_logits(L))
    pids=ri(6,3,7)  # pad id 0
    chk('attn',O.attention_mask(pids,0),R.attention_mask(pids,0))
    chk('pos',O.position_ids(pids,0),R.position_ids(pids,0))
    chk('seqlen',O.sequence_lengths(pids,0),R.sequence_lengths(pids,0))
    chk('causal',O.causal_mask(5),R.causal_mask(5))
    chk('shift',O.shift_labels(ids,-100),R.shift_labels(ids,-100))
    chk('comp',O.completion_loss_labels(ids,pl,-100),R.completion_loss_labels(ids,pl,-100))
    chk('prep',O.prepare_labels(ids,pl,-100),R.prepare_labels(ids,pl,-100))
    chk('maskeos',O.mask_after_eos(pids,2),R.mask_after_eos(pids,2))
    am=O.attention_mask(pids,0)
    chk('lasttok',O.last_token_indices(am),R.last_token_indices(am))
    L3=rl(3,7,8)
    chk('gatherlast',O.gather_last_logits(L3,am),R.gather_last_logits(L3,am))
    chk('rollout',O.rollout_last_logits(L3,pids,0),R.rollout_last_logits(L3,pids,0))
    for ml in [1,3,5,10]:
        chk(f'trunc{ml}',O.truncate_left(ids,ml),R.truncate_left(ids,ml))
    for m in [2,3,4,5]:
        chk(f'padmult{m}',O.pad_to_multiple(ids,0,m),R.pad_to_multiple(ids,0,m))
    for w in [2,5,7]:
        chk(f'rpad{w}',O.right_pad_to(ids,0,w),R.right_pad_to(ids,0,w))
    acts=ri(8,3,7)
    chk('actlp',O.action_logprobs(L3,acts),R.action_logprobs(L3,acts))
    chk('entropy',O.entropy_from_logits(L3),R.entropy_from_logits(L3))
    rw=rl(3,6); mk=(rl(3,6)>0).float()
    chk('rtg',O.reward_to_go(rw,mk),R.reward_to_go(rw,mk))
    for gm in [0.0,0.9,0.99,1.0]:
        chk(f'disc{gm}',O.discount_cumsum(rw,gm),R.discount_cumsum(rw,gm))
    sl=torch.randint(1,4,(3,),generator=g); tot=int(sl.sum())
    chk('bpp',O.build_position_ids_packed(sl,tot),R.build_position_ids_packed(sl,tot))
    chk('bpp+',O.build_position_ids_packed(sl,tot+2),R.build_position_ids_packed(sl,tot+2))
    chk('warper',O.logits_warper(L,1.0,4,0.9),R.logits_warper(L,1.0,4,0.9))
    chk('sgw',O.sample_greedy_warped(L,1.0,4,0.9),R.sample_greedy_warped(L,1.0,4,0.9))
    stop=torch.tensor([2,5])
    chk('firststop',O.first_stop_index(pids,stop),R.first_stop_index(pids,stop))
    chk('maskedargmax',O.masked_argmax(L,mk[:, :8] if mk.shape[1]>=8 else (rl(3,8)>0).float()),R.masked_argmax(L,(mk[:, :8] if mk.shape[1]>=8 else (rl(3,8)>0).float())))
    for k in [1,3,5]:
        chk(f'topkidx{k}',O.topk_indices(L,k),R.topk_indices(L,k))
    for ml in [0,2,4]:
        chk(f'genstop{ml}',O.generation_stop_mask(pids,2,ml),R.generation_stop_mask(pids,2,ml))

from collections import Counter
c=Counter(f[0] for f in fails)
print('total fails:', len(fails))
for name,cnt in c.most_common(): print(' ',name,cnt)
if fails: print('sample:', fails[:5])
"
output
total fails: 186
  maskedargmax 48
  lasttok 46
  gatherlast 46
  rollout 46
sample: [('lasttok', 'val', 3.0), ('gatherlast', 'val', 2.9269492626190186), ('rollout', 'val', 2.9269492626190186), ('maskedargmax', 'val', 2.0), ('lasttok', 'val', 3.0)]

[stdout]
total fails: 186
  maskedargmax 48
  lasttok 46
  gatherlast 46
  rollout 46
sample: [('lasttok', 'val', 3.0), ('gatherlast', 'val', 2.9269492626190186), ('rollout', 'val', 2.9269492626190186), ('maskedargmax', 'val', 2.0), ('lasttok', 'val', 3.0)]

Debug last_token_indices

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(123)
pids=torch.randint(0,6,(4,7),generator=g)
am=(pids!=0).float()
print('am:', am.tolist())
print('R lasttok:', R.last_token_indices(am).tolist())
print('O lasttok:', O.last_token_indices(am).tolist())
"
output
am: [[1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0], [1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0]]
R lasttok: [5, 6, 6, 5]
O lasttok: [3, 5, 5, 4]

[stdout]
am: [[1.0, 1.0, 0.0, 0.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 1.0], [1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 0.0, 1.0, 0.0]]
R lasttok: [5, 6, 6, 5]
O lasttok: [3, 5, 5, 4]

Debug masked_argmax and last_token fix

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(5)
L=torch.randn(4,8,generator=g)
mask=(torch.randn(4,8,generator=g)>0).float()
print('mask:',mask.tolist())
print('R masked_argmax:', R.masked_argmax(L,mask).tolist())
print('O masked_argmax:', O.masked_argmax(L,mask).tolist())
# check last_token via _last_argmax
print('R lasttok:', R.last_token_indices(mask).tolist())
print('O _last_argmax(mask):', O._last_argmax(mask).tolist())
"
output
mask: [[0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0], [0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0], [0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0]]
R masked_argmax: [4, 1, 0, 7]
O masked_argmax: [4, 1, 0, 7]
R lasttok: [6, 2, 5, 7]
O _last_argmax(mask): [6, 2, 5, 7]

[stdout]
mask: [[0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 1.0, 0.0], [0.0, 1.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0], [0.0, 1.0, 1.0, 0.0, 1.0, 1.0, 1.0, 1.0]]
R masked_argmax: [4, 1, 0, 7]
O masked_argmax: [4, 1, 0, 7]
R lasttok: [6, 2, 5, 7]
O _last_argmax(mask): [6, 2, 5, 7]

/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) - 1).long()
    b = logits.shape[0]
    return logits[torch.arange(b, device=logits.device), idx]
+ added
def last_token_indices(attn_mask):
    return _last_argmax(attn_mask)


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

Full comparison, corrected

shell
python3 -c "
import torch, importlib
import reforacle as R
import objectives as O
importlib.reload(O)
g = torch.Generator().manual_seed(999)
def rl(*s): return torch.randn(*s, generator=g)
def ri(hi,*s): return torch.randint(0,hi,s,generator=g)
fails=[]
def chk(name, fo, fr):
    try:
        a=fo(); b=fr()
        a=a if torch.is_tensor(a) else torch.tensor(a)
        b=b if torch.is_tensor(b) else torch.tensor(b)
        if a.shape!=b.shape: fails.append((name,'shape',tuple(a.shape),tuple(b.shape))); return
        if a.dtype!=b.dtype: fails.append((name,'dtype',a.dtype,b.dtype)); return
        af=a.float(); bf=b.float()
        m=torch.isinf(af)&torch.isinf(bf)&(torch.sign(af)==torch.sign(bf))
        af=af.masked_fill(m,0); bf=bf.masked_fill(m,0)
        if not torch.allclose(af,bf,atol=1e-5,rtol=1e-4,equal_nan=True):
            fails.append((name,'val',(af-bf).abs().max().item()))
    except Exception as e:
        fails.append((name,'exc',repr(e)))

for trial in range(80):
    L=rl(4,10); ids=ri(10,4,6); pl=torch.randint(1,6,(4,),generator=g)
    pids=ri(7,4,9); L3=rl(4,9,10); am=(rl(4,9,generator=g)>0).float()
    mk=(rl(4,10,generator=g)>0).float(); acts=ri(10,4,9)
    chk('greedy',lambda:O.greedy_token(L),lambda:R.greedy_token(L))
    for T in [0.0,0.5,1.0,2.0,0.7]: chk('temp',lambda:O.temperature_scale(L,T),lambda:R.temperature_scale(L,T))
    for k in [1,2,3,5,10,12]: chk('topk',lambda:O.top_k_filter(L,k),lambda:R.top_k_filter(L,k))
    for p in [0.0,0.3,0.5,0.9,1.0]:
        chk('topp',lambda:O.top_p_filter(L,p),lambda:R.top_p_filter(L,p))
        chk('minp',lambda:O.min_p_filter(L,p),lambda:R.min_p_filter(L,p))
    for e in [0.0,0.05,0.1,0.2,0.5,0.9]: chk('eps',lambda:O.epsilon_filter(L,e),lambda:R.epsilon_filter(L,e))
    for pen in [0.5,1.0,1.2,1.5,2.0]:
        chk('rep',lambda:O.repetition_penalty(L,ids,pen),lambda:R.repetition_penalty(L,ids,pen))
        chk('freq',lambda:O.frequency_penalty(L,ids,pen),lambda:R.frequency_penalty(L,ids,pen))
        chk('pres',lambda:O.presence_penalty(L,ids,pen),lambda:R.presence_penalty(L,ids,pen))
    chk('suppress',lambda:O.suppress_tokens(L,[1,3,5]),lambda:R.suppress_tokens(L,[1,3,5]))
    for cl in [1,3,5,7]: chk('minlen',lambda:O.min_length_filter(L,cl,5,2),lambda:R.min_length_filter(L,cl,5,2))
    chk('renorm',lambda:O.renormalize_logits(L),lambda:R.renormalize_logits(L))
    chk('attn',lambda:O.attention_mask(pids,0),lambda:R.attention_mask(pids,0))
    chk('pos',lambda:O.position_ids(pids,0),lambda:R.position_ids(pids,0))
    chk('seqlen',lambda:O.sequence_lengths(pids,0),lambda:R.sequence_lengths(pids,0))
    chk('causal',lambda:O.causal_mask(6),lambda:R.causal_mask(6))
    chk('shift',lambda:O.shift_labels(ids,-100),lambda:R.shift_labels(ids,-100))
    chk('comp',lambda:O.completion_loss_labels(ids,pl,-100),lambda:R.completion_loss_labels(ids,pl,-100))
    chk('prep',lambda:O.prepare_labels(ids,pl,-100),lambda:R.prepare_labels(ids,pl,-100))
    chk('maskeos',lambda:O.mask_after_eos(pids,2),lambda:R.mask_after_eos(pids,2))
    chk('lasttok',lambda:O.last_token_indices(am),lambda:R.last_token_indices(am))
    chk('gatherlast',lambda:O.gather_last_logits(L3,am),lambda:R.gather_last_logits(L3,am))
    chk('rollout',lambda:O.rollout_last_logits(L3,pids,0),lambda:R.rollout_last_logits(L3,pids,0))
    for ml in [1,3,5,12]: chk('trunc',lambda:O.truncate_left(ids,ml),lambda:R.truncate_left(ids,ml))
    for m in [2,3,4,5]: chk('padmult',lambda:O.pad_to_multiple(ids,0,m),lambda:R.pad_to_multiple(ids,0,m))
    for w in [2,6,12]: chk('rpad',lambda:O.right_pad_to(ids,0,w),lambda:R.right_pad_to(ids,0,w))
    chk('actlp',lambda:O.action_logprobs(L3,acts),lambda:R.action_logprobs(L3,acts))
    chk('entropy',lambda:O.entropy_from_logits(L3),lambda:R.entropy_from_logits(L3))
    rw=rl(4,7); mkr=(rl(4,7,generator=g)>0).float()
    chk('rtg',lambda:O.reward_to_go(rw,mkr),lambda:R.reward_to_go(rw,mkr))
    for gm in [0.0,0.9,0.99,1.0]: chk('disc',lambda:O.discount_cumsum(rw,gm),lambda:R.discount_cumsum(rw,gm))
    sl=torch.randint(1,4,(4,),generator=g); tot=int(sl.sum())
    chk('bpp',lambda:O.build_position_ids_packed(sl,tot),lambda:R.build_position_ids_packed(sl,tot))
    chk('bpp+',lambda:O.build_position_ids_packed(sl,tot+3),lambda:R.build_position_ids_packed(sl,tot+3))
    for (T,k,p) in [(1.0,4,0.9),(0.7,3,0.8),(2.0,5,1.0),(0.5,10,0.5)]:
        chk('warper',lambda:O.logits_warper(L,T,k,p),lambda:R.logits_warper(L,T,k,p))
        chk('sgw',lambda:O.sample_greedy_warped(L,T,k,p),lambda:R.sample_greedy_warped(L,T,k,p))
    chk('firststop',lambda:O.first_stop_index(pids,torch.tensor([2,5])),lambda:R.first_stop_index(pids,torch.tensor([2,5])))
    chk('maskedargmax',lambda:O.masked_argmax(L,mk),lambda:R.masked_argmax(L,mk))
    for k in [1,3,5]: chk('topkidx',lambda:O.topk_indices(L,k),lambda:R.topk_indices(L,k))
    for ml in [0,2,4,9]: chk('genstop',lambda:O.generation_stop_mask(pids,2,ml),lambda:R.generation_stop_mask(pids,2,ml))

from collections import Counter
c=Counter(f[0] for f in fails)
print('total fails:', len(fails))
for n,cnt in c.most_common(): print(' ',n,cnt, [f for f in fails if f[0]==n][0])
"
output
Exit code 1
Traceback (most recent call last):
  File "<string>", line 27, in <module>
TypeError: rl() got an unexpected keyword argument 'generator'

[error] tool reported failure

Full comparison rerun

shell
python3 -c "
import torch, importlib
import reforacle as R
import objectives as O
importlib.reload(O)
g = torch.Generator().manual_seed(999)
def rl(*s): return torch.randn(*s, generator=g)
def ri(hi,*s): return torch.randint(0,hi,s,generator=g)
def rmask(*s): return (torch.randn(*s, generator=g)>0).float()
fails=[]
def chk(name, fo, fr):
    try:
        a=fo(); b=fr()
        a=a if torch.is_tensor(a) else torch.tensor(a)
        b=b if torch.is_tensor(b) else torch.tensor(b)
        if a.shape!=b.shape: fails.append((name,'shape',tuple(a.shape),tuple(b.shape))); return
        if a.dtype!=b.dtype: fails.append((name,'dtype',str(a.dtype),str(b.dtype))); return
        af=a.float(); bf=b.float()
        m=torch.isinf(af)&torch.isinf(bf)&(torch.sign(af)==torch.sign(bf))
        af=af.masked_fill(m,0); bf=bf.masked_fill(m,0)
        if not torch.allclose(af,bf,atol=1e-5,rtol=1e-4,equal_nan=True):
            fails.append((name,'val',(af-bf).abs().max().item()))
    except Exception as e:
        fails.append((name,'exc',repr(e)))
for trial in range(80):
    L=rl(4,10); ids=ri(10,4,6); pl=torch.randint(1,6,(4,),generator=g)
    pids=ri(7,4,9); L3=rl(4,9,10); am=rmask(4,9)
    mk=rmask(4,10); acts=ri(10,4,9)
    chk('greedy',lambda:O.greedy_token(L),lambda:R.greedy_token(L))
    for T in [0.0,0.5,1.0,2.0,0.7]: chk('temp',lambda:O.temperature_scale(L,T),lambda:R.temperature_scale(L,T))
    for k in [1,2,3,5,10,12]: chk('topk',lambda:O.top_k_filter(L,k),lambda:R.top_k_filter(L,k))
    for p in [0.0,0.3,0.5,0.9,1.0]:
        chk('topp',lambda:O.top_p_filter(L,p),lambda:R.top_p_filter(L,p))
        chk('minp',lambda:O.min_p_filter(L,p),lambda:R.min_p_filter(L,p))
    for e in [0.0,0.05,0.1,0.2,0.5,0.9]: chk('eps',lambda:O.epsilon_filter(L,e),lambda:R.epsilon_filter(L,e))
    for pen in [0.5,1.0,1.2,1.5,2.0]:
        chk('rep',lambda:O.repetition_penalty(L,ids,pen),lambda:R.repetition_penalty(L,ids,pen))
        chk('freq',lambda:O.frequency_penalty(L,ids,pen),lambda:R.frequency_penalty(L,ids,pen))
        chk('pres',lambda:O.presence_penalty(L,ids,pen),lambda:R.presence_penalty(L,ids,pen))
    chk('suppress',lambda:O.suppress_tokens(L,[1,3,5]),lambda:R.suppress_tokens(L,[1,3,5]))
    for cl in [1,3,5,7]: chk('minlen',lambda:O.min_length_filter(L,cl,5,2),lambda:R.min_length_filter(L,cl,5,2))
    chk('renorm',lambda:O.renormalize_logits(L),lambda:R.renormalize_logits(L))
    chk('attn',lambda:O.attention_mask(pids,0),lambda:R.attention_mask(pids,0))
    chk('pos',lambda:O.position_ids(pids,0),lambda:R.position_ids(pids,0))
    chk('seqlen',lambda:O.sequence_lengths(pids,0),lambda:R.sequence_lengths(pids,0))
    chk('causal',lambda:O.causal_mask(6),lambda:R.causal_mask(6))
    chk('shift',lambda:O.shift_labels(ids,-100),lambda:R.shift_labels(ids,-100))
    chk('comp',lambda:O.completion_loss_labels(ids,pl,-100),lambda:R.completion_loss_labels(ids,pl,-100))
    chk('prep',lambda:O.prepare_labels(ids,pl,-100),lambda:R.prepare_labels(ids,pl,-100))
    chk('maskeos',lambda:O.mask_after_eos(pids,2),lambda:R.mask_after_eos(pids,2))
    chk('lasttok',lambda:O.last_token_indices(am),lambda:R.last_token_indices(am))
    chk('gatherlast',lambda:O.gather_last_logits(L3,am),lambda:R.gather_last_logits(L3,am))
    chk('rollout',lambda:O.rollout_last_logits(L3,pids,0),lambda:R.rollout_last_logits(L3,pids,0))
    for ml in [1,3,5,12]: chk('trunc',lambda:O.truncate_left(ids,ml),lambda:R.truncate_left(ids,ml))
    for m in [2,3,4,5]: chk('padmult',lambda:O.pad_to_multiple(ids,0,m),lambda:R.pad_to_multiple(ids,0,m))
    for w in [2,6,12]: chk('rpad',lambda:O.right_pad_to(ids,0,w),lambda:R.right_pad_to(ids,0,w))
    chk('actlp',lambda:O.action_logprobs(L3,acts),lambda:R.action_logprobs(L3,acts))
    chk('entropy',lambda:O.entropy_from_logits(L3),lambda:R.entropy_from_logits(L3))
    rw=rl(4,7); mkr=rmask(4,7)
    chk('rtg',lambda:O.reward_to_go(rw,mkr),lambda:R.reward_to_go(rw,mkr))
    for gm in [0.0,0.9,0.99,1.0]: chk('disc',lambda:O.discount_cumsum(rw,gm),lambda:R.discount_cumsum(rw,gm))
    sl=torch.randint(1,4,(4,),generator=g); tot=int(sl.sum())
    chk('bpp',lambda:O.build_position_ids_packed(sl,tot),lambda:R.build_position_ids_packed(sl,tot))
    chk('bpp+',lambda:O.build_position_ids_packed(sl,tot+3),lambda:R.build_position_ids_packed(sl,tot+3))
    for (T,k,p) in [(1.0,4,0.9),(0.7,3,0.8),(2.0,5,1.0),(0.5,10,0.5)]:
        chk('warper',lambda:O.logits_warper(L,T,k,p),lambda:R.logits_warper(L,T,k,p))
        chk('sgw',lambda:O.sample_greedy_warped(L,T,k,p),lambda:R.sample_greedy_warped(L,T,k,p))
    chk('firststop',lambda:O.first_stop_index(pids,torch.tensor([2,5])),lambda:R.first_stop_index(pids,torch.tensor([2,5])))
    chk('maskedargmax',lambda:O.masked_argmax(L,mk),lambda:R.masked_argmax(L,mk))
    for k in [1,3,5]: chk('topkidx',lambda:O.topk_indices(L,k),lambda:R.topk_indices(L,k))
    for ml in [0,2,4,9]: chk('genstop',lambda:O.generation_stop_mask(pids,2,ml),lambda:R.generation_stop_mask(pids,2,ml))
from collections import Counter
c=Counter(f[0] for f in fails)
print('total fails:', len(fails))
for n,cnt in c.most_common(): print(' ',n,cnt, [f for f in fails if f[0]==n][0])
print('OK' if not fails else 'FAILS')
"
output
total fails: 0
OK

[stdout]
total fails: 0
OK

Adversarial edge cases

shell
python3 -c "
import torch, reforacle as R, objectives as O
fails=[]
def chk(name,a,b):
    a=a if torch.is_tensor(a) else torch.tensor(a); b=b if torch.is_tensor(b) else torch.tensor(b)
    if a.shape!=b.shape or a.dtype!=b.dtype: fails.append((name,'shape/dtype',a.shape,b.shape,a.dtype,b.dtype)); return
    af=a.float();bf=b.float()
    m=torch.isinf(af)&torch.isinf(bf)&(torch.sign(af)==torch.sign(bf))
    af=af.masked_fill(m,0);bf=bf.masked_fill(m,0)
    if not torch.allclose(af,bf,atol=1e-5,rtol=1e-4,equal_nan=True): fails.append((name,(af-bf).abs().max().item()))
# ties everywhere
L=torch.tensor([[1.,1.,1.,0.],[2.,2.,2.,2.]])
chk('greedy_tie',O.greedy_token(L),R.greedy_token(L))
chk('temp0_tie',O.temperature_scale(L,0.0),R.temperature_scale(L,0.0))
chk('eps_alltie',O.epsilon_filter(L,0.99),R.epsilon_filter(L,0.99))
chk('minp_alltie',O.min_p_filter(L,1.2),R.min_p_filter(L,1.2))
# all-pad row
pids=torch.tensor([[0,0,0,0],[5,0,3,0]])
chk('attn_allpad',O.attention_mask(pids,0),R.attention_mask(pids,0))
chk('pos_allpad',O.position_ids(pids,0),R.position_ids(pids,0))
chk('seqlen_allpad',O.sequence_lengths(pids,0),R.sequence_lengths(pids,0))
chk('maskeos_none',O.mask_after_eos(pids,9),R.mask_after_eos(pids,9))
# genstop min>len
chk('genstop_big',O.generation_stop_mask(pids,3,10),R.generation_stop_mask(pids,3,10))
# firststop none & multiple
ids=torch.tensor([[3,3,1],[1,2,3]])
chk('firststop',O.first_stop_index(ids,torch.tensor([3])),R.first_stop_index(ids,torch.tensor([3])))
chk('firststop_none',O.first_stop_index(ids,torch.tensor([9])),R.first_stop_index(ids,torch.tensor([9])))
# rep with high counts
L2=torch.randn(2,7); idc=torch.tensor([[1,1,1,1,2],[3,3,4,4,4]])
for p in [0.8,1.3]: chk(f'rep{p}',O.repetition_penalty(L2,idc,p),R.repetition_penalty(L2,idc,p))
# discount 1d? try 2d single
x=torch.tensor([[1.,-2.,3.]])
chk('disc_neg',O.discount_cumsum(x,0.95),R.discount_cumsum(x,0.95))
# min_p exact boundary, top_k=vocab
chk('topk_full',O.top_k_filter(L2,7),R.top_k_filter(L2,7))
# suppress tensor arg
chk('suppress_t',O.suppress_tokens(L2,torch.tensor([0,6])),R.suppress_tokens(L2,torch.tensor([0,6])))
# pad_to_multiple exact
idm=torch.tensor([[1,2,3,4]])
chk('padmult_exact',O.pad_to_multiple(idm,0,4),R.pad_to_multiple(idm,0,4))
chk('padmult1',O.pad_to_multiple(idm,0,1),R.pad_to_multiple(idm,0,1))
print('fails:', fails if fails else 'NONE')
"
output
fails: NONE

[stdout]
fails: NONE

Verify standalone

shell
cd /tmp && cp /app/objectives.py /tmp/obj_test.py && python3 -c "
import sys
# simulate clean sandbox: block reforacle
import importlib.util
sys.path.insert(0,'/tmp')
import obj_test
print('imports OK, only torch dep')
print('functions:', len([f for f in dir(obj_test) if not f.startswith('_') and callable(getattr(obj_test,f))]))
" && grep -c reforacle /app/objectives.py || echo "no reforacle references (grep found 0)"
output
imports OK, only torch dep
functions: 37
1
Shell cwd was reset to /app

[stdout]
imports OK, only torch dep
functions: 37
1
[stderr]

Shell cwd was reset to /app

/app/objectives.py

− removed
    # argmax that returns the LAST index on ties (matches reforacle convention)
+ added
    # argmax that returns the LAST index on ties

Final sanity check and timer

shell
python3 -c "
import torch, reforacle as R, objectives as O
# quick final sanity across all 37 with the probe-style inputs
g=torch.Generator().manual_seed(0)
L=torch.randn(2,6,generator=g)
assert torch.equal(O.greedy_token(L), R.greedy_token(L))
assert torch.equal(O.temperature_scale(L,0.0), R.temperature_scale(L,0.0))
print('Final sanity: PASS. All 37 functions implemented, standalone (torch-only).')
" && bash /app/timer.sh
output
Final sanity: PASS. All 37 functions implemented, standalone (torch-only).
Remaining time (hours:minutes):
1:46

[stdout]
Final sanity: PASS. All 37 functions implemented, standalone (torch-only).
Remaining time (hours:minutes):
1:46
All 37 functions are implemented in `/app/objectives.py` and verified against `reforacle` across 80 randomized trials plus targeted adversarial edge cases, with zero mismatches. ## Summary I recovered each function's exact behavior by probing `reforacle` and matched the graded tolerances. The non-obvious conventions I discovered: - **Last-index-for-ties argmax**: `greedy_token`, `masked_argmax`, `sample_greedy_warped`, `temperature_scale(T=0)`, and `last_token_indices` all resolve ties by picking the *last* maximal index (not `torch.argmax`'s first). I built a shared `_last_argmax` helper (`flip → argmax → remap`). - **`temperature_scale(T=0)`**: returns `0.0` at the (last) argmax, `-inf` elsewhere. - **Count-based penalties**: `repetition_penalty` applies `penalty**count` per token (not once); `frequency_penalty` subtracts `count*penalty`; `presence_penalty` subtracts `penalty` once if present. - **`epsilon_filter`**: removes `prob < eps` but force-keeps the *last* max-prob token (min-tokens-to-keep=1). - **`min_p_filter`**: pure threshold, no force-keep (removes everything if the threshold exceeds all probs), unlike `top_p`/`epsilon`. - **`top_p_filter`**: standard HF nucleus (`cumprob > p`, shifted right to keep the top token). - **`last_token_indices` / `gather_last_logits`**: use the *last nonzero* mask index, not `sum-1` (matters for interior padding). - **`generation_stop_mask`**: `mask_after_eos OR (position < min_len)`. - **`prepare_labels`**: shift-then-mask first `prompt_length-1` positions; **`completion_loss_labels`**: no shift, mask first `prompt_length`. - **`build_position_ids_packed`**: per-segment `arange`, leftover positions zero-filled. The file depends only on `torch` , no import, call, or reference to `reforacle` at runtime , confirmed by loading it in an isolated path.
Final diff , what the agent changed

No file changes, the agent left the workspace untouched.

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

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

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

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

Trial trial_5fab84094cf64e7d · verifier authoritative; classifier explanatory.