SyncValsverifier → artifact → classifier → verdict
SyncVals · Trajectory

reverse-engineer-decoding

claude-code claude-opus-4-8 ✗ failed GOOD_FAILURE ↑ View task
Solved from the instruction alone, tests/ and solution/ were withheld from the agent's workspace and restored only for grading.
Reward = tests/test.sh exit code (0 → resolved); the classification below is post-hoc and cannot change it.
Classification , post-hoc; cannot change the reward
GOOD_FAILUREHonest miss, the agent ran correctly but couldn't solve it. Expected for a hard task; the task is sound.
SubtypeIncomplete Exploration
EvidenceTest output: 'FUNCTIONS PASSED: 35/37 (partial_score=0.9459)'; 'REWARD = 0 (1 iff ALL 37 functions match the sealed reference exactly)'. Agent trajectory shows: (1) Probed 37 functions systematically with PYTHONPATH=/app; (2) Implemented all 37 with careful tie-breaking logic (e.g., _argmax_last for last-index breaking); (3) Verified on 100+ randomized trials with various shapes; (4) Confirmed standalone execution. Two functions out of 37 fail on held-out grading inputs. Instruction stated: '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.'
Root causeThe agent achieved 94.59% pass rate by reverse-engineering 35 of 37 functions through probing, but failed to perfectly recover the reference behavior for 2 functions on held-out test inputs. The task requires 100% exact match; the instruction explicitly warned that agreement on probed inputs does not guarantee agreement on grading inputs.
RecommendationN/A - task is fine. This is a hard reverse-engineering task by design. The agent's near-complete solution (35/37) demonstrates both strong problem-solving and the inherent challenge of recovering sealed oracle behavior. No task modification needed.
Trajectory
Tool-by-tool agent trajectory
35 tool calls · 4 tool types · 37 steps
Reimplement the post-training decoding-and-collation library in `/app/objectives.py` so that every function reproduces the sealed reference module `reforacle` exactly. `reforacle` is importable and callable (`import reforacle`) but its source cannot be read. It is the only ground truth; there is no written specification of what any function computes or of the conventions it follows. Recover its behavior by calling it. `/app/objectives.py` contains thirty-seven function stubs , logit processors and sampling filters, sequence collation and masking utilities, and rollout helpers used to generate and batch data for RLHF and SFT. Keep every signature exactly as given and implement every function. The names indicate the standard quantity each is meant to produce, but the reference's exact definition is whatever `reforacle` returns , not what you assume. Determine each definition by probing. Grading compares your implementation to `reforacle` on held-out inputs you do not see, function by function, requiring exact numerical agreement (primitives within `atol=1e-6, rtol=1e-5`; composite functions within `atol=1e-5, rtol=1e-4`). Each function passes only if it agrees on every graded input, and the task passes only if all thirty-seven functions pass. Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on. Grading runs your `/app/objectives.py` in a clean sandbox that does not contain `reforacle`, so any submission that imports, wraps, calls, loads, or otherwise depends on the reference at run time fails outright. Your file must be a standalone implementation whose only dependency is `torch`. Do not read, decompile, disassemble, or otherwise extract the compiled reference; recover its behavior only by calling it and comparing. `python3 /app/probe.py` shows the calling pattern; `bash /app/timer.sh` reports remaining budget.

/app/objectives.py

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

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

Remaining time (hours:minutes):
1:59

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

/app/probe.py

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

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

Remaining time (hours:minutes):
1:59

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

Check remaining budget

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

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

Remaining time (hours:minutes):
1:59

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

Run the probe example

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.
['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']

List reforacle functions

shell
python3 -c "import reforacle; print([x for x in dir(reforacle) if not x.startswith('__')])"
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.
['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']

Check constants

shell
python3 -c "
import reforacle
print('NEG_INF', repr(reforacle.NEG_INF))
print('F', repr(reforacle.F))
"
output
NEG_INF -inf
F <module 'torch.nn.functional' from '/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py'>

[stdout]
NEG_INF -inf
F <module 'torch.nn.functional' from '/usr/local/lib/python3.12/dist-packages/torch/nn/functional.py'>

Probe logit processors group 1

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

def show(name, out):
    if isinstance(out, torch.Tensor):
        print(name, out.dtype, tuple(out.shape))
        print("  ", out.tolist())
    else:
        print(name, type(out), out)

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

# greedy_token
show("greedy_token", R.greedy_token(logits))
# temperature_scale normal
show("temp T=2", R.temperature_scale(logits, 2.0))
show("temp T=1", R.temperature_scale(logits, 1.0))
show("temp T=0.5", R.temperature_scale(logits, 0.5))
# top_k
show("top_k k=3", R.top_k_filter(logits, 3))
show("top_k k=1", R.top_k_filter(logits, 1))
show("top_k k=10", R.top_k_filter(logits, 10))
# top_p
show("top_p .9", R.top_p_filter(logits, 0.9))
show("top_p .5", R.top_p_filter(logits, 0.5))
show("top_p 1.0", R.top_p_filter(logits, 1.0))
show("top_p 0.0", R.top_p_filter(logits, 0.0))
# min_p
show("min_p .3", R.min_p_filter(logits, 0.3))
show("min_p .1", R.min_p_filter(logits, 0.1))
EOF
python3 /tmp/p1.py
output
Exit code 1
Traceback (most recent call last):
  File "/tmp/p1.py", line 1, in <module>
    import torch, reforacle as R
ModuleNotFoundError: No module named 'reforacle'

[error] tool reported failure

Probe from /app dir

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

[error] tool reported failure

Find reforacle location

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

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

Probe with PYTHONPATH set

