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.
SubtypeComplexity Overwhelm
EvidenceTest output: 'FUNCTIONS PASSED: 31/37 (partial_score=0.8378) REWARD = 0 (1 iff ALL 37 functions match the sealed reference exactly)'. Agent systematically probed all 37 functions via oracle (p1.py through p10.py), implemented all of them, and achieved 84% pass rate. Agent discovered and fixed masked_argmax bug during verification. Only 6 functions out of 37 failed numerical matching tests with tolerance atol=1e-5, rtol=1e-4 for primitives.
Root causeThe task requires 100% pass rate across 37 complex decoding/collation/RLHF functions with exact numerical agreement. Agent achieved 84% but failed on 6 functions, likely due to subtle edge cases in the reference behavior that required more probing iterations or different input patterns to discover.
RecommendationN/A - task is fine
Trajectory
Tool-by-tool agent trajectory
26 tool calls · 4 tool types · 28 steps
Reimplement the post-training decoding-and-collation library in `/app/objectives.py` so that every function reproduces the sealed reference module `reforacle` exactly. `reforacle` is importable and callable (`import reforacle`) but its source cannot be read. It is the only ground truth; there is no written specification of what any function computes or of the conventions it follows. Recover its behavior by calling it. `/app/objectives.py` contains thirty-seven function stubs , logit processors and sampling filters, sequence collation and masking utilities, and rollout helpers used to generate and batch data for RLHF and SFT. Keep every signature exactly as given and implement every function. The names indicate the standard quantity each is meant to produce, but the reference's exact definition is whatever `reforacle` returns , not what you assume. Determine each definition by probing. Grading compares your implementation to `reforacle` on held-out inputs you do not see, function by function, requiring exact numerical agreement (primitives within `atol=1e-6, rtol=1e-5`; composite functions within `atol=1e-5, rtol=1e-4`). Each function passes only if it agrees on every graded input, and the task passes only if all thirty-seven functions pass. Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on. Grading runs your `/app/objectives.py` in a clean sandbox that does not contain `reforacle`, so any submission that imports, wraps, calls, loads, or otherwise depends on the reference at run time fails outright. Your file must be a standalone implementation whose only dependency is `torch`. Do not read, decompile, disassemble, or otherwise extract the compiled reference; recover its behavior only by calling it and comparing. `python3 /app/probe.py` shows the calling pattern; `bash /app/timer.sh` reports remaining budget.

/app/objectives.py

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

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

Remaining time (hours:minutes):
1:59

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

/app/probe.py

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

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

Remaining time (hours:minutes):
1:59

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

Check remaining budget

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

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

Remaining time (hours:minutes):
1:59

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

Run probe.py

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

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

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

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

Probe batch 1

shell
cat > /tmp/p1.py << 'EOF'
import torch, reforacle as R

def show(name, x):
    if torch.is_tensor(x):
        print(f"{name}: dtype={x.dtype} shape={tuple(x.shape)}")
        print("   ", x.tolist())
    else:
        print(f"{name}: {x!r}")

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

show("greedy_token", R.greedy_token(logits))
show("temp T=2.0", R.temperature_scale(logits, 2.0))
show("temp T=0.5", R.temperature_scale(logits, 0.5))
show("top_k k=2", R.top_k_filter(logits, 2))
show("top_k k=1", R.top_k_filter(logits, 1))
show("top_p p=0.9", R.top_p_filter(logits, 0.9))
show("top_p p=0.5", R.top_p_filter(logits, 0.5))
show("min_p p=0.3", R.min_p_filter(logits, 0.3))
show("min_p p=0.1", R.min_p_filter(logits, 0.1))
show("renormalize", R.renormalize_logits(logits))
show("entropy", R.entropy_from_logits(logits))
show("epsilon 0.05", R.epsilon_filter(logits, 0.05))
show("suppress [0,3]", R.suppress_tokens(logits, [0,3]))
show("topk_indices k=3", R.topk_indices(logits, 3))
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

Run probe from /app

shell
cp /tmp/p1.py /tmp/p1.py && python3 /tmp/p1.py
output
Exit code 1
cp: '/tmp/p1.py' and '/tmp/p1.py' are the same file

[error] tool reported failure

Run 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

List app and run from app

shell
ls /app; echo "---"; python3 /tmp/p1.py
output
Exit code 1
__pycache__
objectives.py
probe.py
reforacle.cpython-312-aarch64-linux-gnu.so
timer.sh
---
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

Run from app with cd

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

Add app to path and run

shell
sed -i '1i import sys; sys.path.insert(0, "/app")' /tmp/p1.py && 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: dtype=torch.int64 shape=(2,)
    [0, 1]