shell
PYTHONPATH=/app python3 /tmp/p1.py
output
LOGITS [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
greedy_token torch.int64 (2,)
   [0, 1]
temp T=2 torch.float32 (2, 6)
   [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
temp T=1 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
temp T=0.5 torch.float32 (2, 6)
   [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854, -2.7971909046173096], [0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828, -1.1932706832885742, 0.3640729784965515]]
top_k k=3 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
top_k k=1 torch.float32 (2, 6)
   [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k k=10 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p .9 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p .5 torch.float32 (2, 6)
   [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_p 1.0 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p 0.0 torch.float32 (2, 6)
   [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
min_p .3 torch.float32 (2, 6)
   [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
min_p .1 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]

[stdout]
LOGITS [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
greedy_token torch.int64 (2,)
   [0, 1]
temp T=2 torch.float32 (2, 6)
   [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
temp T=1 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
temp T=0.5 torch.float32 (2, 6)
   [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854, -2.7971909046173096], [0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828, -1.1932706832885742, 0.3640729784965515]]
top_k k=3 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
top_k k=1 torch.float32 (2, 6)
   [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k k=10 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p .9 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p .5 torch.float32 (2, 6)
   [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_p 1.0 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p 0.0 torch.float32 (2, 6)
   [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
min_p .3 torch.float32 (2, 6)
   [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
min_p .1 torch.float32 (2, 6)
   [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]

Probe penalties/suppress/renorm

shell
cat > /tmp/p2.py << 'EOF'
import torch, reforacle as R
def show(name, out):
    if isinstance(out, torch.Tensor):
        print(name, out.dtype, tuple(out.shape), out.tolist())
    else:
        print(name, repr(out))
g=torch.Generator().manual_seed(1)
logits=torch.randn(2,6,generator=g)
ids=torch.tensor([[0,1,1,3],[2,2,2,5]])
print("LOGITS",logits.tolist())
show("rep_pen 1.5", R.repetition_penalty(logits, ids, 1.5))
show("freq_pen 0.5", R.frequency_penalty(logits, ids, 0.5))
show("pres_pen 0.5", R.presence_penalty(logits, ids, 0.5))
show("suppress [1,3]", R.suppress_tokens(logits, [1,3]))
show("suppress tensor", R.suppress_tokens(logits, torch.tensor([0,5])))
show("min_length cur=2 min=5 eos=4", R.min_length_filter(logits, 2, 5, 4))
show("min_length cur=5 min=5 eos=4", R.min_length_filter(logits, 5, 5, 4))
show("renormalize", R.renormalize_logits(logits))
print("sum exp renorm row0", torch.logsumexp(R.renormalize_logits(logits),dim=-1))
EOF
PYTHONPATH=/app python3 /tmp/p2.py
output
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]]
rep_pen 1.5 torch.float32 (2, 6) [[0.44090142846107483, 0.11863294243812561, 0.06167725846171379, 0.41421154141426086, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -3.468179225921631, -0.563052773475647, -0.8922905325889587, -0.08737526834011078]]
freq_pen 0.5 torch.float32 (2, 6) [[0.16135215759277344, -0.7330758571624756, 0.06167725846171379, 0.12131732702255249, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -2.527608633041382, -0.563052773475647, -0.8922905325889587, -0.5582501888275146]]
pres_pen 0.5 torch.float32 (2, 6) [[0.16135215759277344, -0.23307588696479797, 0.06167725846171379, 0.12131732702255249, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.5276086330413818, -0.563052773475647, -0.8922905325889587, -0.5582501888275146]]
suppress [1,3] torch.float32 (2, 6) [[0.6613521575927734, -inf, 0.06167725846171379, -inf, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, -inf, -1.0276086330413818, -inf, -0.8922905325889587, -0.05825017765164375]]
suppress tensor torch.float32 (2, 6) [[-inf, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -inf], [-inf, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -inf]]
min_length cur=2 min=5 eos=4 torch.float32 (2, 6) [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -inf, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -inf, -0.05825017765164375]]
min_length cur=5 min=5 eos=4 torch.float32 (2, 6) [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
renormalize torch.float32 (2, 6) [[-1.3735755681991577, -1.7680035829544067, -1.973250389099121, -1.4136104583740234, -2.4868335723876953, -2.2010579109191895], [-2.8998494148254395, -0.9953970909118652, -2.4046897888183594, -1.940133810043335, -2.269371509552002, -1.4353312253952026]]
sum exp renorm row0 tensor([0., 0.])

[stdout]
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]]
rep_pen 1.5 torch.float32 (2, 6) [[0.44090142846107483, 0.11863294243812561, 0.06167725846171379, 0.41421154141426086, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -3.468179225921631, -0.563052773475647, -0.8922905325889587, -0.08737526834011078]]
freq_pen 0.5 torch.float32 (2, 6) [[0.16135215759277344, -0.7330758571624756, 0.06167725846171379, 0.12131732702255249, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -2.527608633041382, -0.563052773475647, -0.8922905325889587, -0.5582501888275146]]
pres_pen 0.5 torch.float32 (2, 6) [[0.16135215759277344, -0.23307588696479797, 0.06167725846171379, 0.12131732702255249, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.5276086330413818, -0.563052773475647, -0.8922905325889587, -0.5582501888275146]]
suppress [1,3] torch.float32 (2, 6) [[0.6613521575927734, -inf, 0.06167725846171379, -inf, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, -inf, -1.0276086330413818, -inf, -0.8922905325889587, -0.05825017765164375]]
suppress tensor torch.float32 (2, 6) [[-inf, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -inf], [-inf, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -inf]]
min_length cur=2 min=5 eos=4 torch.float32 (2, 6) [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -inf, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -inf, -0.05825017765164375]]
min_length cur=5 min=5 eos=4 torch.float32 (2, 6) [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
renormalize torch.float32 (2, 6) [[-1.3735755681991577, -1.7680035829544067, -1.973250389099121, -1.4136104583740234, -2.4868335723876953, -2.2010579109191895], [-2.8998494148254395, -0.9953970909118652, -2.4046897888183594, -1.940133810043335, -2.269371509552002, -1.4353312253952026]]
sum exp renorm row0 tensor([0., 0.])

Probe masking group

shell
cat > /tmp/p3.py << 'EOF'
import torch, reforacle as R
def show(name, out):
    if isinstance(out, torch.Tensor):
        print(name, out.dtype, tuple(out.shape), out.tolist())
    elif isinstance(out,(tuple,list)):
        print(name, "SEQ", out)
    else:
        print(name, repr(out))
ids=torch.tensor([[5,6,7,0,0],[1,2,3,4,8]])
print("ids", ids.tolist())
show("attention_mask pad=0", R.attention_mask(ids, 0))
show("position_ids pad=0", R.position_ids(ids, 0))
show("causal_mask t=4", R.causal_mask(4))
show("shift_labels ig=-100", R.shift_labels(ids, -100))
show("completion_loss_labels pl=[2,3] ig=-100", R.completion_loss_labels(ids, torch.tensor([2,3]), -100))
ids2=torch.tensor([[5,6,9,7,3],[1,9,3,4,8]])
show("mask_after_eos eos=9", R.mask_after_eos(ids2, 9))
show("sequence_lengths pad=0", R.sequence_lengths(ids, 0))
am=R.attention_mask(ids,0)
show("last_token_indices", R.last_token_indices(am))
EOF
PYTHONPATH=/app python3 /tmp/p3.py
output
ids [[5, 6, 7, 0, 0], [1, 2, 3, 4, 8]]
attention_mask pad=0 torch.float32 (2, 5) [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
position_ids pad=0 torch.int64 (2, 5) [[0, 1, 2, 2, 2], [0, 1, 2, 3, 4]]
causal_mask t=4 torch.float32 (4, 4) [[0.0, -inf, -inf, -inf], [0.0, 0.0, -inf, -inf], [0.0, 0.0, 0.0, -inf], [0.0, 0.0, 0.0, 0.0]]
shift_labels ig=-100 torch.int64 (2, 5) [[6, 7, 0, 0, -100], [2, 3, 4, 8, -100]]
completion_loss_labels pl=[2,3] ig=-100 torch.int64 (2, 5) [[-100, -100, 7, 0, 0], [-100, -100, -100, 4, 8]]
mask_after_eos eos=9 torch.float32 (2, 5) [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]]
sequence_lengths pad=0 torch.int64 (2,) [3, 5]
last_token_indices torch.int64 (2,) [2, 4]

[stdout]
ids [[5, 6, 7, 0, 0], [1, 2, 3, 4, 8]]
attention_mask pad=0 torch.float32 (2, 5) [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
position_ids pad=0 torch.int64 (2, 5) [[0, 1, 2, 2, 2], [0, 1, 2, 3, 4]]
causal_mask t=4 torch.float32 (4, 4) [[0.0, -inf, -inf, -inf], [0.0, 0.0, -inf, -inf], [0.0, 0.0, 0.0, -inf], [0.0, 0.0, 0.0, 0.0]]
shift_labels ig=-100 torch.int64 (2, 5) [[6, 7, 0, 0, -100], [2, 3, 4, 8, -100]]
completion_loss_labels pl=[2,3] ig=-100 torch.int64 (2, 5) [[-100, -100, 7, 0, 0], [-100, -100, -100, 4, 8]]
mask_after_eos eos=9 torch.float32 (2, 5) [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]]
sequence_lengths pad=0 torch.int64 (2,) [3, 5]
last_token_indices torch.int64 (2,) [2, 4]

Probe group 4

shell
cat > /tmp/p4.py << 'EOF'
import torch, reforacle as R
def show(name, out):
    if isinstance(out, torch.Tensor):
        print(name, out.dtype, tuple(out.shape), out.tolist())
    elif isinstance(out,(tuple,list)):
        print(name, "SEQ", [ (o.tolist() if isinstance(o,torch.Tensor) else o) for o in out])
    else:
        print(name, repr(out))
g=torch.Generator().manual_seed(2)
logits=torch.randn(2,4,5,generator=g)  # B,T,V
am=torch.tensor([[1,1,1,0],[1,1,0,0]])
show("gather_last_logits", R.gather_last_logits(logits, am))
# truncate_left
ids=torch.tensor([[1,2,3,4,5,6],[7,8,9,10,11,12]])
show("truncate_left max=3", R.truncate_left(ids, 3))
show("truncate_left max=10", R.truncate_left(ids, 10))
# pad_to_multiple
ids2=torch.tensor([[1,2,3],[4,5,6]])
show("pad_to_multiple pad=0 mult=4", R.pad_to_multiple(ids2,0,4))
show("pad_to_multiple pad=0 mult=2", R.pad_to_multiple(ids2,0,2))
# action_logprobs
lg=torch.randn(2,3,5,generator=g)
act=torch.tensor([[1,2,0],[3,4,1]])
show("action_logprobs", R.action_logprobs(lg,act))
# entropy
show("entropy", R.entropy_from_logits(torch.randn(2,5,generator=g)))
EOF
PYTHONPATH=/app python3 /tmp/p4.py
output
gather_last_logits torch.float32 (2, 5) [[-0.5713335871696472, -0.2351086586713791, 1.0076055526733398, -0.7528814673423767, -0.22499920427799225], [-0.2656300365924835, 0.16654501855373383, 0.7910981774330139, 0.8560096621513367, 0.40943124890327454]]
truncate_left max=3 torch.int64 (2, 3) [[4, 5, 6], [10, 11, 12]]
truncate_left max=10 torch.int64 (2, 6) [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]]
pad_to_multiple pad=0 mult=4 torch.int64 (2, 4) [[1, 2, 3, 0], [4, 5, 6, 0]]
pad_to_multiple pad=0 mult=2 torch.int64 (2, 4) [[1, 2, 3, 0], [4, 5, 6, 0]]
action_logprobs torch.float32 (2, 3) [[-4.543883800506592, -0.415120393037796, -3.1971359252929688], [-1.3026783466339111, -0.6381038427352905, -1.658724069595337]]
entropy torch.float32 (2,) [1.538684368133545, 1.4150136709213257]

[stdout]
gather_last_logits torch.float32 (2, 5) [[-0.5713335871696472, -0.2351086586713791, 1.0076055526733398, -0.7528814673423767, -0.22499920427799225], [-0.2656300365924835, 0.16654501855373383, 0.7910981774330139, 0.8560096621513367, 0.40943124890327454]]
truncate_left max=3 torch.int64 (2, 3) [[4, 5, 6], [10, 11, 12]]
truncate_left max=10 torch.int64 (2, 6) [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]]
pad_to_multiple pad=0 mult=4 torch.int64 (2, 4) [[1, 2, 3, 0], [4, 5, 6, 0]]
pad_to_multiple pad=0 mult=2 torch.int64 (2, 4) [[1, 2, 3, 0], [4, 5, 6, 0]]
action_logprobs torch.float32 (2, 3) [[-4.543883800506592, -0.415120393037796, -3.1971359252929688], [-1.3026783466339111, -0.6381038427352905, -1.658724069595337]]
entropy torch.float32 (2,) [1.538684368133545, 1.4150136709213257]

Probe group 5

shell
cat > /tmp/p5.py << 'EOF'
import torch, reforacle as R
def show(name, out):
    if isinstance(out, torch.Tensor):
        print(name, out.dtype, tuple(out.shape), out.tolist())
    elif isinstance(out,(tuple,list)):
        print(name, "SEQ", [ (o.tolist() if isinstance(o,torch.Tensor) else o) for o in out])
    else:
        print(name, repr(out))
# reward_to_go
rewards=torch.tensor([[1.,2.,3.,4.],[1.,1.,1.,1.]])
mask=torch.tensor([[1.,1.,1.,0.],[1.,1.,0.,0.]])
show("reward_to_go", R.reward_to_go(rewards, mask))
# discount_cumsum
x=torch.tensor([[1.,2.,3.,4.]])
show("discount_cumsum g=0.9", R.discount_cumsum(x, 0.9))
show("discount_cumsum g=1.0", R.discount_cumsum(x, 1.0))
# build_position_ids_packed
show("build_pos_packed [2,3] tot=5", R.build_position_ids_packed(torch.tensor([2,3]), 5))
show("build_pos_packed [1,4] tot=5", R.build_position_ids_packed(torch.tensor([1,4]), 5))
# epsilon_filter
g=torch.Generator().manual_seed(3)
logits=torch.randn(2,6,generator=g)
print("LOGITS",logits.tolist())
show("epsilon_filter 0.1", R.epsilon_filter(logits, 0.1))
show("epsilon_filter 0.05", R.epsilon_filter(logits, 0.05))
# first_stop_index
ids=torch.tensor([[5,6,9,7,3],[1,2,3,4,8]])
show("first_stop_index [9,3]", R.first_stop_index(ids, [9,3]))
show("first_stop_index [99]", R.first_stop_index(ids, [99]))
# masked_argmax
mask2=torch.tensor([[1.,0.,1.,0.,1.,0.],[0.,0.,1.,1.,1.,1.]])
show("masked_argmax", R.masked_argmax(logits, mask2))
# topk_indices
show("topk_indices k=3", R.topk_indices(logits, 3))
EOF
PYTHONPATH=/app python3 /tmp/p5.py
output
Exit code 1
reward_to_go torch.float32 (2, 4) [[6.0, 5.0, 3.0, 0.0], [2.0, 1.0, 0.0, 0.0]]
discount_cumsum g=0.9 torch.float32 (1, 4) [[8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]]
discount_cumsum g=1.0 torch.float32 (1, 4) [[10.0, 9.0, 7.0, 4.0]]
build_pos_packed [2,3] tot=5 torch.int64 (5,) [0, 1, 0, 1, 2]
build_pos_packed [1,4] tot=5 torch.int64 (5,) [0, 0, 1, 2, 3]
LOGITS [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -0.6137180328369141, 0.04618244990706444, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -1.435179352760315, 0.9774317741394043, 0.5220499634742737, 1.2379167079925537]]
epsilon_filter 0.1 torch.float32 (2, 6) [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -inf, 0.04618244990706444, -inf], [0.3374950885772705, 1.0111159086227417, -inf, 0.9774317741394043, 0.5220499634742737, 1.2379167079925537]]
epsilon_filter 0.05 torch.float32 (2, 6) [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -0.6137180328369141, 0.04618244990706444, -inf], [0.3374950885772705, 1.0111159086227417, -inf, 0.9774317741394043, 0.5220499634742737, 1.2379167079925537]]
Traceback (most recent call last):
  File "/tmp/p5.py", line 28, in <module>
    show("first_stop_index [9,3]", R.first_stop_index(ids, [9,3]))
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "reforacle.py", line 243, in reforacle.first_stop_index
AttributeError: 'list' object has no attribute 'tolist'

[error] tool reported failure

Probe group 6

shell
cat > /tmp/p6.py << 'EOF'
import torch, reforacle as R
def show(name, out):
    if isinstance(out, torch.Tensor):
        print(name, out.dtype, tuple(out.shape), out.tolist())
    else:
        print(name, repr(out))
g=torch.Generator().manual_seed(3)
logits=torch.randn(2,6,generator=g)
ids=torch.tensor([[5,6,9,7,3],[1,2,3,4,8]])
show("first_stop_index tensor[9,3]", R.first_stop_index(ids, torch.tensor([9,3])))
show("first_stop_index tensor[99]", R.first_stop_index(ids, torch.tensor([99])))
mask2=torch.tensor([[1.,0.,1.,0.,1.,0.],[0.,0.,1.,1.,1.,1.]])
show("masked_argmax", R.masked_argmax(logits, mask2))
show("topk_indices k=3", R.topk_indices(logits, 3))
show("topk_indices k=1", R.topk_indices(logits, 1))
# right_pad_to
ids2=torch.tensor([[1,2,3],[4,5,6]])
show("right_pad_to pad=0 w=5", R.right_pad_to(ids2,0,5))
show("right_pad_to pad=0 w=2", R.right_pad_to(ids2,0,2))
# generation_stop_mask
ids3=torch.tensor([[5,6,4,7,3],[1,4,3,4,8]])
show("gen_stop_mask eos=4 min=0", R.generation_stop_mask(ids3,4,0))
show("gen_stop_mask eos=4 min=3", R.generation_stop_mask(ids3,4,3))
# reward_to_go middle mask test
rewards=torch.tensor([[1.,2.,3.,4.]])
maskm=torch.tensor([[1.,0.,1.,1.]])
show("reward_to_go midmask", R.reward_to_go(rewards, maskm))
# build_position_ids_packed total>sum
show("build_pos_packed [2,2] tot=5", R.build_position_ids_packed(torch.tensor([2,2]), 5))
EOF
PYTHONPATH=/app python3 /tmp/p6.py
output
first_stop_index tensor[9,3] torch.int64 (2,) [2, 2]
first_stop_index tensor[99] torch.int64 (2,) [-1, -1]
masked_argmax torch.int64 (2,) [0, 5]
topk_indices k=3 torch.int64 (2, 3) [[0, 1, 2], [5, 1, 3]]
topk_indices k=1 torch.int64 (2, 1) [[0], [5]]
right_pad_to pad=0 w=5 torch.int64 (2, 5) [[1, 2, 3, 0, 0], [4, 5, 6, 0, 0]]
right_pad_to pad=0 w=2 torch.int64 (2, 2) [[1, 2], [4, 5]]
gen_stop_mask eos=4 min=0 torch.float32 (2, 5) [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]]
gen_stop_mask eos=4 min=3 torch.float32 (2, 5) [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0, 0.0]]
reward_to_go midmask torch.float32 (1, 4) [[8.0, 0.0, 7.0, 4.0]]
build_pos_packed [2,2] tot=5 torch.int64 (5,) [0, 1, 0, 1, 0]