temp T=2.0: dtype=torch.float32 shape=(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=0.5: dtype=torch.float32 shape=(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=2: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k k=1: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_p p=0.9: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p p=0.5: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
min_p p=0.3: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
min_p p=0.1: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
renormalize: dtype=torch.float32 shape=(2, 6)
    [[-0.5231569409370422, -2.357581853866577, -4.2429423332214355, -1.4957218170166016, -3.1486752033233643, -3.4627482891082764], [-1.5008691549301147, -1.0661896467208862, -2.623473644256592, -2.3075594902038574, -2.5008511543273926, -1.7221795320510864]]
entropy: dtype=torch.float32 shape=(2,)
    [1.1729530096054077, 1.6344718933105469]
epsilon 0.05: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
suppress [0,3]: dtype=torch.float32 shape=(2, 6)
    [[-inf, -0.293428897857666, -2.1787893772125244, -inf, -1.0845223665237427, -1.3985954523086548], [-inf, 0.8380263447761536, -0.7192575931549072, -inf, -0.5966353416442871, 0.18203648924827576]]
topk_indices k=3: dtype=torch.int64 shape=(2, 3)
    [[0, 3, 1], [1, 0, 5]]

[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: dtype=torch.int64 shape=(2,)
    [0, 1]
temp T=2.0: dtype=torch.float32 shape=(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=0.5: dtype=torch.float32 shape=(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=2: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k k=1: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_p p=0.9: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p p=0.5: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
min_p p=0.3: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
min_p p=0.1: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
renormalize: dtype=torch.float32 shape=(2, 6)
    [[-0.5231569409370422, -2.357581853866577, -4.2429423332214355, -1.4957218170166016, -3.1486752033233643, -3.4627482891082764], [-1.5008691549301147, -1.0661896467208862, -2.623473644256592, -2.3075594902038574, -2.5008511543273926, -1.7221795320510864]]
entropy: dtype=torch.float32 shape=(2,)
    [1.1729530096054077, 1.6344718933105469]
epsilon 0.05: dtype=torch.float32 shape=(2, 6)
    [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
suppress [0,3]: dtype=torch.float32 shape=(2, 6)
    [[-inf, -0.293428897857666, -2.1787893772125244, -inf, -1.0845223665237427, -1.3985954523086548], [-inf, 0.8380263447761536, -0.7192575931549072, -inf, -0.5966353416442871, 0.18203648924827576]]
topk_indices k=3: dtype=torch.int64 shape=(2, 3)
    [[0, 3, 1], [1, 0, 5]]

Probe penalties and min_length

shell
cat > /tmp/p2.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R
def show(name, x):
    if torch.is_tensor(x): print(f"{name}: dtype={x.dtype} shape={tuple(x.shape)}\n    {x.tolist()}")
    else: print(f"{name}: {x!r}")

g=torch.Generator().manual_seed(1)
logits=torch.randn(2,6,generator=g)
print("LOGITS",logits.tolist())
ids=torch.tensor([[0,1,1,3,1],[2,2,4,4,4]])
show("rep_pen 1.5", R.repetition_penalty(logits.clone(), ids, 1.5))
show("freq_pen 0.5", R.frequency_penalty(logits.clone(), ids, 0.5))
show("pres_pen 0.5", R.presence_penalty(logits.clone(), ids, 0.5))
# min_length
show("min_len cur=2 min=5 eos=3", R.min_length_filter(logits.clone(), 2, 5, 3))
show("min_len cur=5 min=5 eos=3", R.min_length_filter(logits.clone(), 5, 5, 3))
show("min_len cur=6 min=5 eos=3", R.min_length_filter(logits.clone(), 6, 5, 3))

# negative logits for rep pen sign check
neg=torch.tensor([[-1.0,2.0,-3.0,4.0]])
show("rep_pen neg 2.0", R.repetition_penalty(neg.clone(), torch.tensor([[0,2]]), 2.0))
EOF
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: dtype=torch.float32 shape=(2, 6)
    [[0.44090142846107483, 0.07908862829208374, 0.06167725846171379, 0.41421154141426086, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -2.312119483947754, -0.563052773475647, -3.0114805698394775, -0.05825017765164375]]
freq_pen 0.5: dtype=torch.float32 shape=(2, 6)
    [[0.16135215759277344, -1.2330758571624756, 0.06167725846171379, 0.12131732702255249, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -2.027608633041382, -0.563052773475647, -2.3922905921936035, -0.05825017765164375]]
pres_pen 0.5: dtype=torch.float32 shape=(2, 6)
    [[0.16135215759277344, -0.23307588696479797, 0.06167725846171379, 0.12131732702255249, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.5276086330413818, -0.563052773475647, -1.3922905921936035, -0.05825017765164375]]
min_len cur=2 min=5 eos=3: dtype=torch.float32 shape=(2, 6)
    [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, -inf, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -inf, -0.8922905325889587, -0.05825017765164375]]
min_len cur=5 min=5 eos=3: dtype=torch.float32 shape=(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]]
min_len cur=6 min=5 eos=3: dtype=torch.float32 shape=(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]]
rep_pen neg 2.0: dtype=torch.float32 shape=(1, 4)
    [[-2.0, 2.0, -6.0, 4.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: dtype=torch.float32 shape=(2, 6)
    [[0.44090142846107483, 0.07908862829208374, 0.06167725846171379, 0.41421154141426086, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -2.312119483947754, -0.563052773475647, -3.0114805698394775, -0.05825017765164375]]
freq_pen 0.5: dtype=torch.float32 shape=(2, 6)
    [[0.16135215759277344, -1.2330758571624756, 0.06167725846171379, 0.12131732702255249, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -2.027608633041382, -0.563052773475647, -2.3922905921936035, -0.05825017765164375]]
pres_pen 0.5: dtype=torch.float32 shape=(2, 6)
    [[0.16135215759277344, -0.23307588696479797, 0.06167725846171379, 0.12131732702255249, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.5276086330413818, -0.563052773475647, -1.3922905921936035, -0.05825017765164375]]
min_len cur=2 min=5 eos=3: dtype=torch.float32 shape=(2, 6)
    [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, -inf, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -inf, -0.8922905325889587, -0.05825017765164375]]
min_len cur=5 min=5 eos=3: dtype=torch.float32 shape=(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]]
min_len cur=6 min=5 eos=3: dtype=torch.float32 shape=(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]]
rep_pen neg 2.0: dtype=torch.float32 shape=(1, 4)
    [[-2.0, 2.0, -6.0, 4.0]]

Probe collation functions

shell
cat > /tmp/p3.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R
def show(name, x):
    if torch.is_tensor(x): print(f"{name}: dtype={x.dtype} shape={tuple(x.shape)}\n    {x.tolist()}")
    else: print(f"{name}: {x!r}")

# padding conventions: assume right-padded with pad_id=0
ids = torch.tensor([[5,6,7,0,0],[8,9,0,0,0],[1,2,3,4,5]])
show("attention_mask pad0", R.attention_mask(ids, 0))
show("position_ids pad0", R.position_ids(ids, 0))
show("sequence_lengths pad0", R.sequence_lengths(ids, 0))
# left padded?
idsL = torch.tensor([[0,0,5,6,7],[0,0,0,8,9]])
show("attention_mask leftpad", R.attention_mask(idsL, 0))
show("position_ids leftpad", R.position_ids(idsL, 0))
show("sequence_lengths leftpad", R.sequence_lengths(idsL, 0))
show("causal_mask 4", R.causal_mask(4))
show("shift_labels ignore=-100", R.shift_labels(ids, -100))
show("mask_after_eos eos=7", R.mask_after_eos(ids, 7))
show("last_token_indices", R.last_token_indices(R.attention_mask(ids,0)))
lg = torch.arange(3*5*4).float().reshape(3,5,4)
show("gather_last_logits", R.gather_last_logits(lg, R.attention_mask(ids,0)))
show("truncate_left max=3", R.truncate_left(ids, 3))
show("pad_to_multiple pad0 mult4", R.pad_to_multiple(ids, 0, 4))
show("right_pad_to pad0 w7", R.right_pad_to(ids, 0, 7))
EOF
python3 /tmp/p3.py
output
attention_mask pad0: dtype=torch.float32 shape=(3, 5)
    [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
position_ids pad0: dtype=torch.int64 shape=(3, 5)
    [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1], [0, 1, 2, 3, 4]]
sequence_lengths pad0: dtype=torch.int64 shape=(3,)
    [3, 2, 5]
attention_mask leftpad: dtype=torch.float32 shape=(2, 5)
    [[0.0, 0.0, 1.0, 1.0, 1.0], [0.0, 0.0, 0.0, 1.0, 1.0]]
position_ids leftpad: dtype=torch.int64 shape=(2, 5)
    [[0, 0, 0, 1, 2], [0, 0, 0, 0, 1]]
sequence_lengths leftpad: dtype=torch.int64 shape=(2,)
    [3, 2]
causal_mask 4: dtype=torch.float32 shape=(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 ignore=-100: dtype=torch.int64 shape=(3, 5)
    [[6, 7, 0, 0, -100], [9, 0, 0, 0, -100], [2, 3, 4, 5, -100]]
mask_after_eos eos=7: dtype=torch.float32 shape=(3, 5)
    [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
last_token_indices: dtype=torch.int64 shape=(3,)
    [2, 1, 4]
gather_last_logits: dtype=torch.float32 shape=(3, 4)
    [[8.0, 9.0, 10.0, 11.0], [24.0, 25.0, 26.0, 27.0], [56.0, 57.0, 58.0, 59.0]]
truncate_left max=3: dtype=torch.int64 shape=(3, 3)
    [[7, 0, 0], [0, 0, 0], [3, 4, 5]]
pad_to_multiple pad0 mult4: dtype=torch.int64 shape=(3, 8)
    [[5, 6, 7, 0, 0, 0, 0, 0], [8, 9, 0, 0, 0, 0, 0, 0], [1, 2, 3, 4, 5, 0, 0, 0]]
right_pad_to pad0 w7: dtype=torch.int64 shape=(3, 7)
    [[5, 6, 7, 0, 0, 0, 0], [8, 9, 0, 0, 0, 0, 0], [1, 2, 3, 4, 5, 0, 0]]

[stdout]
attention_mask pad0: dtype=torch.float32 shape=(3, 5)
    [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
position_ids pad0: dtype=torch.int64 shape=(3, 5)
    [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1], [0, 1, 2, 3, 4]]
sequence_lengths pad0: dtype=torch.int64 shape=(3,)
    [3, 2, 5]
attention_mask leftpad: dtype=torch.float32 shape=(2, 5)
    [[0.0, 0.0, 1.0, 1.0, 1.0], [0.0, 0.0, 0.0, 1.0, 1.0]]
position_ids leftpad: dtype=torch.int64 shape=(2, 5)
    [[0, 0, 0, 1, 2], [0, 0, 0, 0, 1]]
sequence_lengths leftpad: dtype=torch.int64 shape=(2,)
    [3, 2]
causal_mask 4: dtype=torch.float32 shape=(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 ignore=-100: dtype=torch.int64 shape=(3, 5)
    [[6, 7, 0, 0, -100], [9, 0, 0, 0, -100], [2, 3, 4, 5, -100]]
mask_after_eos eos=7: dtype=torch.float32 shape=(3, 5)
    [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
last_token_indices: dtype=torch.int64 shape=(3,)
    [2, 1, 4]
gather_last_logits: dtype=torch.float32 shape=(3, 4)
    [[8.0, 9.0, 10.0, 11.0], [24.0, 25.0, 26.0, 27.0], [56.0, 57.0, 58.0, 59.0]]
truncate_left max=3: dtype=torch.int64 shape=(3, 3)
    [[7, 0, 0], [0, 0, 0], [3, 4, 5]]
pad_to_multiple pad0 mult4: dtype=torch.int64 shape=(3, 8)
    [[5, 6, 7, 0, 0, 0, 0, 0], [8, 9, 0, 0, 0, 0, 0, 0], [1, 2, 3, 4, 5, 0, 0, 0]]
right_pad_to pad0 w7: dtype=torch.int64 shape=(3, 7)
    [[5, 6, 7, 0, 0, 0, 0], [8, 9, 0, 0, 0, 0, 0], [1, 2, 3, 4, 5, 0, 0]]

Probe RL and warper functions

shell
cat > /tmp/p4.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R
def show(name, x):
    if torch.is_tensor(x): print(f"{name}: dtype={x.dtype} shape={tuple(x.shape)}\n    {x.tolist()}")
    else: print(f"{name}: {x!r}")

ids = torch.tensor([[5,6,7,8,9],[1,2,3,4,0]])
show("completion_loss pl=[2,3] ig-100", R.completion_loss_labels(ids, torch.tensor([2,3]), -100))
show("prepare_labels pl=[2,3] ig-100", R.prepare_labels(ids, torch.tensor([2,3]), -100))

g=torch.Generator().manual_seed(3)
lg=torch.randn(2,3,4,generator=g)
acts=torch.tensor([[1,2,0],[3,1,2]])
show("action_logprobs", R.action_logprobs(lg, acts))

rew=torch.tensor([[1.0,2.0,3.0,4.0],[5.0,6.0,7.0,8.0]])
msk=torch.tensor([[1.0,1.0,1.0,0.0],[1.0,1.0,0.0,0.0]])
show("reward_to_go", R.reward_to_go(rew,msk))
x=torch.tensor([1.0,2.0,3.0,4.0])
show("discount_cumsum g=0.9", R.discount_cumsum(x,0.9))
x2=torch.tensor([[1.0,2.0,3.0],[4.0,5.0,6.0]])
show("discount_cumsum 2d g=0.5", R.discount_cumsum(x2,0.5))
show("build_pos_packed [2,3,1] tot6", R.build_position_ids_packed(torch.tensor([2,3,1]),6))

lw = torch.randn(2,6,generator=torch.Generator().manual_seed(5))
print("LW logits", lw.tolist())
show("logits_warper T=0.5 k=3 p=0.9", R.logits_warper(lw.clone(),0.5,3,0.9))
show("sample_greedy_warped T=0.5 k=3 p=0.9", R.sample_greedy_warped(lw.clone(),0.5,3,0.9))
show("first_stop_index stop[7,3]", R.first_stop_index(ids,[7,3]))
mask2=torch.tensor([[1.0,0.0,1.0,0.0,1.0,0.0],[0.0,0.0,1.0,1.0,1.0,1.0]])
show("masked_argmax", R.masked_argmax(lw, mask2))
show("generation_stop_mask eos=7 min=1", R.generation_stop_mask(ids,7,1))
show("rollout_last_logits", R.rollout_last_logits(torch.arange(2*5*4).float().reshape(2,5,4), ids, 0))
EOF
python3 /tmp/p4.py
output
Exit code 1
completion_loss pl=[2,3] ig-100: dtype=torch.int64 shape=(2, 5)
    [[-100, -100, 7, 8, 9], [-100, -100, -100, 4, 0]]
prepare_labels pl=[2,3] ig-100: dtype=torch.int64 shape=(2, 5)
    [[-100, 7, 8, 9, -100], [-100, -100, 4, 0, -100]]
action_logprobs: dtype=torch.float32 shape=(2, 3)
    [[-0.998786985874176, -0.45303505659103394, -1.4590954780578613], [-3.039785385131836, -1.8894623517990112, -0.9366484880447388]]
reward_to_go: dtype=torch.float32 shape=(2, 4)
    [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]]
discount_cumsum g=0.9: dtype=torch.float32 shape=(4,)
    [8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]
discount_cumsum 2d g=0.5: dtype=torch.float32 shape=(2, 3)
    [[2.75, 3.5, 3.0], [8.0, 8.0, 6.0]]
build_pos_packed [2,3,1] tot6: dtype=torch.int64 shape=(6,)
    [0, 1, 0, 1, 2, 0]
LW logits [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
logits_warper T=0.5 k=3 p=0.9: dtype=torch.float32 shape=(2, 6)
    [[-inf, -inf, -inf, -inf, -inf, 3.8855669498443604], [-inf, -inf, 1.2624268531799316, -inf, -inf, 1.3814375400543213]]
sample_greedy_warped T=0.5 k=3 p=0.9: dtype=torch.int64 shape=(2,)
    [5, 5]
Traceback (most recent call last):
  File "/tmp/p4.py", line 29, in <module>
    show("first_stop_index stop[7,3]", R.first_stop_index(ids,[7,3]))
                                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "reforacle.py", line 243, in reforacle.first_stop_index
AttributeError: 'list' object has no attribute 'tolist'

[error] tool reported failure

Probe tail functions

shell
cat > /tmp/p5.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R
def show(name, x):
    if torch.is_tensor(x): print(f"{name}: dtype={x.dtype} shape={tuple(x.shape)}\n    {x.tolist()}")
    else: print(f"{name}: {x!r}")

ids = torch.tensor([[5,6,7,8,9],[1,2,3,4,0]])
show("first_stop_index stop[7,3]", R.first_stop_index(ids, torch.tensor([7,3])))
show("first_stop_index stop[99]", R.first_stop_index(ids, torch.tensor([99])))
show("first_stop_index stop[9,4]", R.first_stop_index(ids, torch.tensor([9,4])))
lw = torch.randn(2,6,generator=torch.Generator().manual_seed(5))
mask2=torch.tensor([[1.0,0.0,1.0,0.0,1.0,0.0],[0.0,0.0,1.0,1.0,1.0,1.0]])
show("masked_argmax", R.masked_argmax(lw, mask2))
show("masked_argmax boolmask", R.masked_argmax(lw, mask2.bool()))
show("gen_stop_mask eos=7 min=1", R.generation_stop_mask(ids,7,1))
show("gen_stop_mask eos=7 min=4", R.generation_stop_mask(ids,7,4))
show("gen_stop_mask eos=9 min=1", R.generation_stop_mask(ids,9,1))
show("rollout_last_logits", R.rollout_last_logits(torch.arange(2*5*4).float().reshape(2,5,4), ids, 0))
# action_logprobs numeric check
g=torch.Generator().manual_seed(3); lg=torch.randn(2,3,4,generator=g)
print("check", torch.log_softmax(lg,-1).gather(-1, torch.tensor([[1,2,0],[3,1,2]]).unsqueeze(-1)).squeeze(-1).tolist())
EOF
python3 /tmp/p5.py
output
first_stop_index stop[7,3]: dtype=torch.int64 shape=(2,)
    [2, 2]
first_stop_index stop[99]: dtype=torch.int64 shape=(2,)
    [-1, -1]
first_stop_index stop[9,4]: dtype=torch.int64 shape=(2,)
    [4, 3]
masked_argmax: dtype=torch.int64 shape=(2,)
    [4, 5]
masked_argmax boolmask: dtype=torch.int64 shape=(2,)
    [4, 5]
gen_stop_mask eos=7 min=1: dtype=torch.float32 shape=(2, 5)
    [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
gen_stop_mask eos=7 min=4: dtype=torch.float32 shape=(2, 5)
    [[1.0, 1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
gen_stop_mask eos=9 min=1: dtype=torch.float32 shape=(2, 5)
    [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
rollout_last_logits: dtype=torch.float32 shape=(2, 4)
    [[16.0, 17.0, 18.0, 19.0], [32.0, 33.0, 34.0, 35.0]]
check [[-0.998786985874176, -0.45303505659103394, -1.4590954780578613], [-3.039785385131836, -1.8894623517990112, -0.9366484880447388]]

[stdout]
first_stop_index stop[7,3]: dtype=torch.int64 shape=(2,)
    [2, 2]
first_stop_index stop[99]: dtype=torch.int64 shape=(2,)
    [-1, -1]
first_stop_index stop[9,4]: dtype=torch.int64 shape=(2,)
    [4, 3]
masked_argmax: dtype=torch.int64 shape=(2,)
    [4, 5]
masked_argmax boolmask: dtype=torch.int64 shape=(2,)
    [4, 5]
gen_stop_mask eos=7 min=1: dtype=torch.float32 shape=(2, 5)
    [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
gen_stop_mask eos=7 min=4: dtype=torch.float32 shape=(2, 5)
    [[1.0, 1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
gen_stop_mask eos=9 min=1: dtype=torch.float32 shape=(2, 5)
    [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
rollout_last_logits: dtype=torch.float32 shape=(2, 4)
    [[16.0, 17.0, 18.0, 19.0], [32.0, 33.0, 34.0, 35.0]]
check [[-0.998786985874176, -0.45303505659103394, -1.4590954780578613], [-3.039785385131836, -1.8894623517990112, -0.9366484880447388]]

Probe edge cases

shell
cat > /tmp/p6.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R
def show(name, x):
    if torch.is_tensor(x): print(f"{name}: {x.tolist()}")
    else: print(f"{name}: {x!r}")
lw = torch.randn(2,6,generator=torch.Generator().manual_seed(5))
# epsilon larger than max prob
show("epsilon 0.9 (all below?)", R.epsilon_filter(lw.clone(),0.9))
show("epsilon 0.5", R.epsilon_filter(lw.clone(),0.5))
show("min_p 1.5 large", R.min_p_filter(lw.clone(),1.5))
show("min_p 0.999", R.min_p_filter(lw.clone(),0.999))
show("top_k 0", R.top_k_filter(lw.clone(),0))
show("top_k 10 (>vocab)", R.top_k_filter(lw.clone(),10))
show("top_p 0.0", R.top_p_filter(lw.clone(),0.0))
show("top_p 1.0", R.top_p_filter(lw.clone(),1.0))
# multiple eos in mask_after_eos
ids=torch.tensor([[7,1,7,2,3]])
show("mask_after_eos two eos", R.mask_after_eos(ids,7))
show("gen_stop_mask two eos min1", R.generation_stop_mask(ids,7,1))
# right_pad_to width smaller than current -> truncate?
show("right_pad_to w3 (smaller)", R.right_pad_to(torch.tensor([[1,2,3,4,5]]),0,3))
# temperature negative
show("temp T=-1", R.temperature_scale(torch.tensor([[1.0,2.0,3.0]]),-1.0))
EOF
python3 /tmp/p6.py
output
Exit code 1
epsilon 0.9 (all below?): [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
epsilon 0.5: [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
min_p 1.5 large: [[-inf, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, -inf]]
min_p 0.999: [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
Traceback (most recent call last):
  File "/tmp/p6.py", line 12, in <module>
    show("top_k 0", R.top_k_filter(lw.clone(),0))
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  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

More edge cases

shell
cat > /tmp/p7.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R
def show(name, x):
    if torch.is_tensor(x): print(f"{name}: {x.tolist()}")
    else: print(f"{name}: {x!r}")
lw = torch.randn(2,6,generator=torch.Generator().manual_seed(5))
show("top_k 10 (>vocab)", R.top_k_filter(lw.clone(),10))
show("top_p 1.0", R.top_p_filter(lw.clone(),1.0))
show("top_p 0.0", R.top_p_filter(lw.clone(),0.0))
ids=torch.tensor([[7,1,7,2,3]])
show("mask_after_eos two eos", R.mask_after_eos(ids,7))
show("gen_stop_mask two eos min1", R.generation_stop_mask(ids,7,1))
show("right_pad_to w3 smaller", R.right_pad_to(torch.tensor([[1,2,3,4,5]]),0,3))
show("temp T=-1", R.temperature_scale(torch.tensor([[1.0,2.0,3.0]]),-1.0))
show("truncate_left max=10 bigger", R.truncate_left(torch.tensor([[1,2,3]]),10))
# first_stop when stop empty? and first at index 0
show("first_stop idx0", R.first_stop_index(torch.tensor([[7,1,2]]),torch.tensor([7])))
# min_length eos as list? 
show("min_len eos negative cur", R.min_length_filter(lw.clone(),0,3,2))
# discount_cumsum 1d confirm gamma 0
show("discount g=0", R.discount_cumsum(torch.tensor([1.,2.,3.]),0.0))
EOF
python3 /tmp/p7.py
output
top_k 10 (>vocab): [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
top_p 1.0: [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
top_p 0.0: [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
mask_after_eos two eos: [[1.0, 0.0, 0.0, 0.0, 0.0]]
gen_stop_mask two eos min1: [[1.0, 0.0, 0.0, 0.0, 0.0]]
right_pad_to w3 smaller: [[1, 2, 3]]
temp T=-1: [[-1.0, -2.0, -3.0]]
truncate_left max=10 bigger: [[1, 2, 3]]
first_stop idx0: [0]
min_len eos negative cur: [[-0.4867807626724243, -0.6038215756416321, -inf, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, -inf, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
discount g=0: [1.0, 2.0, 3.0]

[stdout]
top_k 10 (>vocab): [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
top_p 1.0: [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
top_p 0.0: [[-inf, -inf, -inf, -inf, -inf, 1.9427834749221802], [-inf, -inf, -inf, -inf, -inf, 0.6907187700271606]]
mask_after_eos two eos: [[1.0, 0.0, 0.0, 0.0, 0.0]]
gen_stop_mask two eos min1: [[1.0, 0.0, 0.0, 0.0, 0.0]]
right_pad_to w3 smaller: [[1, 2, 3]]
temp T=-1: [[-1.0, -2.0, -3.0]]
truncate_left max=10 bigger: [[1, 2, 3]]
first_stop idx0: [0]
min_len eos negative cur: [[-0.4867807626724243, -0.6038215756416321, -inf, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, -inf, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
discount g=0: [1.0, 2.0, 3.0]

Pin down index conventions

shell
cat > /tmp/p8.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R
def show(name, x):
    if torch.is_tensor(x): print(f"{name}: {x.tolist()}")
    else: print(f"{name}: {x!r}")
# left padded / non-contiguous masks
m = torch.tensor([[0.,0.,1.,1.,1.],[0.,1.,1.,0.,0.],[1.,1.,1.,1.,1.]])
show("last_token_indices noncontig", R.last_token_indices(m))
lg=torch.arange(3*5*4).float().reshape(3,5,4)
show("gather_last noncontig", R.gather_last_logits(lg, m))
# pad_to_multiple already multiple
show("pad_to_mult already", R.pad_to_multiple(torch.tensor([[1,2,3,4]]),0,4))
show("pad_to_mult mult1", R.pad_to_multiple(torch.tensor([[1,2,3]]),0,1))
# sequence_lengths with pad in middle (non-contiguous)
show("seqlen noncontig", R.sequence_lengths(torch.tensor([[1,0,2,0,0]]),0))
# rollout_last with all pad row
show("rollout allpad", R.rollout_last_logits(torch.arange(1*3*4).float().reshape(1,3,4), torch.tensor([[0,0,0]]),0))
EOF
python3 /tmp/p8.py
output
last_token_indices noncontig: [4, 2, 4]
gather_last noncontig: [[16.0, 17.0, 18.0, 19.0], [28.0, 29.0, 30.0, 31.0], [56.0, 57.0, 58.0, 59.0]]
pad_to_mult already: [[1, 2, 3, 4]]
pad_to_mult mult1: [[1, 2, 3]]
seqlen noncontig: [2]
rollout allpad: [[0.0, 1.0, 2.0, 3.0]]

[stdout]
last_token_indices noncontig: [4, 2, 4]
gather_last noncontig: [[16.0, 17.0, 18.0, 19.0], [28.0, 29.0, 30.0, 31.0], [56.0, 57.0, 58.0, 59.0]]
pad_to_mult already: [[1, 2, 3, 4]]
pad_to_mult mult1: [[1, 2, 3]]
seqlen noncontig: [2]
rollout allpad: [[0.0, 1.0, 2.0, 3.0]]

rollout vs count conventions

shell
cat > /tmp/p9.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R
def show(name, x):
    if torch.is_tensor(x): print(f"{name}: {x.tolist()}")
    else: print(f"{name}: {x!r}")
# rollout with non-contiguous pad
show("rollout noncontig ids[1,0,2,0,0]", R.rollout_last_logits(torch.arange(1*5*4).float().reshape(1,5,4), torch.tensor([[1,0,2,0,0]]),0))
# last_token_indices all zero
try: show("last_tok allzero", R.last_token_indices(torch.tensor([[0.,0.,0.]])))
except Exception as e: print("last_tok allzero ERR", e)
# gather_last all zero mask
try: show("gather allzero mask", R.gather_last_logits(torch.arange(1*3*4).float().reshape(1,3,4), torch.tensor([[0.,0.,0.]])))
except Exception as e: print("gather allzero ERR", e)
# sequence_lengths noncontig gave 2 (count). last_token_indices uses last idx. rollout?
EOF
python3 /tmp/p9.py
output
rollout noncontig ids[1,0,2,0,0]: [[8.0, 9.0, 10.0, 11.0]]
last_tok allzero: [-1]
gather allzero mask: [[0.0, 1.0, 2.0, 3.0]]

[stdout]
rollout noncontig ids[1,0,2,0,0]: [[8.0, 9.0, 10.0, 11.0]]
last_tok allzero: [-1]
gather allzero mask: [[0.0, 1.0, 2.0, 3.0]]

Final convention checks

shell
cat > /tmp/p10.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R
def show(name, x):
    if torch.is_tensor(x): print(f"{name}: {x.tolist()} dtype={x.dtype}")
    else: print(f"{name}: {x!r}")
show("packed sum<total (tot=8)", R.build_position_ids_packed(torch.tensor([2,3,1]),8))
show("packed sum==total", R.build_position_ids_packed(torch.tensor([4]),4))
# entropy 3d?
g=torch.Generator().manual_seed(2); lg=torch.randn(2,3,4,generator=g)
show("entropy 3d", R.entropy_from_logits(lg))
# action_logprobs 2d logits
show("action_logprobs 2d", R.action_logprobs(torch.randn(3,5,generator=torch.Generator().manual_seed(9)), torch.tensor([0,4,2])))
# suppress with tensor arg
show("suppress tensor", R.suppress_tokens(torch.zeros(1,4), torch.tensor([1,3])))
EOF
python3 /tmp/p10.py
output
packed sum<total (tot=8): [0, 1, 0, 1, 2, 0, 0, 0] dtype=torch.int64
packed sum==total: [0, 1, 2, 3] dtype=torch.int64
entropy 3d: [[0.8910926580429077, 0.9417887926101685, 1.1524007320404053], [1.3753901720046997, 1.1342965364456177, 1.280873417854309]] dtype=torch.float32
action_logprobs 2d: [-2.5763306617736816, -1.146837830543518, -2.037553071975708] dtype=torch.float32
suppress tensor: [[0.0, -inf, 0.0, -inf]] dtype=torch.float32

[stdout]
packed sum<total (tot=8): [0, 1, 0, 1, 2, 0, 0, 0] dtype=torch.int64
packed sum==total: [0, 1, 2, 3] dtype=torch.int64
entropy 3d: [[0.8910926580429077, 0.9417887926101685, 1.1524007320404053], [1.3753901720046997, 1.1342965364456177, 1.280873417854309]] dtype=torch.float32
action_logprobs 2d: [-2.5763306617736816, -1.146837830543518, -2.037553071975708] dtype=torch.float32
suppress tensor: [[0.0, -inf, 0.0, -inf]] dtype=torch.float32
write · /app/objectives.py
from __future__ import annotations

import torch

NEG_INF = float("-inf")


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


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


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


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


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


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


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


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


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


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


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


def attention_mask(input_ids, pad_id):
    return (input_ids != pad_id).float()


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


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


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


def completion_loss_labels(input_ids, prompt_lengths, ignore_index):
    labels = input_ids.clone()
    positions = torch.arange(input_ids.size(-1), device=input_ids.device)
    mask = positions.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
    has = is_eos.any(dim=-1)
    first = is_eos.int().argmax(dim=-1)
    T = input_ids.size(-1)
    first = torch.where(has, first, torch.full_like(first, T))
    positions = torch.arange(T, device=input_ids.device)
    return (positions.unsqueeze(0) <= first.unsqueeze(-1)).float()


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


def last_token_indices(attn_mask):
    positions = torch.arange(attn_mask.size(-1), device=attn_mask.device)
    masked = torch.where(attn_mask > 0, positions, torch.full_like(positions, -1))
    return masked.max(dim=-1).values


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


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


def pad_to_multiple(input_ids, pad_id, multiple):
    T = input_ids.size(-1)
    pad_amt = (-T) % multiple
    if pad_amt == 0:
        return input_ids
    pad = torch.full((input_ids.size(0), pad_amt), pad_id, dtype=input_ids.dtype, device=input_ids.device)
    return torch.cat([input_ids, pad], dim=-1)


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


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


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


def discount_cumsum(x, gamma):
    out = x.clone()
    T = x.size(-1)
    for t in range(T - 2, -1, -1):
        out[..., t] = x[..., t] + gamma * out[..., t + 1]
    return out


def build_position_ids_packed(seq_lengths, total_len):
    parts = [torch.arange(int(l)) for l in seq_lengths]
    packed = torch.cat(parts) if parts else torch.zeros(0, dtype=torch.long)
    out = torch.zeros(total_len, dtype=torch.long)
    n = min(packed.numel(), total_len)
    out[:n] = packed[:n]
    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 = torch.full_like(input_ids, ignore_index)
    labels[..., :-1] = input_ids[..., 1:]
    positions = torch.arange(input_ids.size(-1), device=input_ids.device)
    mask = positions.unsqueeze(0) < (prompt_lengths - 1).unsqueeze(-1)
    labels[mask] = ignore_index
    return labels


def rollout_last_logits(logits, input_ids, pad_id):
    idx = last_token_indices((input_ids != pad_id)).clamp(min=0)
    b = torch.arange(logits.size(0), device=logits.device)
    return logits[b, idx]


def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    remove = probs < epsilon
    top1 = torch.topk(probs, 1, dim=-1)[0][..., -1, None]
    remove = remove & (probs < top1)
    return logits.masked_fill(remove, NEG_INF)


def first_stop_index(input_ids, stop_ids):
    is_stop = (input_ids.unsqueeze(-1) == stop_ids).any(dim=-1)
    has = is_stop.any(dim=-1)
    first = is_stop.int().argmax(dim=-1)
    return torch.where(has, first, torch.full_like(first, -1))


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


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


def right_pad_to(input_ids, pad_id, width):
    T = input_ids.size(-1)
    if T == width:
        return input_ids
    if T > width:
        return input_ids[..., :width]
    pad = torch.full((input_ids.size(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
    has = is_eos.any(dim=-1)
    first = is_eos.int().argmax(dim=-1)
    T = input_ids.size(-1)
    stop = torch.clamp(first, min=min_len - 1)
    stop = torch.where(has, stop, torch.full_like(first, T))
    positions = torch.arange(T, device=input_ids.device)
    return (positions.unsqueeze(0) <= stop.unsqueeze(-1)).float()

Comprehensive verification

shell
cat > /tmp/verify.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R, objectives as O

def eq(a, b, atol=1e-5, rtol=1e-4):
    if torch.is_tensor(a):
        if a.dtype != b.dtype: return f"DTYPE {a.dtype} vs {b.dtype}"
        if a.shape != b.shape: return f"SHAPE {a.shape} vs {b.shape}"
        af, bf = a.float(), b.float()
        # handle inf
        m = torch.isinf(af) | torch.isinf(bf)
        if not torch.equal(torch.isinf(af)&(af<0), torch.isinf(bf)&(bf<0)): return "NEGINF mismatch"
        if not torch.equal(torch.isinf(af)&(af>0), torch.isinf(bf)&(bf>0)): return "POSINF mismatch"
        af=af.clone(); bf=bf.clone(); af[m]=0; bf[m]=0
        if torch.allclose(af, bf, atol=atol, rtol=rtol): return None
        return f"VAL max diff {(af-bf).abs().max().item()}"
    return None if a==b else f"{a} vs {b}"

fails=[]
def check(name, ref, mine):
    r = eq(ref, mine)
    if r: fails.append((name, r))

G=lambda s: torch.Generator().manual_seed(s)
for s in range(30):
    B=torch.randint(1,4,(1,),generator=G(s)).item()
    V=torch.randint(2,10,(1,),generator=G(s+100)).item()
    T=torch.randint(1,7,(1,),generator=G(s+200)).item()
    lg=torch.randn(B,V,generator=G(s+1))
    ids=torch.randint(0,V,(B,T),generator=G(s+2))
    # decoding
    check("greedy",R.greedy_token(lg),O.greedy_token(lg))
    for Tp in [0.0,0.5,1.0,2.0,-1.0]:
        check(f"temp{Tp}",R.temperature_scale(lg,Tp),O.temperature_scale(lg,Tp))
    for k in [1,2,V]:
        check(f"topk{k}",R.top_k_filter(lg.clone(),k),O.top_k_filter(lg.clone(),k))
        check(f"topki{k}",R.topk_indices(lg,k),O.topk_indices(lg,k))
    for p in [0.0,0.3,0.5,0.9,1.0]:
        check(f"topp{p}",R.top_p_filter(lg.clone(),p),O.top_p_filter(lg.clone(),p))
        check(f"minp{p}",R.min_p_filter(lg.clone(),p),O.min_p_filter(lg.clone(),p))
        check(f"eps{p}",R.epsilon_filter(lg.clone(),p),O.epsilon_filter(lg.clone(),p))
    for pen in [0.5,1.0,1.5,2.0]:
        check(f"rep{pen}",R.repetition_penalty(lg.clone(),ids,pen),O.repetition_penalty(lg.clone(),ids,pen))
        check(f"freq{pen}",R.frequency_penalty(lg.clone(),ids,pen),O.frequency_penalty(lg.clone(),ids,pen))
        check(f"pres{pen}",R.presence_penalty(lg.clone(),ids,pen),O.presence_penalty(lg.clone(),ids,pen))
    check("suppress",R.suppress_tokens(lg.clone(),[0, V-1]),O.suppress_tokens(lg.clone(),[0,V-1]))
    for cl,ml in [(0,3),(3,3),(5,3)]:
        check(f"minlen{cl}",R.min_length_filter(lg.clone(),cl,ml,0),O.min_length_filter(lg.clone(),cl,ml,0))
    check("renorm",R.renormalize_logits(lg),O.renormalize_logits(lg))
    check("entropy",R.entropy_from_logits(lg),O.entropy_from_logits(lg))
    check("warper",R.logits_warper(lg.clone(),0.7,2,0.9),O.logits_warper(lg.clone(),0.7,2,0.9))
    check("samplegw",R.sample_greedy_warped(lg.clone(),0.7,2,0.9),O.sample_greedy_warped(lg.clone(),0.7,2,0.9))
    # collation - pad both left and right style: use raw ids with pad 0
    padded=ids.clone()
    check("attn",R.attention_mask(padded,0),O.attention_mask(padded,0))
    check("posids",R.position_ids(padded,0),O.position_ids(padded,0))
    check("seqlen",R.sequence_lengths(padded,0),O.sequence_lengths(padded,0))
    check("causal",R.causal_mask(T),O.causal_mask(T))
    check("shift",R.shift_labels(padded,-100),O.shift_labels(padded,-100))
    pl=torch.randint(0,T+1,(B,),generator=G(s+7))
    check("comploss",R.completion_loss_labels(padded,pl,-100),O.completion_loss_labels(padded,pl,-100))
    check("prep",R.prepare_labels(padded,pl,-100),O.prepare_labels(padded,pl,-100))
    check("maskeos",R.mask_after_eos(padded,0),O.mask_after_eos(padded,0))
    m=(padded!=0).float()
    check("lasttok",R.last_token_indices(m),O.last_token_indices(m))
    lg3=torch.randn(B,T,V,generator=G(s+9))
    check("gatherlast",R.gather_last_logits(lg3,m),O.gather_last_logits(lg3,m))
    check("rollout",R.rollout_last_logits(lg3,padded,0),O.rollout_last_logits(lg3,padded,0))
    check("trunc",R.truncate_left(padded,2),O.truncate_left(padded,2))
    check("pad2mult",R.pad_to_multiple(padded,0,4),O.pad_to_multiple(padded,0,4))
    check("rpad",R.right_pad_to(padded,0,T+2),O.right_pad_to(padded,0,T+2))
    if T>=2: check("rpad2",R.right_pad_to(padded,0,T-1),O.right_pad_to(padded,0,T-1))
    acts=torch.randint(0,V,(B,T),generator=G(s+11))
    check("actlp",R.action_logprobs(lg3,acts),O.action_logprobs(lg3,acts))
    rew=torch.randn(B,T,generator=G(s+12)); msk=(torch.rand(B,T,generator=G(s+13))>0.3).float()
    check("rtg",R.reward_to_go(rew,msk),O.reward_to_go(rew,msk))
    check("disc",R.discount_cumsum(rew,0.95),O.discount_cumsum(rew,0.95))
    sl=torch.randint(0,4,(3,),generator=G(s+14))
    tot=int(sl.sum().item())+2
    check("packed",R.build_position_ids_packed(sl,tot),O.build_position_ids_packed(sl,tot))
    stop=torch.tensor([0,V-1])
    check("firststop",R.first_stop_index(padded,stop),O.first_stop_index(padded,stop))
    check("maskargmax",R.masked_argmax(lg,m if m.shape==lg.shape else (torch.rand(B,V,generator=G(s+15))>0.4).float()),O.masked_argmax(lg,m if m.shape==lg.shape else (torch.rand(B,V,generator=G(s+15))>0.4).float()))
    for ml2 in [1,2,T]:
        check(f"genstop{ml2}",R.generation_stop_mask(padded,0,ml2),O.generation_stop_mask(padded,0,ml2))

from collections import Counter
c=Counter(f[0].rstrip('0123456789.-') for f in fails)
print("TOTAL FAILS:", len(fails))
for k,v in c.items(): print("  ",k,v)
for f in fails[:20]: print("   detail",f)
EOF
python3 /tmp/verify.py
output
TOTAL FAILS: 3
   maskargmax 3
   detail ('maskargmax', 'VAL max diff 1.0')
   detail ('maskargmax', 'VAL max diff 1.0')
   detail ('maskargmax', 'VAL max diff 1.0')

[stdout]
TOTAL FAILS: 3
   maskargmax 3
   detail ('maskargmax', 'VAL max diff 1.0')
   detail ('maskargmax', 'VAL max diff 1.0')
   detail ('maskargmax', 'VAL max diff 1.0')

Isolate masked_argmax

shell
cat > /tmp/vm.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R, objectives as O
bad=0
for s in range(200):
    g=torch.Generator().manual_seed(s)
    B=torch.randint(1,4,(1,),generator=g).item()
    V=torch.randint(2,8,(1,),generator=g).item()
    lg=torch.randn(B,V,generator=g)
    mask=(torch.rand(B,V,generator=g)>0.4).float()
    r=R.masked_argmax(lg,mask); m=O.masked_argmax(lg,mask)
    if not torch.equal(r,m):
        bad+=1
        if bad<=5:
            print("MISMATCH", lg.tolist(), mask.tolist(), "ref",r.tolist(),"mine",m.tolist())
    # bool mask
    r2=R.masked_argmax(lg,mask.bool()); m2=O.masked_argmax(lg,mask.bool())
    if not torch.equal(r2,m2): print("BOOL mismatch", r2.tolist(), m2.tolist())
print("bad", bad)
EOF
python3 /tmp/vm.py
output
MISMATCH [[-0.0396014042198658, -0.34204018115997314], [2.282931327819824, -0.851894736289978]] [[0.0, 0.0], [1.0, 1.0]] ref [-1, 0] mine [0, 0]
BOOL mismatch [-1, 0] [0, 0]
MISMATCH [[-1.3282146453857422, 0.5064482688903809]] [[0.0, 0.0]] ref [-1] mine [0]
BOOL mismatch [-1] [0]
MISMATCH [[-0.6650162935256958, -1.452513575553894, -1.1388238668441772], [0.9992223978042603, 1.979245662689209, 0.10843940079212189], [-0.6909270882606506, 0.7321369647979736, 0.2862148880958557]] [[0.0, 1.0, 1.0], [0.0, 0.0, 0.0], [1.0, 0.0, 1.0]] ref [2, -1, 2] mine [2, 0, 2]
BOOL mismatch [2, -1, 2] [2, 0, 2]
MISMATCH [[0.7080374956130981, -1.5891871452331543], [0.28969448804855347, -0.6696841716766357]] [[0.0, 1.0], [0.0, 0.0]] ref [1, -1] mine [1, 0]
BOOL mismatch [1, -1] [1, 0]
MISMATCH [[-1.3775849342346191, 0.7566263675689697, -0.026567790657281876, -0.1914340853691101, 1.2311701774597168, 0.7105574011802673], [1.794904112815857, 1.01513671875, -0.40477046370506287, 1.0466219186782837, 0.05193345621228218, 0.9268918037414551]] [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 1.0, 1.0]] ref [-1, 5] mine [0, 5]
BOOL mismatch [-1, 5] [0, 5]
BOOL mismatch [1, 2, -1] [1, 2, 0]
BOOL mismatch [1, -1, 1] [1, 0, 1]
BOOL mismatch [2, -1, 2] [2, 0, 2]
BOOL mismatch [-1, 0] [0, 0]
BOOL mismatch [-1, 1] [0, 1]
BOOL mismatch [-1, 1] [0, 1]
BOOL mismatch [-1] [0]
bad 12

[stdout]
MISMATCH [[-0.0396014042198658, -0.34204018115997314], [2.282931327819824, -0.851894736289978]] [[0.0, 0.0], [1.0, 1.0]] ref [-1, 0] mine [0, 0]
BOOL mismatch [-1, 0] [0, 0]
MISMATCH [[-1.3282146453857422, 0.5064482688903809]] [[0.0, 0.0]] ref [-1] mine [0]
BOOL mismatch [-1] [0]
MISMATCH [[-0.6650162935256958, -1.452513575553894, -1.1388238668441772], [0.9992223978042603, 1.979245662689209, 0.10843940079212189], [-0.6909270882606506, 0.7321369647979736, 0.2862148880958557]] [[0.0, 1.0, 1.0], [0.0, 0.0, 0.0], [1.0, 0.0, 1.0]] ref [2, -1, 2] mine [2, 0, 2]
BOOL mismatch [2, -1, 2] [2, 0, 2]
MISMATCH [[0.7080374956130981, -1.5891871452331543], [0.28969448804855347, -0.6696841716766357]] [[0.0, 1.0], [0.0, 0.0]] ref [1, -1] mine [1, 0]
BOOL mismatch [1, -1] [1, 0]
MISMATCH [[-1.3775849342346191, 0.7566263675689697, -0.026567790657281876, -0.1914340853691101, 1.2311701774597168, 0.7105574011802673], [1.794904112815857, 1.01513671875, -0.40477046370506287, 1.0466219186782837, 0.05193345621228218, 0.9268918037414551]] [[0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 1.0, 1.0]] ref [-1, 5] mine [0, 5]
BOOL mismatch [-1, 5] [0, 5]
BOOL mismatch [1, 2, -1] [1, 2, 0]
BOOL mismatch [1, -1, 1] [1, 0, 1]
BOOL mismatch [2, -1, 2] [2, 0, 2]
BOOL mismatch [-1, 0] [0, 0]
BOOL mismatch [-1, 1] [0, 1]
BOOL mismatch [-1, 1] [0, 1]
BOOL mismatch [-1] [0]
bad 12

/app/objectives.py

− removed
def masked_argmax(logits, mask):
    filled = logits.masked_fill(mask == 0, NEG_INF)
    return filled.argmax(dim=-1)
+ added
def masked_argmax(logits, mask):
    filled = logits.masked_fill(mask == 0, NEG_INF)
    idx = filled.argmax(dim=-1)
    valid = (mask > 0).any(dim=-1)
    return torch.where(valid, idx, torch.full_like(idx, -1))

Re-verify masked_argmax and full suite

shell
python3 /tmp/vm.py 2>&1 | tail -2; echo "=== full ==="; python3 /tmp/verify.py 2>&1 | head -5
output
bad 0
=== full ===
TOTAL FAILS: 0

[stdout]
bad 0
=== full ===
TOTAL FAILS: 0

Hard stress test

shell
cat > /tmp/v2.py << 'EOF'
import sys; sys.path.insert(0, "/app")
import torch, reforacle as R, objectives as O
from collections import Counter
def eq(a,b,atol=1e-5,rtol=1e-4):
    if torch.is_tensor(a):
        if a.dtype!=b.dtype: return f"DTYPE {a.dtype}/{b.dtype}"
        if a.shape!=b.shape: return f"SHAPE {a.shape}/{b.shape}"
        af,bf=a.float(),b.float()
        if not torch.equal(af.isinf()&(af<0),bf.isinf()&(bf<0)): return "NEGINF"
        if not torch.equal(af.isinf()&(af>0),bf.isinf()&(bf>0)): return "POSINF"
        m=af.isinf()|bf.isinf(); af=af.clone();bf=bf.clone();af[m]=0;bf[m]=0
        return None if torch.allclose(af,bf,atol=atol,rtol=rtol) else f"VAL {(af-bf).abs().max().item()}"
    return None if a==b else f"{a}/{b}"
fails=[]
def ck(n,r,m):
    e=eq(r,m)
    if e: fails.append((n,e,))
G=lambda s: torch.Generator().manual_seed(s)
for s in range(60):
    g=G(s); B=torch.randint(1,5,(1,),generator=g).item(); V=torch.randint(2,20,(1,),generator=g).item(); T=torch.randint(1,9,(1,),generator=g).item()
    pad=torch.randint(0,V,(1,),generator=g).item()
    lg=torch.randn(B,V,generator=G(s+1))*3
    ids=torch.randint(0,V,(B,T),generator=G(s+2))
    # force some all-pad rows
    if B>1: ids[0]=pad
    m=(ids!=pad).float()
    for fn in ['attention_mask','position_ids','sequence_lengths']:
        ck(fn+str(pad), getattr(R,fn)(ids,pad), getattr(O,fn)(ids,pad))
    ck("lasttok",R.last_token_indices(m),O.last_token_indices(m))
    lg3=torch.randn(B,T,V,generator=G(s+3))
    ck("gather",R.gather_last_logits(lg3,m),O.gather_last_logits(lg3,m))
    ck("rollout",R.rollout_last_logits(lg3,ids,pad),O.rollout_last_logits(lg3,ids,pad))
    ck("maskeos",R.mask_after_eos(ids,pad),O.mask_after_eos(ids,pad))
    for ml in [0,1,2,T,T+3]:
        ck("genstop",R.generation_stop_mask(ids,pad,ml),O.generation_stop_mask(ids,pad,ml))
    for pen in [0.7,1.3,3.0]:
        ck("rep",R.repetition_penalty(lg.clone(),ids,pen),O.repetition_penalty(lg.clone(),ids,pen))
        ck("freq",R.frequency_penalty(lg.clone(),ids,pen),O.frequency_penalty(lg.clone(),ids,pen))
        ck("pres",R.presence_penalty(lg.clone(),ids,pen),O.presence_penalty(lg.clone(),ids,pen))
    for p in [0.01,0.2,0.7,0.95,0.999]:
        ck("topp",R.top_p_filter(lg.clone(),p),O.top_p_filter(lg.clone(),p))
        ck("minp",R.min_p_filter(lg.clone(),p),O.min_p_filter(lg.clone(),p))
        ck("eps",R.epsilon_filter(lg.clone(),p),O.epsilon_filter(lg.clone(),p))
    for k in range(1,V+1):
        ck("topk",R.top_k_filter(lg.clone(),k),O.top_k_filter(lg.clone(),k))
        ck("topki",R.topk_indices(lg,k),O.topk_indices(lg,k))
    # 1D discount / rtg
    x1=torch.randn(T,generator=G(s+5))
    for gm in [0.0,0.5,0.99,1.0]:
        ck("disc1d",R.discount_cumsum(x1,gm),O.discount_cumsum(x1,gm))
    rew=torch.randn(B,T,generator=G(s+6)); msk=(torch.rand(B,T,generator=G(s+7))>0.5).float()
    ck("rtg",R.reward_to_go(rew,msk),O.reward_to_go(rew,msk))
    ck("disc2d",R.discount_cumsum(rew,0.9),O.discount_cumsum(rew,0.9))
    # action logprobs 2D
    ck("actlp2d",R.action_logprobs(lg,ids[:,0]),O.action_logprobs(lg,ids[:,0]))
    ck("actlp3d",R.action_logprobs(lg3,ids),O.action_logprobs(lg3,ids))
    ck("entropy",R.entropy_from_logits(lg),O.entropy_from_logits(lg))
    ck("entropy3d",R.entropy_from_logits(lg3),O.entropy_from_logits(lg3))
    pl=torch.randint(0,T+1,(B,),generator=G(s+8))
    ck("comp",R.completion_loss_labels(ids,pl,-100),O.completion_loss_labels(ids,pl,-100))
    ck("prep",R.prepare_labels(ids,pl,-100),O.prepare_labels(ids,pl,-100))
    ck("shift",R.shift_labels(ids,-100),O.shift_labels(ids,-100))
    for w in [1,T,T+3]:
        ck("rpad",R.right_pad_to(ids,pad,w),O.right_pad_to(ids,pad,w))
    for mu in [1,3,7]:
        ck("p2m",R.pad_to_multiple(ids,pad,mu),O.pad_to_multiple(ids,pad,mu))
    for mx in [1,T,T+2]:
        ck("trunc",R.truncate_left(ids,mx),O.truncate_left(ids,mx))
    ck("warp",R.logits_warper(lg.clone(),0.8,max(1,V//2),0.9),O.logits_warper(lg.clone(),0.8,max(1,V//2),0.9))
    ck("sgw",R.sample_greedy_warped(lg.clone(),0.8,max(1,V//2),0.9),O.sample_greedy_warped(lg.clone(),0.8,max(1,V//2),0.9))
    sl=torch.randint(0,5,(4,),generator=G(s+9)); tot=int(sl.sum())+torch.randint(0,4,(1,),generator=g).item()
    ck("packed",R.build_position_ids_packed(sl,tot),O.build_position_ids_packed(sl,tot))
    ck("suppress",R.suppress_tokens(lg.clone(),torch.tensor([0,V-1])),O.suppress_tokens(lg.clone(),torch.tensor([0,V-1])))
    ck("causal",R.causal_mask(T),O.causal_mask(T))
    ck("renorm",R.renormalize_logits(lg),O.renormalize_logits(lg))
    ck("greedy",R.greedy_token(lg),O.greedy_token(lg))
    for Tp in [0.0,0.3,1.0,2.5,-2.0]:
        ck("temp",R.temperature_scale(lg,Tp),O.temperature_scale(lg,Tp))
    stop=torch.tensor([pad,V-1])
    ck("fstop",R.first_stop_index(ids,stop),O.first_stop_index(ids,stop))
    mm=(torch.rand(B,V,generator=G(s+11))>0.5).float()
    ck("margmax",R.masked_argmax(lg,mm),O.masked_argmax(lg,mm))
    for cl in [0,T,T+2]:
        ck("minlen",R.min_length_filter(lg.clone(),cl,T,pad),O.min_length_filter(lg.clone(),cl,T,pad))
c=Counter(f[0] for f in fails)
print("FAILS",len(fails), dict(c))
for f in fails[:15]: print(f)
EOF
python3 /tmp/v2.py
output
FAILS 0 {}

[stdout]
FAILS 0 {}

Tight tolerance + standalone check

shell
cd /app 2>/dev/null; sed 's/atol=1e-5,rtol=1e-4/atol=1e-6,rtol=1e-5/' /tmp/v2.py > /tmp/v3.py && python3 /tmp/v3.py; echo "--- import check (no reforacle) ---"; grep -n "reforacle\|import" /app/objectives.py; echo "--- compiles standalone ---"; python3 -c "import ast; ast.parse(open('/app/objectives.py').read()); print('syntax OK')"
output
FAILS 0 {}
--- import check (no reforacle) ---
1:from __future__ import annotations
3:import torch
--- compiles standalone ---
syntax OK

[stdout]
FAILS 0 {}
--- import check (no reforacle) ---
1:from __future__ import annotations
3:import torch
--- compiles standalone ---
syntax OK
All 37 functions match `reforacle` exactly , even under the tight primitive tolerance (`atol=1e-6, rtol=1e-5`) and across a broad randomized stress test covering the tricky edge cases. The file depends only on `torch` (no `reforacle` reference at runtime). ## Summary I reverse-engineered each function by probing `reforacle` and reimplemented all 37 in `/app/objectives.py`. Key behaviors I recovered that differ from naive assumptions: - **`temperature_scale`** , special-cases `T==0` to a one-hot logit vector (argmax→`0`, rest `-inf`); otherwise `logits/T`. - **`repetition_penalty`** , *count-based* (not present/absent): factor is `penalty**count`, applied as divide if logit>0 else multiply. `frequency_penalty` subtracts `penalty*count`; `presence_penalty` subtracts `penalty` once if present. - **`top_p`** , HF nucleus with right-shift (keeps the boundary token, ≥1 token). **`min_p`** , keeps `prob ≥ p·max_prob` with **no** floor (all `-inf` when `p>1`). **`epsilon_filter`** , keeps `prob ≥ ε` but **floors at top-1**. **`top_k`** clamps `k` to vocab size. - **`position_ids`** , `clamp(cumsum(mask)-1, min=0)`. **`last_token_indices`** , last index where mask>0 (`-1` if none). **`gather_last_logits`/`rollout_last_logits`** clamp that index to `≥0`. **`masked_argmax`** returns `-1` for fully-masked rows. - **`mask_after_eos`/`generation_stop_mask`** , first EOS inclusive; the latter forces `stop_pos = max(eos_idx, min_len-1)`. - **`prepare_labels`** , shift-left labels with first `prompt_length-1` positions masked; **`completion_loss_labels`** , no shift, mask first `prompt_length`. - **`reward_to_go`** (masked reverse cumsum), **`discount_cumsum`** (γ-discounted reverse), **`build_position_ids_packed`** (per-segment aranges, zero-padded to `total_len`), **`right_pad_to`** (truncates when width is smaller). Verified with two independent randomized harnesses (90 seeds total) exercising non-zero/non-contiguous pads, all-pad rows, ties, extreme temperatures, and full `k`/`p` sweeps , 0 mismatches.
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: 31/37  (partial_score=0.8378)
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_e3ee55f38a8f43db. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.

Trial trial_e3ee55f38a8f43db · verifier authoritative; classifier explanatory.