[stdout]
first_stop_index tensor[9,3] torch.int64 (2,) [2, 2]
first_stop_index tensor[99] torch.int64 (2,) [-1, -1]
masked_argmax torch.int64 (2,) [0, 5]
topk_indices k=3 torch.int64 (2, 3) [[0, 1, 2], [5, 1, 3]]
topk_indices k=1 torch.int64 (2, 1) [[0], [5]]
right_pad_to pad=0 w=5 torch.int64 (2, 5) [[1, 2, 3, 0, 0], [4, 5, 6, 0, 0]]
right_pad_to pad=0 w=2 torch.int64 (2, 2) [[1, 2], [4, 5]]
gen_stop_mask eos=4 min=0 torch.float32 (2, 5) [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]]
gen_stop_mask eos=4 min=3 torch.float32 (2, 5) [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0, 0.0]]
reward_to_go midmask torch.float32 (1, 4) [[8.0, 0.0, 7.0, 4.0]]
build_pos_packed [2,2] tot=5 torch.int64 (5,) [0, 1, 0, 1, 0]

Probe group 7

shell
cat > /tmp/p7.py << 'EOF'
import torch, reforacle as R
def show(name, out):
    if isinstance(out, torch.Tensor):
        print(name, out.dtype, tuple(out.shape), out.tolist())
    else:
        print(name, repr(out))
g=torch.Generator().manual_seed(4)
logits=torch.randn(2,6,generator=g)
print("LOGITS",logits.tolist())
show("logits_warper T=2,k=4,p=0.9", R.logits_warper(logits,2.0,4,0.9))
show("sample_greedy_warped T=2,k=4,p=0.9", R.sample_greedy_warped(logits,2.0,4,0.9))
show("sample_greedy_warped T=1,k=0,p=1.0", R.sample_greedy_warped(logits,1.0,0,1.0))
# prepare_labels vs completion_loss_labels
ids=torch.tensor([[5,6,7,8,9],[1,2,3,4,0]])
show("prepare_labels pl=[2,3]", R.prepare_labels(ids, torch.tensor([2,3]), -100))
show("completion_loss_labels pl=[2,3]", R.completion_loss_labels(ids, torch.tensor([2,3]), -100))
# rollout_last_logits
L=torch.randn(2,5,7,generator=g)
ids4=torch.tensor([[5,6,7,0,0],[1,2,3,4,8]])
show("rollout_last_logits pad=0", R.rollout_last_logits(L, ids4, 0))
# disambiguate gen_stop_mask
show("gsm [4,7,8,9,9] eos4 min3", R.generation_stop_mask(torch.tensor([[4,7,8,9,9]]),4,3))
show("gsm noeos min0", R.generation_stop_mask(torch.tensor([[5,6,7,8,9]]),4,0))
# build_pos leftover
show("build_pos [2,2] tot6", R.build_position_ids_packed(torch.tensor([2,2]),6))
show("build_pos [3] tot5", R.build_position_ids_packed(torch.tensor([3]),5))
EOF
PYTHONPATH=/app python3 /tmp/p7.py
output
Exit code 1
LOGITS [[-1.605276346206665, 0.23248571157455444, 2.239870071411133, 0.8472937941551208, 1.2006442546844482, -0.4015503227710724], [-1.4260196685791016, 0.903931736946106, 0.8557155728340149, 0.6888809204101562, 0.8849856853485107, 1.7706429958343506]]
logits_warper T=2,k=4,p=0.9 torch.float32 (2, 6) [[-inf, 0.11624285578727722, 1.1199350357055664, 0.4236468970775604, 0.6003221273422241, -inf], [-inf, 0.451965868473053, 0.42785778641700745, -inf, 0.44249284267425537, 0.8853214979171753]]
sample_greedy_warped T=2,k=4,p=0.9 torch.int64 (2,) [2, 5]
Traceback (most recent call last):
  File "/tmp/p7.py", line 12, in <module>
    show("sample_greedy_warped T=1,k=0,p=1.0", R.sample_greedy_warped(logits,1.0,0,1.0))
                                               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "reforacle.py", line 217, in reforacle.sample_greedy_warped
  File "reforacle.py", line 211, in reforacle.logits_warper
  File "reforacle.py", line 30, in reforacle.top_k_filter
RuntimeError: The size of tensor a (6) must match the size of tensor b (0) at non-singleton dimension 1

[error] tool reported failure

Probe group 8

shell
cat > /tmp/p8.py << 'EOF'
import torch, reforacle as R
def show(name, out):
    if isinstance(out, torch.Tensor):
        print(name, out.dtype, tuple(out.shape), out.tolist())
    else:
        print(name, repr(out))
g=torch.Generator().manual_seed(4)
logits=torch.randn(2,6,generator=g)
ids=torch.tensor([[5,6,7,8,9],[1,2,3,4,0]])
show("prepare_labels pl=[2,3]", R.prepare_labels(ids, torch.tensor([2,3]), -100))
show("completion_loss_labels pl=[2,3]", R.completion_loss_labels(ids, torch.tensor([2,3]), -100))
L=torch.randn(2,5,7,generator=g)
ids4=torch.tensor([[5,6,7,0,0],[1,2,3,4,8]])
show("rollout_last_logits pad=0", R.rollout_last_logits(L, ids4, 0))
show("gsm [4,7,8,9,9] eos4 min3", R.generation_stop_mask(torch.tensor([[4,7,8,9,9]]),4,3))
show("gsm noeos min0", R.generation_stop_mask(torch.tensor([[5,6,7,8,9]]),4,0))
show("build_pos [2,2] tot6", R.build_position_ids_packed(torch.tensor([2,2]),6))
show("build_pos [3] tot5", R.build_position_ids_packed(torch.tensor([3]),5))
# prepare_labels: shifted or not? compare to shift
show("prepare_labels pl=[0,0]", R.prepare_labels(ids, torch.tensor([0,0]), -100))
EOF
PYTHONPATH=/app python3 /tmp/p8.py
output
prepare_labels pl=[2,3] torch.int64 (2, 5) [[-100, 7, 8, 9, -100], [-100, -100, 4, 0, -100]]
completion_loss_labels pl=[2,3] torch.int64 (2, 5) [[-100, -100, 7, 8, 9], [-100, -100, -100, 4, 0]]
rollout_last_logits pad=0 torch.float32 (2, 7) [[-1.2768372297286987, -1.4579845666885376, -0.7042272090911865, 2.012578248977661, 0.9120109677314758, 0.27727454900741577, 0.2083573192358017], [0.24560418725013733, 0.16178056597709656, -0.9001774191856384, 0.18334539234638214, 0.5111587047576904, 0.29877257347106934, 1.5619761943817139]]
gsm [4,7,8,9,9] eos4 min3 torch.float32 (1, 5) [[1.0, 1.0, 1.0, 0.0, 0.0]]
gsm noeos min0 torch.float32 (1, 5) [[1.0, 1.0, 1.0, 1.0, 1.0]]
build_pos [2,2] tot6 torch.int64 (6,) [0, 1, 0, 1, 0, 0]
build_pos [3] tot5 torch.int64 (5,) [0, 1, 2, 0, 0]
prepare_labels pl=[0,0] torch.int64 (2, 5) [[6, 7, 8, 9, -100], [2, 3, 4, 0, -100]]

[stdout]
prepare_labels pl=[2,3] torch.int64 (2, 5) [[-100, 7, 8, 9, -100], [-100, -100, 4, 0, -100]]
completion_loss_labels pl=[2,3] torch.int64 (2, 5) [[-100, -100, 7, 8, 9], [-100, -100, -100, 4, 0]]
rollout_last_logits pad=0 torch.float32 (2, 7) [[-1.2768372297286987, -1.4579845666885376, -0.7042272090911865, 2.012578248977661, 0.9120109677314758, 0.27727454900741577, 0.2083573192358017], [0.24560418725013733, 0.16178056597709656, -0.9001774191856384, 0.18334539234638214, 0.5111587047576904, 0.29877257347106934, 1.5619761943817139]]
gsm [4,7,8,9,9] eos4 min3 torch.float32 (1, 5) [[1.0, 1.0, 1.0, 0.0, 0.0]]
gsm noeos min0 torch.float32 (1, 5) [[1.0, 1.0, 1.0, 1.0, 1.0]]
build_pos [2,2] tot6 torch.int64 (6,) [0, 1, 0, 1, 0, 0]
build_pos [3] tot5 torch.int64 (5,) [0, 1, 2, 0, 0]
prepare_labels pl=[0,0] torch.int64 (2, 5) [[6, 7, 8, 9, -100], [2, 3, 4, 0, -100]]
write · /app/objectives.py
from __future__ import annotations

import torch

NEG_INF = float("-inf")


def _token_counts(logits, input_ids):
    vocab = logits.shape[-1]
    counts = torch.zeros(
        input_ids.shape[0], vocab, dtype=logits.dtype, device=logits.device
    )
    counts.scatter_add_(1, input_ids, torch.ones_like(input_ids, dtype=logits.dtype))
    return counts


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):
    kth = torch.topk(logits, k, dim=-1)[0][..., -1, None]
    return logits.masked_fill(logits < kth, NEG_INF)


def top_p_filter(logits, p):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
    cum_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
    sorted_remove = cum_probs > 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_prob = probs.max(dim=-1, keepdim=True)[0]
    remove = probs < (p * top_prob)
    return logits.masked_fill(remove, NEG_INF)


def repetition_penalty(logits, input_ids, penalty):
    counts = _token_counts(logits, input_ids)
    factor = penalty ** counts
    return torch.where(logits < 0, logits * factor, logits / factor)


def frequency_penalty(logits, input_ids, penalty):
    counts = _token_counts(logits, input_ids)
    return logits - penalty * counts


def presence_penalty(logits, input_ids, penalty):
    counts = _token_counts(logits, input_ids)
    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:
        return logits.index_fill(
            -1, torch.tensor([eos_id], device=logits.device), NEG_INF
        )
    return logits.clone()


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


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


def position_ids(input_ids, pad_id):
    mask = (input_ids != pad_id).to(torch.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()
    T = input_ids.shape[-1]
    ar = torch.arange(T, 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
    after = is_eos.cumsum(dim=-1) - is_eos.to(torch.long)
    return (after == 0).to(torch.float32)


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


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


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


def truncate_left(input_ids, max_len):
    if input_ids.shape[-1] <= max_len:
        return input_ids
    return input_ids[..., -max_len:]


def pad_to_multiple(input_ids, pad_id, multiple):
    T = input_ids.shape[-1]
    new_T = ((T + multiple - 1) // multiple) * multiple
    if new_T == T:
        return input_ids
    pad = torch.full(
        (input_ids.shape[0], new_T - T),
        pad_id,
        dtype=input_ids.dtype,
        device=input_ids.device,
    )
    return torch.cat([input_ids, pad], dim=-1)


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


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


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


def discount_cumsum(x, gamma):
    out = torch.zeros_like(x)
    running = torch.zeros_like(x[..., 0])
    T = x.shape[-1]
    for t in range(T - 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, device=seq_lengths.device)
    pos = 0
    for l in seq_lengths.tolist():
        end = min(pos + l, total_len)
        n = end - pos
        if n > 0:
            out[pos:end] = torch.arange(n, dtype=torch.long, device=seq_lengths.device)
        pos = end
        if pos >= total_len:
            break
    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 logits_warper(logits, temperature, top_k, top_p).argmax(dim=-1)


def prepare_labels(input_ids, prompt_lengths, ignore_index):
    labels = shift_labels(input_ids, ignore_index)
    T = input_ids.shape[-1]
    ar = torch.arange(T, 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).to(torch.long) - 1
    b = torch.arange(logits.shape[0], device=logits.device)
    return logits[b, idx]


def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    return logits.masked_fill(probs < epsilon, NEG_INF)


def first_stop_index(input_ids, stop_ids):
    stop_ids = torch.as_tensor(stop_ids, device=input_ids.device)
    is_stop = (input_ids.unsqueeze(-1) == stop_ids).any(dim=-1)
    T = input_ids.shape[-1]
    ar = torch.arange(T, device=input_ids.device)
    idx = torch.where(is_stop, ar, torch.full_like(ar, T))
    first = idx.min(dim=-1)[0]
    return torch.where(first == T, torch.full_like(first, -1), first)


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


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


def right_pad_to(input_ids, pad_id, width):
    T = input_ids.shape[-1]
    if T >= width:
        return input_ids[..., :width]
    pad = torch.full(
        (input_ids.shape[0], width - T),
        pad_id,
        dtype=input_ids.dtype,
        device=input_ids.device,
    )
    return torch.cat([input_ids, pad], dim=-1)


def generation_stop_mask(input_ids, eos_id, min_len):
    is_eos = input_ids == eos_id
    after = is_eos.cumsum(dim=-1) - is_eos.to(torch.long)
    stop = after == 0
    T = input_ids.shape[-1]
    ar = torch.arange(T, device=input_ids.device)
    keep_min = ar.unsqueeze(0) < min_len
    return (stop | keep_min).to(torch.float32)

Comprehensive comparison harness

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

def eq(a,b,name):
    if isinstance(a,torch.Tensor):
        if a.shape!=b.shape:
            return f"SHAPE {a.shape} vs {b.shape}"
        if a.dtype!=b.dtype:
            return f"DTYPE {a.dtype} vs {b.dtype}"
        af=a.float(); bf=b.float()
        # handle -inf
        mask_inf = torch.isinf(af)|torch.isinf(bf)
        if not torch.equal(torch.isinf(af)&(af<0), torch.isinf(bf)&(bf<0)):
            return "INF mismatch"
        if not torch.equal(torch.isinf(af)&(af>0), torch.isinf(bf)&(bf>0)):
            return "PINF mismatch"
        fin=~mask_inf
        if fin.any():
            if not torch.allclose(af[fin],bf[fin],atol=1e-5,rtol=1e-4):
                d=(af[fin]-bf[fin]).abs().max().item()
                return f"VAL maxdiff {d}"
        return None
    else:
        return None if a==b else f"scalar {a} vs {b}"

fails={}
def test(name, fn_r, fn_o, *args):
    try: r=fn_r(*args)
    except Exception as e: r=("ERR",str(type(e).__name__))
    try: o=fn_o(*args)
    except Exception as e: o=("ERR",str(type(e).__name__))
    if isinstance(r,tuple) and r and r[0]=="ERR":
        if isinstance(o,tuple) and o and o[0]=="ERR": return
        fails.setdefault(name,[]).append(f"ref err {r}, ours {type(o)}"); return
    res=eq(r,o,name)
    if res: fails.setdefault(name,[]).append(res)

g=torch.Generator().manual_seed(0)
for trial in range(30):
    B=int(torch.randint(1,4,(1,),generator=g))
    V=int(torch.randint(2,10,(1,),generator=g))
    T=int(torch.randint(2,8,(1,),generator=g))
    logits=torch.randn(B,V,generator=g)
    ids=torch.randint(0,V,(B,T),generator=g)
    pad=0
    test("greedy_token",R.greedy_token,O.greedy_token,logits)
    for Tp in [0.0,0.5,1.0,2.0,0.7]:
        test("temperature_scale",R.temperature_scale,O.temperature_scale,logits,Tp)
    for k in [1,2,V,max(1,V-1)]:
        test("top_k_filter",R.top_k_filter,O.top_k_filter,logits,k)
        test("topk_indices",R.topk_indices,O.topk_indices,logits,k)
    for p in [0.0,0.3,0.5,0.9,1.0,0.75]:
        test("top_p_filter",R.top_p_filter,O.top_p_filter,logits,p)
    for p in [0.05,0.1,0.3,0.5,0.9]:
        test("min_p_filter",R.min_p_filter,O.min_p_filter,logits,p)
        test("epsilon_filter",R.epsilon_filter,O.epsilon_filter,logits,p)
    for pen in [0.5,1.0,1.2,1.5,2.0]:
        test("repetition_penalty",R.repetition_penalty,O.repetition_penalty,logits,ids,pen)
        test("frequency_penalty",R.frequency_penalty,O.frequency_penalty,logits,ids,pen)
        test("presence_penalty",R.presence_penalty,O.presence_penalty,logits,ids,pen)
    tk=torch.randint(0,V,(min(V,3),),generator=g).unique()
    test("suppress_tokens",R.suppress_tokens,O.suppress_tokens,logits,tk)
    test("suppress_tokens",R.suppress_tokens,O.suppress_tokens,logits,tk.tolist())
    for cl in [0,3,5]:
        test("min_length_filter",R.min_length_filter,O.min_length_filter,logits,cl,4,int(torch.randint(0,V,(1,))))
    test("renormalize_logits",R.renormalize_logits,O.renormalize_logits,logits)
    test("attention_mask",R.attention_mask,O.attention_mask,ids,pad)
    test("position_ids",R.position_ids,O.position_ids,ids,pad)
    test("causal_mask",R.causal_mask,O.causal_mask,T)
    test("shift_labels",R.shift_labels,O.shift_labels,ids,-100)
    pl=torch.randint(0,T+1,(B,),generator=g)
    test("completion_loss_labels",R.completion_loss_labels,O.completion_loss_labels,ids,pl,-100)
    test("prepare_labels",R.prepare_labels,O.prepare_labels,ids,pl,-100)
    test("mask_after_eos",R.mask_after_eos,O.mask_after_eos,ids,int(torch.randint(0,V,(1,))))
    test("sequence_lengths",R.sequence_lengths,O.sequence_lengths,ids,pad)
    am=(ids!=pad).float()
    # ensure at least one valid per row
    am[:,0]=1.0
    test("last_token_indices",R.last_token_indices,O.last_token_indices,am)
    logits3=torch.randn(B,T,V,generator=g)
    test("gather_last_logits",R.gather_last_logits,O.gather_last_logits,logits3,am)
    for ml in [1,3,T,T+2]:
        test("truncate_left",R.truncate_left,O.truncate_left,ids,ml)
    for m in [1,2,3,4]:
        test("pad_to_multiple",R.pad_to_multiple,O.pad_to_multiple,ids,pad,m)
    act=torch.randint(0,V,(B,T),generator=g)
    test("action_logprobs",R.action_logprobs,O.action_logprobs,logits3,act)
    test("entropy_from_logits",R.entropy_from_logits,O.entropy_from_logits,logits)
    rew=torch.randn(B,T,generator=g); mk=(torch.rand(B,T,generator=g)>0.3).float()
    test("reward_to_go",R.reward_to_go,O.reward_to_go,rew,mk)
    for gm in [0.0,0.9,0.99,1.0]:
        test("discount_cumsum",R.discount_cumsum,O.discount_cumsum,rew,gm)
    sl=torch.randint(1,4,(B,),generator=g); tot=int(sl.sum())+int(torch.randint(0,3,(1,)))
    test("build_position_ids_packed",R.build_position_ids_packed,O.build_position_ids_packed,sl,tot)
    for Tp in [0.5,1.0,2.0]:
      for k in [1,2,V]:
        for p in [0.5,0.9,1.0]:
            test("logits_warper",R.logits_warper,O.logits_warper,logits,Tp,k,p)
            test("sample_greedy_warped",R.sample_greedy_warped,O.sample_greedy_warped,logits,Tp,k,p)
    test("rollout_last_logits",R.rollout_last_logits,O.rollout_last_logits,logits3,ids_nz:=torch.where((ids!=pad).any(1,keepdim=True),ids,ids+1),pad)
    sid=torch.randint(0,V,(2,),generator=g)
    test("first_stop_index",R.first_stop_index,O.first_stop_index,ids,sid)
    mmask=(torch.rand(B,V,generator=g)>0.3).float(); mmask[:,0]=1.0
    test("masked_argmax",R.masked_argmax,O.masked_argmax,logits,mmask)
    for w in [1,T,T+2]:
        test("right_pad_to",R.right_pad_to,O.right_pad_to,ids,pad,w)
    for mn in [0,2,T]:
        test("generation_stop_mask",R.generation_stop_mask,O.generation_stop_mask,ids,int(torch.randint(0,V,(1,))),mn)

if not fails:
    print("ALL PASS")
else:
    for k,v in fails.items():
        print("FAIL",k,v[:3])
EOF
PYTHONPATH=/app python3 /tmp/cmp.py
output
FAIL epsilon_filter ['INF mismatch', 'INF mismatch', 'INF mismatch']
FAIL last_token_indices ['VAL maxdiff 1.0', 'VAL maxdiff 1.0', 'VAL maxdiff 1.0']
FAIL gather_last_logits ['VAL maxdiff 2.2464613914489746', 'VAL maxdiff 2.2746009826660156', 'VAL maxdiff 2.6070845127105713']
FAIL rollout_last_logits ['VAL maxdiff 2.2464613914489746', 'VAL maxdiff 2.2746009826660156', 'VAL maxdiff 1.367235541343689']

[stdout]
FAIL epsilon_filter ['INF mismatch', 'INF mismatch', 'INF mismatch']
FAIL last_token_indices ['VAL maxdiff 1.0', 'VAL maxdiff 1.0', 'VAL maxdiff 1.0']
FAIL gather_last_logits ['VAL maxdiff 2.2464613914489746', 'VAL maxdiff 2.2746009826660156', 'VAL maxdiff 2.6070845127105713']
FAIL rollout_last_logits ['VAL maxdiff 2.2464613914489746', 'VAL maxdiff 2.2746009826660156', 'VAL maxdiff 1.367235541343689']

Probe last-index and epsilon semantics

shell
cat > /tmp/pe.py << 'EOF'
import torch, reforacle as R
am=torch.tensor([[1.,0.,1.,0.],[1.,1.,0.,1.]])
print("last_token_indices gappy", R.last_token_indices(am).tolist())
ids=torch.tensor([[5,0,7,0],[1,2,0,8]])
print("rollout last, ids gappy pad0", R.rollout_last_logits(torch.arange(2*4*3.).reshape(2,4,3), ids, 0).tolist())
print("gather_last gappy am", R.gather_last_logits(torch.arange(2*4*3.).reshape(2,4,3), am).tolist())
# epsilon exact
g=torch.Generator().manual_seed(5)
lg=torch.randn(1,6,generator=g)
p=torch.softmax(lg,-1)
print("probs", p.tolist())
for eps in [0.1,0.15,0.2]:
    print("eps",eps, R.epsilon_filter(lg,eps).tolist())
EOF
PYTHONPATH=/app python3 /tmp/pe.py
output
last_token_indices gappy [2, 3]
rollout last, ids gappy pad0 [[6.0, 7.0, 8.0], [21.0, 22.0, 23.0]]
gather_last gappy am [[6.0, 7.0, 8.0], [21.0, 22.0, 23.0]]
probs [[0.05352732166647911, 0.047615163028240204, 0.04984293133020401, 0.16977892816066742, 0.07149002701044083, 0.6077456474304199]]
eps 0.1 [[-inf, -inf, -inf, 0.6675242781639099, -inf, 1.9427834749221802]]
eps 0.15 [[-inf, -inf, -inf, 0.6675242781639099, -inf, 1.9427834749221802]]
eps 0.2 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802]]

[stdout]
last_token_indices gappy [2, 3]
rollout last, ids gappy pad0 [[6.0, 7.0, 8.0], [21.0, 22.0, 23.0]]
gather_last gappy am [[6.0, 7.0, 8.0], [21.0, 22.0, 23.0]]
probs [[0.05352732166647911, 0.047615163028240204, 0.04984293133020401, 0.16977892816066742, 0.07149002701044083, 0.6077456474304199]]
eps 0.1 [[-inf, -inf, -inf, 0.6675242781639099, -inf, 1.9427834749221802]]
eps 0.15 [[-inf, -inf, -inf, 0.6675242781639099, -inf, 1.9427834749221802]]
eps 0.2 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802]]

Check large-threshold keep-at-least-one behavior

shell
cat > /tmp/pe2.py << 'EOF'
import torch, reforacle as R
g=torch.Generator().manual_seed(5)
lg=torch.randn(2,6,generator=g)
for eps in [0.3,0.5,0.9,0.99]:
    print("eps",eps, R.epsilon_filter(lg,eps).tolist())
# top_k min_tokens? top_p keep at least 1 confirmed. min_p large:
for p in [0.9,0.99,1.0]:
    print("min_p",p, R.min_p_filter(lg,p).tolist())
EOF
PYTHONPATH=/app python3 /tmp/pe2.py
output
eps 0.3 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, 0.6312134265899658, -inf, -inf, 0.6907187700271606]]
eps 0.5 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
eps 0.9 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
eps 0.99 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
min_p 0.9 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, 0.6312134265899658, -inf, -inf, 0.6907187700271606]]
min_p 0.99 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
min_p 1.0 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]

[stdout]
eps 0.3 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, 0.6312134265899658, -inf, -inf, 0.6907187700271606]]
eps 0.5 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
eps 0.9 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
eps 0.99 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
min_p 0.9 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, 0.6312134265899658, -inf, -inf, 0.6907187700271606]]
min_p 0.99 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
min_p 1.0 [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]

/app/objectives.py

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


def gather_last_logits(logits, attn_mask):
    idx = attn_mask.sum(dim=-1).to(torch.long) - 1
    b = torch.arange(logits.shape[0], device=logits.device)
    return logits[b, idx]
+ added
def _last_true_index(mask):
    T = mask.shape[-1]
    ar = torch.arange(T, device=mask.device)
    idx = torch.where(mask != 0, ar, torch.full_like(ar, -1))
    return idx.max(dim=-1)[0]


def last_token_indices(attn_mask):
    return _last_true_index(attn_mask)


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

/app/objectives.py

− removed
def rollout_last_logits(logits, input_ids, pad_id):
    idx = (input_ids != pad_id).sum(dim=-1).to(torch.long) - 1
    b = torch.arange(logits.shape[0], device=logits.device)
    return logits[b, idx]


def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    return logits.masked_fill(probs < epsilon, NEG_INF)
+ added
def rollout_last_logits(logits, input_ids, pad_id):
    idx = _last_true_index(input_ids != pad_id)
    b = torch.arange(logits.shape[0], device=logits.device)
    return logits[b, idx]


def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    remove = probs < epsilon
    top_val = logits.max(dim=-1, keepdim=True)[0]
    remove = remove & (logits < top_val)
    return logits.masked_fill(remove, NEG_INF)

Rerun full comparison

shell
PYTHONPATH=/app python3 /tmp/cmp.py
output
ALL PASS

[stdout]
ALL PASS

Extended stress test

shell
cd /tmp && sed -i 's/manual_seed(0)/manual_seed(123)/; s/range(30)/range(120)/; s/torch.randint(2,10/torch.randint(2,20/; s/torch.randint(2,8/torch.randint(2,14/' cmp.py
# add warper T=0 test
python3 - << 'EOF'
import re
s=open('/tmp/cmp.py').read()
s=s.replace("for Tp in [0.5,1.0,2.0]:\n      for k in [1,2,V]:","for Tp in [0.0,0.5,1.0,2.0]:\n      for k in [1,2,V]:")
open('/tmp/cmp.py','w').write(s)
EOF
PYTHONPATH=/app python3 /tmp/cmp.py
output
ALL PASS
Shell cwd was reset to /app

[stdout]
ALL PASS
[stderr]

Shell cwd was reset to /app

Edge-case conventions

shell
cat > /tmp/pedge.py << 'EOF'
import torch, reforacle as R
import importlib.util
spec=importlib.util.spec_from_file_location("obj","/app/objectives.py")
O=importlib.util.module_from_spec(spec); spec.loader.exec_module(O)
def cmp(name,a,b):
    try:
        ok = a.shape==b.shape and a.dtype==b.dtype and torch.allclose(a.float(),b.float(),atol=1e-5,rtol=1e-4,equal_nan=True)
    except Exception as e: ok=f"ERR {e}"
    print(name, "OK" if ok is True else ("FAIL "+str(ok)), "ref",a.dtype,tuple(a.shape))
# 1D RL
x=torch.randn(6)
cmp("discount1d",R.discount_cumsum(x,0.9),O.discount_cumsum(x,0.9))
r=torch.randn(6); m=(torch.rand(6)>0.3).float()
cmp("rtg1d",R.reward_to_go(r,m),O.reward_to_go(r,m))
# packed truncation sum>total
sl=torch.tensor([3,4]); 
cmp("packed_trunc",R.build_position_ids_packed(sl,5),O.build_position_ids_packed(sl,5))
cmp("packed_exact",R.build_position_ids_packed(sl,7),O.build_position_ids_packed(sl,7))
# entropy alt formula check large
lg=torch.randn(4,50); cmp("entropy_big",R.entropy_from_logits(lg),O.entropy_from_logits(lg))
# temp T=0 tie
tie=torch.tensor([[1.0,1.0,0.5]])
cmp("temp0tie",R.temperature_scale(tie,0.0),O.temperature_scale(tie,0.0))
cmp("greedytie",R.greedy_token(tie),O.greedy_token(tie))
# suppress empty
cmp("suppress_empty",R.suppress_tokens(lg,torch.tensor([],dtype=torch.long)),O.suppress_tokens(lg,torch.tensor([],dtype=torch.long)))
# min_length negative eos edge already fine
# action_logprobs 1D? actions shape (B,T), logits (B,T,V) tested. try (B,V)/(B,) 
l2=torch.randn(3,7); a2=torch.randint(0,7,(3,))
cmp("action2d",R.action_logprobs(l2,a2),O.action_logprobs(l2,a2))
EOF
PYTHONPATH=/app python3 /tmp/pedge.py
output
discount1d OK ref torch.float32 (6,)
rtg1d OK ref torch.float32 (6,)
packed_trunc OK ref torch.int64 (5,)
packed_exact OK ref torch.int64 (7,)
entropy_big OK ref torch.float32 (4,)
temp0tie FAIL False ref torch.float32 (1, 3)
greedytie FAIL False ref torch.int64 (1,)
suppress_empty OK ref torch.float32 (4, 50)
action2d OK ref torch.float32 (3,)

[stdout]
discount1d OK ref torch.float32 (6,)
rtg1d OK ref torch.float32 (6,)
packed_trunc OK ref torch.int64 (5,)
packed_exact OK ref torch.int64 (7,)
entropy_big OK ref torch.float32 (4,)
temp0tie FAIL False ref torch.float32 (1, 3)
greedytie FAIL False ref torch.int64 (1,)
suppress_empty OK ref torch.float32 (4, 50)
action2d OK ref torch.float32 (3,)

Check tie-breaking

shell
cat > /tmp/ptie.py << 'EOF'
import torch, reforacle as R
tie=torch.tensor([[1.0,1.0,0.5],[0.5,2.0,2.0]])
print("greedy",R.greedy_token(tie).tolist())
print("temp0",R.temperature_scale(tie,0.0).tolist())
print("masked_argmax", R.masked_argmax(tie, torch.ones_like(tie)).tolist())
print("torch.argmax", tie.argmax(-1).tolist())
EOF
PYTHONPATH=/app python3 /tmp/ptie.py
output
greedy [1, 2]
temp0 [[-inf, 0.0, -inf], [-inf, -inf, 0.0]]
masked_argmax [1, 2]
torch.argmax [0, 1]

[stdout]
greedy [1, 2]
temp0 [[-inf, 0.0, -inf], [-inf, -inf, 0.0]]
masked_argmax [1, 2]
torch.argmax [0, 1]

topk and first_stop tie behavior

shell
cat > /tmp/ptie2.py << 'EOF'
import torch, reforacle as R
tie=torch.tensor([[3.0,1.0,3.0,2.0,3.0]])
print("topk_indices k=2", R.topk_indices(tie,2).tolist())
print("torch.topk", tie.topk(2,-1)[1].tolist())
print("top_k_filter k=2", R.top_k_filter(tie,2).tolist())
print("first_stop tie", R.first_stop_index(torch.tensor([[3,3,1,3]]), torch.tensor([3])).tolist())
EOF
PYTHONPATH=/app python3 /tmp/ptie2.py
output
topk_indices k=2 [[0, 2]]
torch.topk [[2, 4]]
top_k_filter k=2 [[3.0, -inf, 3.0, -inf, 3.0]]
first_stop tie [0]

[stdout]
topk_indices k=2 [[0, 2]]
torch.topk [[2, 4]]
top_k_filter k=2 [[3.0, -inf, 3.0, -inf, 3.0]]
first_stop tie [0]

Confirm argmax-last convention

shell
cat > /tmp/ptie3.py << 'EOF'
import torch, reforacle as R
print("greedy [3,1,3]", R.greedy_token(torch.tensor([[3.,1.,3.]])).tolist())
print("greedy [3,3,1,3]", R.greedy_token(torch.tensor([[3.,3.,1.,3.]])).tolist())
print("masked_argmax [3,1,3] mask[1,1,0]", R.masked_argmax(torch.tensor([[3.,1.,3.]]),torch.tensor([[1.,1.,0.]])).tolist())
# sample_greedy_warped tie: temp huge to flatten? use k that keeps ties
print("sgw tie", R.sample_greedy_warped(torch.tensor([[3.,1.,3.,3.]]),1.0,4,1.0).tolist())
EOF
PYTHONPATH=/app python3 /tmp/ptie3.py
output
greedy [3,1,3] [2]
greedy [3,3,1,3] [3]
masked_argmax [3,1,3] mask[1,1,0] [0]
sgw tie [3]

[stdout]
greedy [3,1,3] [2]
greedy [3,3,1,3] [3]
masked_argmax [3,1,3] mask[1,1,0] [0]
sgw tie [3]

/app/objectives.py

− removed
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
+ added
def _argmax_last(x):
    n = x.shape[-1]
    return n - 1 - torch.flip(x, [-1]).argmax(dim=-1)


def greedy_token(logits):
    return _argmax_last(logits)


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

/app/objectives.py

− removed
def sample_greedy_warped(logits, temperature, top_k, top_p):
    return logits_warper(logits, temperature, top_k, top_p).argmax(dim=-1)
+ added
def sample_greedy_warped(logits, temperature, top_k, top_p):
    return _argmax_last(logits_warper(logits, temperature, top_k, top_p))

/app/objectives.py

− removed
def masked_argmax(logits, mask):
    masked = logits.masked_fill(mask == 0, NEG_INF)
    return masked.argmax(dim=-1)


def topk_indices(logits, k):
    return torch.topk(logits, k, dim=-1)[1]
+ added
def masked_argmax(logits, mask):
    masked = logits.masked_fill(mask == 0, NEG_INF)
    return _argmax_last(masked)


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

Rerun harness plus tie probes

shell
PYTHONPATH=/app python3 /tmp/cmp.py && PYTHONPATH=/app python3 /tmp/ptie.py && PYTHONPATH=/app python3 /tmp/ptie2.py && PYTHONPATH=/app python3 /tmp/ptie3.py 2>&1 | tail -20
output
ALL PASS
greedy [1, 2]
temp0 [[-inf, 0.0, -inf], [-inf, -inf, 0.0]]
masked_argmax [1, 2]
torch.argmax [0, 1]
topk_indices k=2 [[0, 2]]
torch.topk [[2, 4]]
top_k_filter k=2 [[3.0, -inf, 3.0, -inf, 3.0]]
first_stop tie [0]
greedy [3,1,3] [2]
greedy [3,3,1,3] [3]
masked_argmax [3,1,3] mask[1,1,0] [0]
sgw tie [3]

[stdout]
ALL PASS
greedy [1, 2]
temp0 [[-inf, 0.0, -inf], [-inf, -inf, 0.0]]
masked_argmax [1, 2]
torch.argmax [0, 1]
topk_indices k=2 [[0, 2]]
torch.topk [[2, 4]]
top_k_filter k=2 [[3.0, -inf, 3.0, -inf, 3.0]]
first_stop tie [0]
greedy [3,1,3] [2]
greedy [3,3,1,3] [3]
masked_argmax [3,1,3] mask[1,1,0] [0]
sgw tie [3]

Direct tie-case comparison

shell
cat > /tmp/ptie_cmp.py << 'EOF'
import torch, reforacle as R, importlib.util
spec=importlib.util.spec_from_file_location("obj","/app/objectives.py"); O=importlib.util.module_from_spec(spec); spec.loader.exec_module(O)
cases=[
 ("greedy", lambda M: M.greedy_token(torch.tensor([[1.,1.,0.5],[0.5,2,2],[3,1,3]]))),
 ("temp0", lambda M: M.temperature_scale(torch.tensor([[1.,1.,0.5],[0.5,2,2]]),0.0)),
 ("masked_argmax", lambda M: M.masked_argmax(torch.tensor([[3.,1.,3.],[3,3,1]]),torch.tensor([[1.,1.,0.],[1,1,1]]))),
 ("topk_indices", lambda M: M.topk_indices(torch.tensor([[3.,1.,3.,2.,3.]]),3)),
 ("sgw", lambda M: M.sample_greedy_warped(torch.tensor([[3.,1.,3.,3.]]),1.0,4,1.0)),
]
for n,f in cases:
    a=f(R); b=f(O); print(n, "OK" if torch.equal(a,b) else f"FAIL ref={a.tolist()} ours={b.tolist()}")
EOF
PYTHONPATH=/app python3 /tmp/ptie_cmp.py
output
greedy OK
temp0 OK
masked_argmax OK
topk_indices OK
sgw OK

[stdout]
greedy OK
temp0 OK
masked_argmax OK
topk_indices OK
sgw OK

Fresh-seed stress + standalone import check

shell
cd /tmp && sed -i 's/manual_seed(123)/manual_seed(7777)/' cmp.py && PYTHONPATH=/app python3 /tmp/cmp.py
echo "--- import check (no reforacle) ---"
grep -n "import" /app/objectives.py
echo "--- standalone run without reforacle on path ---"
cd / && python3 -c "
import importlib.util
spec=importlib.util.spec_from_file_location('obj','/app/objectives.py')
O=importlib.util.module_from_spec(spec); spec.loader.exec_module(O)
import torch
print('loads standalone; greedy:', O.greedy_token(torch.randn(2,5)).shape)
print('reforacle importable here?', __import__('importlib').util.find_spec('reforacle'))
"
output
ALL PASS
--- import check (no reforacle) ---
1:from __future__ import annotations
3:import torch
--- standalone run without reforacle on path ---
loads standalone; greedy: torch.Size([2])
reforacle importable here? None
Shell cwd was reset to /app

[stdout]
ALL PASS
--- import check (no reforacle) ---
1:from __future__ import annotations
3:import torch
--- standalone run without reforacle on path ---
loads standalone; greedy: torch.Size([2])
reforacle importable here? None
[stderr]

Shell cwd was reset to /app
All 37 functions match the reference exactly across extensive randomized testing, and the file is fully standalone (only `torch`, loads and runs without `reforacle` on the path). ## Summary I reverse-engineered `reforacle` by probing it and reimplemented all 37 functions in `/app/objectives.py`. Key conventions I recovered (none of which were guessable from the names): - **`temperature_scale(T=0)`** → one-hot in log-space (`0.0` at argmax, `-inf` elsewhere). - **Argmax tie-breaking is toward the *last* index** (not torch's default first) , affects `greedy_token`, `temperature_scale(T=0)`, `masked_argmax`, `sample_greedy_warped`. - **`topk_indices` breaks ties toward the *lowest* index** (stable descending sort, not `torch.topk`). - **Penalties are count-based**: `repetition` = `penalty**count` (divide if logit≥0, multiply if <0); `frequency` = `logits - penalty*count`; `presence` = `logits - penalty*1{present}`. - **`epsilon_filter`** keeps at least the top-1 token even when the threshold exceeds all probs. - **Nucleus/last-token semantics**: `top_p` keeps the first token above threshold; `last_token_indices`/`gather_last_logits`/`rollout_last_logits` use the *last nonzero position* (not the mask count) , matters for gappy masks. - **`completion_loss_labels`** masks the prompt without shifting; **`prepare_labels`** shifts then masks `i < prompt_len-1`. - **`generation_stop_mask`** keeps positions `≤ first_eos` OR `< min_len`. - **`build_position_ids_packed`** zero-pads leftover positions beyond the packed segments. - Dtypes matched exactly (float32 masks, int64 indices). Verified against `reforacle` on 100+ randomized trials across varied batch/vocab/sequence shapes plus targeted tie and edge cases, all within the required tolerances.
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_5d8a32a83dbb450a. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.

Trial trial_5d8a32a83dbb450a · verifier authoritative; classifier explanatory.