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 Understanding
EvidenceTest output shows: 'FUNCTIONS PASSED: 35/37 (partial_score=0.9459) REWARD = 0'. Agent completed systematic reverse-engineering via probing, achieving 100% pass rate on multiple local test runs ('0 failing functions out of 37' across 4 test runs, seeds 123, 999 with 120 trials). However, 2 functions failed on grading's held-out inputs. The instruction explicitly states: 'Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on. Grading compares your implementation to reforacle on held-out inputs you do not see.' The agent successfully implemented 35/37 functions with exact behavior matching on all inputs they probed, but missed edge cases or sub-conditions in 2 functions.
Root causeThe task requires reverse-engineering 37 functions with exact numerical agreement on hidden test inputs. The agent's probing-based approach was sound and uncovered sophisticated non-obvious conventions (last-max tie-breaking, stable sorting, complex penalty calculations), but probing is inherently limited, undiscovered edge cases in 2 functions caused failures on held-out inputs outside the agent's test coverage.
RecommendationN/A - task is fine. The agent's failure is expected for this difficulty level. Pass rate of 94.6% (35/37) on a hard reverse-engineering task with no specification is reasonable performance. The task is well-specified and the grading is fair, the agent simply couldn't discover all edge cases through probing alone.
Trajectory
Tool-by-tool agent trajectory
46 tool calls · 4 tool types · 48 steps
Reimplement the post-training decoding-and-collation library in `/app/objectives.py` so that every function reproduces the sealed reference module `reforacle` exactly. `reforacle` is importable and callable (`import reforacle`) but its source cannot be read. It is the only ground truth; there is no written specification of what any function computes or of the conventions it follows. Recover its behavior by calling it. `/app/objectives.py` contains thirty-seven function stubs , logit processors and sampling filters, sequence collation and masking utilities, and rollout helpers used to generate and batch data for RLHF and SFT. Keep every signature exactly as given and implement every function. The names indicate the standard quantity each is meant to produce, but the reference's exact definition is whatever `reforacle` returns , not what you assume. Determine each definition by probing. Grading compares your implementation to `reforacle` on held-out inputs you do not see, function by function, requiring exact numerical agreement (primitives within `atol=1e-6, rtol=1e-5`; composite functions within `atol=1e-5, rtol=1e-4`). Each function passes only if it agrees on every graded input, and the task passes only if all thirty-seven functions pass. Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on. Grading runs your `/app/objectives.py` in a clean sandbox that does not contain `reforacle`, so any submission that imports, wraps, calls, loads, or otherwise depends on the reference at run time fails outright. Your file must be a standalone implementation whose only dependency is `torch`. Do not read, decompile, disassemble, or otherwise extract the compiled reference; recover its behavior only by calling it and comparing. `python3 /app/probe.py` shows the calling pattern; `bash /app/timer.sh` reports remaining budget.

/app/objectives.py

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

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

Remaining time (hours:minutes):
1:59

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

/app/probe.py

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

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

Remaining time (hours:minutes):
1:59

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

Check remaining budget

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

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

Remaining time (hours:minutes):
1:59

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

List reforacle attributes

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

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

Run probe example

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

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

Check constants

shell
python3 -c "import reforacle; print('NEG_INF', reforacle.NEG_INF); print('F', 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 greedy and temperature_scale

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,6,generator=g)
print('logits', logits.tolist())
# temperature_scale
print('T=2', R.temperature_scale(logits,2.0).tolist())
print('T=0.5', R.temperature_scale(logits,0.5).tolist())
print('T=1', R.temperature_scale(logits,1.0).tolist())
# greedy shape/dtype
gt=R.greedy_token(logits); print('greedy', gt.tolist(), gt.dtype, gt.shape)
# 1D logits greedy
l1=torch.randn(6,generator=g); print('l1',l1.tolist()); 
try:
  print('greedy1d', R.greedy_token(l1).tolist(), R.greedy_token(l1).shape)
except Exception as e: print('greedy1d err',e)
try:
  print('tscale1d', R.temperature_scale(l1,2.0).tolist())
except Exception as e: print('tscale1d err', e)
"
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]]
T=2 [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
T=0.5 [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854, -2.7971909046173096], [0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828, -1.1932706832885742, 0.3640729784965515]]
T=1 [[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 [0, 1] torch.int64 torch.Size([2])
l1 [-0.8566746115684509, 1.1006041765213013, -1.0711873769760132, 0.1227012425661087, -0.5663174986839294, 0.3731146454811096]
greedy1d 1 torch.Size([])
tscale1d [-0.42833730578422546, 0.5503020882606506, -0.5355936884880066, 0.06135062128305435, -0.2831587493419647, 0.1865573227405548]

[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]]
T=2 [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
T=0.5 [[3.0819921493530273, -0.586857795715332, -4.357578754425049, 1.1368625164031982, -2.1690447330474854, -2.7971909046173096], [0.8066936731338501, 1.6760526895523071, -1.4385151863098145, -0.80668705701828, -1.1932706832885742, 0.3640729784965515]]
T=1 [[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 [0, 1] torch.int64 torch.Size([2])
l1 [-0.8566746115684509, 1.1006041765213013, -1.0711873769760132, 0.1227012425661087, -0.5663174986839294, 0.3731146454811096]
greedy1d 1 torch.Size([])
tscale1d [-0.42833730578422546, 0.5503020882606506, -0.5355936884880066, 0.06135062128305435, -0.2831587493419647, 0.1865573227405548]

Probe top_k and top_p

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,6,generator=g)
print('logits', logits.tolist())
print('top_k k=2', R.top_k_filter(logits,2).tolist())
print('top_k k=1', R.top_k_filter(logits,1).tolist())
print('top_k k=10', R.top_k_filter(logits,10).tolist())
# top_p
print('top_p 0.9', R.top_p_filter(logits,0.9).tolist())
print('top_p 0.5', R.top_p_filter(logits,0.5).tolist())
print('top_p 0.0', R.top_p_filter(logits,0.0).tolist())
print('top_p 1.0', R.top_p_filter(logits,1.0).tolist())
"
output
logits [[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_k k=2 [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k k=1 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k k=10 [[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.9 [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p 0.5 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_p 0.0 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_p 1.0 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]

[stdout]
logits [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_k k=2 [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k k=1 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_k k=10 [[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.9 [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
top_p 0.5 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_p 0.0 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]]
top_p 1.0 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]

Probe min_p, epsilon, renormalize, suppress

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,6,generator=g)
print('min_p 0.5', R.min_p_filter(logits,0.5).tolist())
print('min_p 0.1', R.min_p_filter(logits,0.1).tolist())
print('min_p 0.0', R.min_p_filter(logits,0.0).tolist())
print('eps 0.1', R.epsilon_filter(logits,0.1).tolist())
print('eps 0.05', R.epsilon_filter(logits,0.05).tolist())
print('renorm', R.renormalize_logits(logits).tolist())
print('suppress [0,2]', R.suppress_tokens(logits,[0,2]).tolist())
print('suppress tensor', R.suppress_tokens(logits,torch.tensor([1,3])).tolist())
"
output
min_p 0.5 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
min_p 0.1 [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
min_p 0.0 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
eps 0.1 [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
eps 0.05 [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
renorm [[-0.5231569409370422, -2.357581853866577, -4.2429423332214355, -1.4957218170166016, -3.1486752033233643, -3.4627482891082764], [-1.5008691549301147, -1.0661896467208862, -2.623473644256592, -2.3075594902038574, -2.5008511543273926, -1.7221795320510864]]
suppress [0,2] [[-inf, -0.293428897857666, -inf, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [-inf, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
suppress tensor [[1.5409960746765137, -inf, -2.1787893772125244, -inf, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, -inf, -0.7192575931549072, -inf, -0.5966353416442871, 0.18203648924827576]]

[stdout]
min_p 0.5 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
min_p 0.1 [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
min_p 0.0 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
eps 0.1 [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]]
eps 0.05 [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
renorm [[-0.5231569409370422, -2.357581853866577, -4.2429423332214355, -1.4957218170166016, -3.1486752033233643, -3.4627482891082764], [-1.5008691549301147, -1.0661896467208862, -2.623473644256592, -2.3075594902038574, -2.5008511543273926, -1.7221795320510864]]
suppress [0,2] [[-inf, -0.293428897857666, -inf, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [-inf, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
suppress tensor [[1.5409960746765137, -inf, -2.1787893772125244, -inf, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, -inf, -0.7192575931549072, -inf, -0.5966353416442871, 0.18203648924827576]]

Probe penalties and min_length

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,6,generator=g)
ids=torch.tensor([[0,0,1,3],[2,2,2,5]])
print('rep 1.5', R.repetition_penalty(logits,ids,1.5).tolist())
print('freq 0.5', R.frequency_penalty(logits,ids,0.5).tolist())
print('pres 0.5', R.presence_penalty(logits,ids,0.5).tolist())
print('minlen cur=2 min=5 eos=3', R.min_length_filter(logits,2,5,3).tolist())
print('minlen cur=5 min=5 eos=3', R.min_length_filter(logits,5,5,3).tolist())
print('minlen cur=6 min=5 eos=3', R.min_length_filter(logits,6,5,3).tolist())
"
output
rep 1.5 [[0.684887170791626, -0.440143346786499, -2.1787893772125244, 0.3789541721343994, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -2.4274942874908447, -0.40334352850914, -0.5966353416442871, 0.12135765701532364]]
freq 0.5 [[0.5409960746765137, -0.793428897857666, -2.1787893772125244, 0.06843125820159912, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -2.2192575931549072, -0.40334352850914, -0.5966353416442871, -0.31796351075172424]]
pres 0.5 [[1.0409960746765137, -0.793428897857666, -2.1787893772125244, 0.06843125820159912, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -1.2192575931549072, -0.40334352850914, -0.5966353416442871, -0.31796351075172424]]
minlen cur=2 min=5 eos=3 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, -inf, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -inf, -0.5966353416442871, 0.18203648924827576]]
minlen cur=5 min=5 eos=3 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
minlen cur=6 min=5 eos=3 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]

[stdout]
rep 1.5 [[0.684887170791626, -0.440143346786499, -2.1787893772125244, 0.3789541721343994, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -2.4274942874908447, -0.40334352850914, -0.5966353416442871, 0.12135765701532364]]
freq 0.5 [[0.5409960746765137, -0.793428897857666, -2.1787893772125244, 0.06843125820159912, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -2.2192575931549072, -0.40334352850914, -0.5966353416442871, -0.31796351075172424]]
pres 0.5 [[1.0409960746765137, -0.793428897857666, -2.1787893772125244, 0.06843125820159912, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -1.2192575931549072, -0.40334352850914, -0.5966353416442871, -0.31796351075172424]]
minlen cur=2 min=5 eos=3 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, -inf, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -inf, -0.5966353416442871, 0.18203648924827576]]
minlen cur=5 min=5 eos=3 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
minlen cur=6 min=5 eos=3 [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]

Probe entropy, action_logprobs, masked_argmax, topk_indices

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,6,generator=g)
print('entropy', R.entropy_from_logits(logits).tolist())
acts=torch.tensor([0,3])
print('action_logprobs', R.action_logprobs(logits,acts).tolist())
# 3D action logprobs
lg3=torch.randn(2,3,6,generator=g); acts3=torch.tensor([[0,1,2],[3,4,5]])
print('alp3', R.action_logprobs(lg3,acts3).tolist(), R.action_logprobs(lg3,acts3).shape)
print('ent3', R.entropy_from_logits(lg3).tolist(), R.entropy_from_logits(lg3).shape)
mask=torch.tensor([[1,1,0,1,0,0],[0,1,1,0,1,1]])
print('masked_argmax', R.masked_argmax(logits,mask.bool()).tolist())
print('topk_indices k=3', R.topk_indices(logits,3).tolist(), R.topk_indices(logits,3).dtype)
"
output
entropy [1.1729530096054077, 1.6344718933105469]
action_logprobs [-0.5231569409370422, -2.3075594902038574]
alp3 [[-1.9803643226623535, -2.902956247329712, -2.3063485622406006], [-0.922078013420105, -3.8747358322143555, -2.4907402992248535]] torch.Size([2, 3])
ent3 [[1.6387826204299927, 1.4712406396865845, 1.521892786026001], [1.602628469467163, 1.4939178228378296, 1.3839476108551025]] torch.Size([2, 3])
masked_argmax [0, 1]
topk_indices k=3 [[0, 3, 1], [1, 0, 5]] torch.int64

[stdout]
entropy [1.1729530096054077, 1.6344718933105469]
action_logprobs [-0.5231569409370422, -2.3075594902038574]
alp3 [[-1.9803643226623535, -2.902956247329712, -2.3063485622406006], [-0.922078013420105, -3.8747358322143555, -2.4907402992248535]] torch.Size([2, 3])
ent3 [[1.6387826204299927, 1.4712406396865845, 1.521892786026001], [1.602628469467163, 1.4939178228378296, 1.3839476108551025]] torch.Size([2, 3])
masked_argmax [0, 1]
topk_indices k=3 [[0, 3, 1], [1, 0, 5]] torch.int64

Probe mask utils

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

[stdout]
attn pad0 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 0.0]] torch.float32
pos pad0 [[0, 1, 2, 2, 2], [0, 1, 2, 3, 3]] torch.int64
seqlen [3, 4] torch.int64
last_token_indices [2, 3] torch.int64
causal_mask 4 [[0.0, -inf, -inf, -inf], [0.0, 0.0, -inf, -inf], [0.0, 0.0, 0.0, -inf], [0.0, 0.0, 0.0, 0.0]] torch.float32

Probe edge cases

shell
python3 -c "
import torch, reforacle as R
# non-contiguous pad
ids=torch.tensor([[5,0,7,0,9]])
print('attn', R.attention_mask(ids,0).tolist())
print('pos', R.position_ids(ids,0).tolist())
print('seqlen', R.sequence_lengths(ids,0).tolist())
am=R.attention_mask(ids,0)
print('lasttok', R.last_token_indices(am).tolist())
# last_token_indices with float mask having middle zeros
m=torch.tensor([[1.,1,0,1,0]])
print('lasttok2', R.last_token_indices(m).tolist())
# shift_labels
ids2=torch.tensor([[5,6,7,8]])
print('shift', R.shift_labels(ids2,-100).tolist(), R.shift_labels(ids2,-100).shape)
"
output
attn [[1.0, 0.0, 1.0, 0.0, 1.0]]
pos [[0, 0, 1, 1, 2]]
seqlen [3]
lasttok [4]
lasttok2 [3]
shift [[6, 7, 8, -100]] torch.Size([1, 4])

[stdout]
attn [[1.0, 0.0, 1.0, 0.0, 1.0]]
pos [[0, 0, 1, 1, 2]]
seqlen [3]
lasttok [4]
lasttok2 [3]
shift [[6, 7, 8, -100]] torch.Size([1, 4])

Probe label and eos functions

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

[stdout]
completion_loss_labels [[-100, -100, 7, 8, 9], [-100, -100, -100, 4, 0]]
prepare_labels [[-100, 7, 8, 9, -100], [-100, -100, 4, 0, -100]]
mask_after_eos eos=2 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] torch.float32

Probe gather/truncate/pad

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,5,4,generator=g)
am=torch.tensor([[1.,1,1,0,0],[1,1,1,1,0]])
gl=R.gather_last_logits(logits,am)
print('gather_last_logits', gl.tolist(), gl.shape)
ids=torch.tensor([[5,6,7,8,9],[1,2,3,4,0]])
print('truncate_left ml=3', R.truncate_left(ids,3).tolist())
print('truncate_left ml=10', R.truncate_left(ids,10).tolist())
print('pad_to_multiple pad0 m=4', R.pad_to_multiple(ids,0,4).tolist())
print('pad_to_multiple pad0 m=3', R.pad_to_multiple(ids,0,3).tolist())
print('right_pad_to pad0 w=7', R.right_pad_to(ids,0,7).tolist())
print('right_pad_to pad0 w=4', R.right_pad_to(ids,0,4).tolist())
"
output
gather_last_logits [[0.32227492332458496, -1.2633347511291504, 0.34998318552970886, 0.30813390016555786], [1.2032238245010376, 0.08453487604856491, -1.2001395225524902, -0.004785861354321241]] torch.Size([2, 4])
truncate_left ml=3 [[7, 8, 9], [3, 4, 0]]
truncate_left ml=10 [[5, 6, 7, 8, 9], [1, 2, 3, 4, 0]]
pad_to_multiple pad0 m=4 [[5, 6, 7, 8, 9, 0, 0, 0], [1, 2, 3, 4, 0, 0, 0, 0]]
pad_to_multiple pad0 m=3 [[5, 6, 7, 8, 9, 0], [1, 2, 3, 4, 0, 0]]
right_pad_to pad0 w=7 [[5, 6, 7, 8, 9, 0, 0], [1, 2, 3, 4, 0, 0, 0]]
right_pad_to pad0 w=4 [[5, 6, 7, 8], [1, 2, 3, 4]]

[stdout]
gather_last_logits [[0.32227492332458496, -1.2633347511291504, 0.34998318552970886, 0.30813390016555786], [1.2032238245010376, 0.08453487604856491, -1.2001395225524902, -0.004785861354321241]] torch.Size([2, 4])
truncate_left ml=3 [[7, 8, 9], [3, 4, 0]]
truncate_left ml=10 [[5, 6, 7, 8, 9], [1, 2, 3, 4, 0]]
pad_to_multiple pad0 m=4 [[5, 6, 7, 8, 9, 0, 0, 0], [1, 2, 3, 4, 0, 0, 0, 0]]
pad_to_multiple pad0 m=3 [[5, 6, 7, 8, 9, 0], [1, 2, 3, 4, 0, 0]]
right_pad_to pad0 w=7 [[5, 6, 7, 8, 9, 0, 0], [1, 2, 3, 4, 0, 0, 0]]
right_pad_to pad0 w=4 [[5, 6, 7, 8], [1, 2, 3, 4]]

Probe rollout, first_stop, gen_stop

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,5,4,generator=g)
ids=torch.tensor([[5,6,7,0,0],[1,2,3,4,0]])
rl=R.rollout_last_logits(logits,ids,0)
print('rollout_last_logits', rl.tolist(), rl.shape)
sids=torch.tensor([2,4])
seq=torch.tensor([[5,2,7,4,9],[1,2,3,4,2]])
print('first_stop_index', R.first_stop_index(seq,sids).tolist(), R.first_stop_index(seq,sids).dtype)
seq2=torch.tensor([[5,6,7],[1,2,3]])
print('first_stop_index none', R.first_stop_index(seq2,sids).tolist())
print('gen_stop_mask eos=2 minlen=2', R.generation_stop_mask(seq,2,2).tolist(), R.generation_stop_mask(seq,2,2).dtype)
"
output
rollout_last_logits [[0.32227492332458496, -1.2633347511291504, 0.34998318552970886, 0.30813390016555786], [1.2032238245010376, 0.08453487604856491, -1.2001395225524902, -0.004785861354321241]] torch.Size([2, 4])
first_stop_index [1, 1] torch.int64
first_stop_index none [-1, 1]
gen_stop_mask eos=2 minlen=2 [[1.0, 1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] torch.float32

[stdout]
rollout_last_logits [[0.32227492332458496, -1.2633347511291504, 0.34998318552970886, 0.30813390016555786], [1.2032238245010376, 0.08453487604856491, -1.2001395225524902, -0.004785861354321241]] torch.Size([2, 4])
first_stop_index [1, 1] torch.int64
first_stop_index none [-1, 1]
gen_stop_mask eos=2 minlen=2 [[1.0, 1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] torch.float32

Probe generation_stop_mask min_len effect

shell
python3 -c "
import torch, reforacle as R
# eos before min_len
seq=torch.tensor([[2,5,6,2,9],[5,6,7,8,2]])
print('a eos=2 min=3', R.generation_stop_mask(seq,2,3).tolist())
print('b eos=2 min=1', R.generation_stop_mask(seq,2,1).tolist())
print('c eos=2 min=0', R.generation_stop_mask(seq,2,0).tolist())
seq2=torch.tensor([[2,2,2,2,2]])
print('d eos=2 min=2', R.generation_stop_mask(seq2,2,2).tolist())
seq3=torch.tensor([[5,6,7,8,9]])
print('e no eos min=2', R.generation_stop_mask(seq3,2,2).tolist())
"
output
a eos=2 min=3 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
b eos=2 min=1 [[1.0, 0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
c eos=2 min=0 [[1.0, 0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
d eos=2 min=2 [[1.0, 1.0, 0.0, 0.0, 0.0]]
e no eos min=2 [[1.0, 1.0, 1.0, 1.0, 1.0]]

[stdout]
a eos=2 min=3 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
b eos=2 min=1 [[1.0, 0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
c eos=2 min=0 [[1.0, 0.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]]
d eos=2 min=2 [[1.0, 1.0, 0.0, 0.0, 0.0]]
e no eos min=2 [[1.0, 1.0, 1.0, 1.0, 1.0]]

Probe packed pos, reward_to_go, discount_cumsum

shell
python3 -c "
import torch, reforacle as R
sl=torch.tensor([2,3,1])
print('build_pos_packed total=6', R.build_position_ids_packed(sl,6).tolist(), R.build_position_ids_packed(sl,6).dtype)
print('build_pos_packed total=8', R.build_position_ids_packed(sl,8).tolist())
r=torch.tensor([[1.,2,3,4],[5,6,7,8]])
m=torch.tensor([[1.,1,1,0],[1,1,0,0]])
print('reward_to_go', R.reward_to_go(r,m).tolist())
x=torch.tensor([[1.,1,1,1],[1,2,3,4]])
print('discount_cumsum 0.9', R.discount_cumsum(x,0.9).tolist())
print('discount_cumsum 1d', R.discount_cumsum(torch.tensor([1.,2,3]),0.5).tolist())
"
output
build_pos_packed total=6 [0, 1, 0, 1, 2, 0] torch.int64
build_pos_packed total=8 [0, 1, 0, 1, 2, 0, 0, 0]
reward_to_go [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]]
discount_cumsum 0.9 [[3.438999891281128, 2.7100000381469727, 1.899999976158142, 1.0], [8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]]
discount_cumsum 1d [2.75, 3.5, 3.0]

[stdout]
build_pos_packed total=6 [0, 1, 0, 1, 2, 0] torch.int64
build_pos_packed total=8 [0, 1, 0, 1, 2, 0, 0, 0]
reward_to_go [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]]
discount_cumsum 0.9 [[3.438999891281128, 2.7100000381469727, 1.899999976158142, 1.0], [8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]]
discount_cumsum 1d [2.75, 3.5, 3.0]

Test reward_to_go non-contiguous mask

shell
python3 -c "
import torch, reforacle as R
r=torch.tensor([[1.,2,3,4,5]])
m=torch.tensor([[1.,0,1,0,1]])
print('rtg noncontig', R.reward_to_go(r,m).tolist())
# manual: rm=[1,0,3,0,5], revcumsum=[9,8,8,5,5], *mask=[9,0,8,0,5]
"
output
rtg noncontig [[9.0, 0.0, 8.0, 0.0, 5.0]]

[stdout]
rtg noncontig [[9.0, 0.0, 8.0, 0.0, 5.0]]

Probe logits_warper and sample_greedy_warped

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,6,generator=g)
print('logits', logits.tolist())
print('warper T=2 k=3 p=0.9', R.logits_warper(logits,2.0,3,0.9).tolist())
print('warper T=1 k=0 p=1.0', R.logits_warper(logits,1.0,0,1.0).tolist())
print('warper T=1 k=6 p=1.0', R.logits_warper(logits,1.0,6,1.0).tolist())
print('sample_greedy_warped T=2 k=3 p=0.9', R.sample_greedy_warped(logits,2.0,3,0.9).tolist())
"
output
Exit code 1
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]]
warper T=2 k=3 p=0.9 [[0.7704980373382568, -0.146714448928833, -inf, 0.28421562910079956, -inf, -inf], [0.20167341828346252, 0.4190131723880768, -inf, -inf, -inf, 0.09101824462413788]]
Traceback (most recent call last):
  File "<string>", line 7, in <module>
  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 sample_greedy_warped

shell
python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,6,generator=g)
print('sgw T=2 k=3 p=0.9', R.sample_greedy_warped(logits,2.0,3,0.9).tolist(), R.sample_greedy_warped(logits,2.0,3,0.9).dtype, R.sample_greedy_warped(logits,2.0,3,0.9).shape)
print('sgw T=0.5 k=6 p=1.0', R.sample_greedy_warped(logits,0.5,6,1.0).tolist())
# does warper order T,topk,topp? Test with k=6 (no effect) p=0.5
print('warper T=1 k=6 p=0.5', R.logits_warper(logits,1.0,6,0.5).tolist())
print('topp only 0.5', R.top_p_filter(logits,0.5).tolist())
"
output
sgw T=2 k=3 p=0.9 [0, 1] torch.int64 torch.Size([2])
sgw T=0.5 k=6 p=1.0 [0, 1]
warper T=1 k=6 p=0.5 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
topp only 0.5 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]

[stdout]
sgw T=2 k=3 p=0.9 [0, 1] torch.int64 torch.Size([2])
sgw T=0.5 k=6 p=1.0 [0, 1]
warper T=1 k=6 p=0.5 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
topp only 0.5 [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]]
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:
        idx = logits.argmax(dim=-1, keepdim=True)
        out = torch.full_like(logits, NEG_INF)
        out.scatter_(-1, idx, 0.0)
        return out
    return logits / temperature


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


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


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


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):
    out = logits.clone()
    out[..., token_ids] = NEG_INF
    return out


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


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


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


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


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


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


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


def mask_after_eos(input_ids, eos_id):
    is_eos = (input_ids == eos_id).long()
    excl = is_eos.cumsum(dim=-1) - is_eos
    return (excl == 0).float()


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


def last_token_indices(attn_mask):
    L = attn_mask.size(-1)
    return L - 1 - attn_mask.flip(-1).argmax(dim=-1)


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


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


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


def action_logprobs(logits, actions):
    logp = 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):
    rm = rewards * mask
    rtg = rm.flip(-1).cumsum(dim=-1).flip(-1)
    return rtg * mask


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


def build_position_ids_packed(seq_lengths, total_len):
    parts = [torch.arange(int(l), device=seq_lengths.device) for l in seq_lengths]
    out = torch.zeros(total_len, dtype=torch.long, device=seq_lengths.device)
    if parts:
        cat = torch.cat(parts)
        out[: cat.size(0)] = cat
    return out


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


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


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


def rollout_last_logits(logits, input_ids, pad_id):
    return gather_last_logits(logits, attention_mask(input_ids, pad_id))


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):
    if not torch.is_tensor(stop_ids):
        stop_ids = torch.tensor(stop_ids, device=input_ids.device)
    is_stop = (input_ids.unsqueeze(-1) == stop_ids).any(dim=-1)
    has = is_stop.any(dim=-1)
    idx = is_stop.int().argmax(dim=-1)
    return torch.where(has, idx, torch.full_like(idx, -1))


def masked_argmax(logits, mask):
    return logits.masked_fill(~mask.bool(), NEG_INF).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):
    L = input_ids.size(-1)
    if L >= width:
        return input_ids[..., :width]
    return torch.nn.functional.pad(input_ids, (0, width - L), value=pad_id)


def generation_stop_mask(input_ids, eos_id, min_len):
    L = input_ids.size(-1)
    is_eos = (input_ids == eos_id)
    has = is_eos.any(dim=-1)
    first = is_eos.int().argmax(dim=-1)
    length = torch.where(has, first + 1, torch.full_like(first, L))
    length = length.clamp(min=min_len, max=L)
    arange = torch.arange(L, device=input_ids.device)
    return (arange.unsqueeze(0) < length.unsqueeze(-1)).float()
write · /app/test_all.py
import torch
import reforacle as R
import objectives as O

torch.manual_seed(123)

def eq(a, b, atol=1e-6, rtol=1e-5):
    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
        if torch.equal(af.isinf(), bf.isinf()) and torch.equal(af.isnan(), bf.isnan()):
            fin = ~af.isinf() & ~af.isnan()
            if fin.sum()==0: return None
            if torch.allclose(af[fin], bf[fin], atol=atol, rtol=rtol):
                # check inf signs
                if torch.equal(af.masked_fill(~af.isinf(),0), bf.masked_fill(~bf.isinf(),0)):
                    return None
            return f"VAL max diff {(af[fin]-bf[fin]).abs().max().item()}"
        return "INF/NAN mismatch"
    return None if a==b else f"{a} vs {b}"

results = {}
def check(name, mine, ref, **kw):
    r = eq(mine, ref, **kw)
    results.setdefault(name, [])
    if r is not None:
        results[name].append(r)

for trial in range(50):
    B = torch.randint(1,4,(1,)).item()
    V = torch.randint(2,8,(1,)).item()
    Lseq = torch.randint(2,7,(1,)).item()
    logits = torch.randn(B,V)
    check("greedy_token", O.greedy_token(logits), R.greedy_token(logits))
    for T in [0.0,0.5,1.0,2.0,0.7]:
        check("temperature_scale", O.temperature_scale(logits,T), R.temperature_scale(logits,T))
    for k in [1,2,V]:
        check("top_k_filter", O.top_k_filter(logits,k), R.top_k_filter(logits,k))
        check("topk_indices", O.topk_indices(logits,k), R.topk_indices(logits,k))
    for p in [0.0,0.3,0.5,0.9,1.0]:
        check("top_p_filter", O.top_p_filter(logits,p), R.top_p_filter(logits,p))
        check("min_p_filter", O.min_p_filter(logits,p), R.min_p_filter(logits,p))
    for e in [0.01,0.1,0.3]:
        check("epsilon_filter", O.epsilon_filter(logits,e), R.epsilon_filter(logits,e))
    ids = torch.randint(0,V,(B,Lseq))
    for pen in [1.0,1.2,1.5,2.0]:
        check("repetition_penalty", O.repetition_penalty(logits,ids,pen), R.repetition_penalty(logits,ids,pen))
    for pen in [0.0,0.5,1.0]:
        check("frequency_penalty", O.frequency_penalty(logits,ids,pen), R.frequency_penalty(logits,ids,pen))
        check("presence_penalty", O.presence_penalty(logits,ids,pen), R.presence_penalty(logits,ids,pen))
    tk = torch.randint(0,V,(torch.randint(1,3,(1,)).item(),)).unique()
    check("suppress_tokens", O.suppress_tokens(logits,tk), R.suppress_tokens(logits,tk))
    check("suppress_tokens", O.suppress_tokens(logits,tk.tolist()), R.suppress_tokens(logits,tk.tolist()))
    for cur in [0,2,5]:
        check("min_length_filter", O.min_length_filter(logits,cur,3,min(V-1,1)), R.min_length_filter(logits,cur,3,min(V-1,1)))
    check("renormalize_logits", O.renormalize_logits(logits), R.renormalize_logits(logits))
    check("entropy_from_logits", O.entropy_from_logits(logits), R.entropy_from_logits(logits))
    # 3d logits
    l3 = torch.randn(B,Lseq,V)
    a3 = torch.randint(0,V,(B,Lseq))
    check("action_logprobs", O.action_logprobs(l3,a3), R.action_logprobs(l3,a3))
    check("entropy_from_logits", O.entropy_from_logits(l3), R.entropy_from_logits(l3))
    check("action_logprobs", O.action_logprobs(logits,torch.randint(0,V,(B,))), R.action_logprobs(logits,torch.randint(0,V,(B,))))

    # pad-based
    pad = 0
    pids = torch.randint(0,V,(B,Lseq))
    check("attention_mask", O.attention_mask(pids,pad), R.attention_mask(pids,pad))
    check("position_ids", O.position_ids(pids,pad), R.position_ids(pids,pad))
    check("sequence_lengths", O.sequence_lengths(pids,pad), R.sequence_lengths(pids,pad))
    am = R.attention_mask(pids,pad)
    if (am.sum(-1)>0).all():
        check("last_token_indices", O.last_token_indices(am), R.last_token_indices(am))
        check("gather_last_logits", O.gather_last_logits(l3,am), R.gather_last_logits(l3,am))
        check("rollout_last_logits", O.rollout_last_logits(l3,pids,pad), R.rollout_last_logits(l3,pids,pad))
    check("causal_mask", O.causal_mask(Lseq), R.causal_mask(Lseq))
    check("shift_labels", O.shift_labels(pids,-100), R.shift_labels(pids,-100))
    pl = torch.randint(1,Lseq+1,(B,))
    check("completion_loss_labels", O.completion_loss_labels(pids,pl,-100), R.completion_loss_labels(pids,pl,-100))
    check("prepare_labels", O.prepare_labels(pids,pl,-100), R.prepare_labels(pids,pl,-100))
    check("mask_after_eos", O.mask_after_eos(pids,1), R.mask_after_eos(pids,1))
    for ml in [1,3,10]:
        check("truncate_left", O.truncate_left(pids,ml), R.truncate_left(pids,ml))
    for m in [2,3,4]:
        check("pad_to_multiple", O.pad_to_multiple(pids,pad,m), R.pad_to_multiple(pids,pad,m))
    for w in [1,Lseq,Lseq+3]:
        check("right_pad_to", O.right_pad_to(pids,pad,w), R.right_pad_to(pids,pad,w))

    # RL
    rew = torch.randn(B,Lseq)
    msk = (torch.rand(B,Lseq)>0.3).float()
    check("reward_to_go", O.reward_to_go(rew,msk), R.reward_to_go(rew,msk))
    for gm in [0.0,0.5,0.9,0.99,1.0]:
        check("discount_cumsum", O.discount_cumsum(rew,gm), R.discount_cumsum(rew,gm))
    check("discount_cumsum", O.discount_cumsum(rew[0],0.9), R.discount_cumsum(rew[0],0.9))
    sl = torch.randint(1,4,(torch.randint(1,4,(1,)).item(),))
    tot = int(sl.sum().item())+torch.randint(0,3,(1,)).item()
    check("build_position_ids_packed", O.build_position_ids_packed(sl,tot), R.build_position_ids_packed(sl,tot))

    # composite
    for T in [0.5,1.0,2.0]:
        for k in [1,2,V]:
            for p in [0.5,0.9,1.0]:
                check("logits_warper", O.logits_warper(logits,T,k,p), R.logits_warper(logits,T,k,p), atol=1e-5, rtol=1e-4)
                check("sample_greedy_warped", O.sample_greedy_warped(logits,T,k,p), R.sample_greedy_warped(logits,T,k,p))
    # masked_argmax
    mmask = (torch.rand(B,V)>0.3)
    for b in range(B):
        if mmask[b].sum()==0: mmask[b,0]=True
    check("masked_argmax", O.masked_argmax(logits,mmask), R.masked_argmax(logits,mmask))
    # first_stop / gen_stop
    sids = torch.randint(0,V,(2,)).unique()
    check("first_stop_index", O.first_stop_index(pids,sids), R.first_stop_index(pids,sids))
    for ml in [0,1,2,3,10]:
        check("generation_stop_mask", O.generation_stop_mask(pids,1,ml), R.generation_stop_mask(pids,1,ml))

allnames = ['greedy_token','temperature_scale','top_k_filter','top_p_filter','min_p_filter','repetition_penalty','frequency_penalty','presence_penalty','suppress_tokens','min_length_filter','renormalize_logits','attention_mask','position_ids','causal_mask','shift_labels','completion_loss_labels','mask_after_eos','sequence_lengths','last_token_indices','gather_last_logits','truncate_left','pad_to_multiple','action_logprobs','entropy_from_logits','reward_to_go','discount_cumsum','build_position_ids_packed','logits_warper','sample_greedy_warped','prepare_labels','rollout_last_logits','epsilon_filter','first_stop_index','masked_argmax','topk_indices','right_pad_to','generation_stop_mask']
fails = 0
for n in allnames:
    errs = results.get(n, None)
    if errs is None:
        print(f"?? {n}: NOT TESTED")
    elif errs:
        fails += 1
        print(f"FAIL {n}: {errs[:3]} ({len(errs)} fails)")
print("="*40)
print(f"{fails} failing functions out of {len(allnames)}")

Run comparison harness

shell
python3 test_all.py
output
FAIL action_logprobs: ['VAL max diff 2.1933975219726562', 'VAL max diff 2.415400266647339', 'VAL max diff 1.6721497774124146'] (41 fails)
FAIL epsilon_filter: ['INF/NAN mismatch'] (1 fails)
========================================
2 failing functions out of 37

[stdout]
FAIL action_logprobs: ['VAL max diff 2.1933975219726562', 'VAL max diff 2.415400266647339', 'VAL max diff 1.6721497774124146'] (41 fails)
FAIL epsilon_filter: ['INF/NAN mismatch'] (1 fails)
========================================
2 failing functions out of 37

Investigate action_logprobs

shell
python3 -c "
import torch, reforacle as R, objectives as O
torch.manual_seed(5)
l=torch.randn(2,5); a=torch.randint(0,5,(2,))
print('logits',l.tolist()); print('a',a.tolist())
print('ref', R.action_logprobs(l,a).tolist())
print('mine', O.action_logprobs(l,a).tolist())
print('logsoftmax gather', l.log_softmax(-1).gather(-1,a.unsqueeze(-1)).squeeze(-1).tolist())
print('softmax log gather', l.softmax(-1).log().gather(-1,a.unsqueeze(-1)).squeeze(-1).tolist())
# maybe raw logit gather?
print('raw gather', l.gather(-1,a.unsqueeze(-1)).squeeze(-1).tolist())
"
output
logits [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549], [1.9427834749221802, -1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096]]
a [4, 2]
ref [-1.70235276222229, -3.0629005432128906]
mine [-1.70235276222229, -3.0629005432128906]
logsoftmax gather [-1.70235276222229, -3.0629005432128906]
softmax log gather [-1.70235276222229, -3.0629007816314697]
raw gather [-0.19741509854793549, -0.7625571489334106]

[stdout]
logits [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549], [1.9427834749221802, -1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096]]
a [4, 2]
ref [-1.70235276222229, -3.0629005432128906]
mine [-1.70235276222229, -3.0629005432128906]
logsoftmax gather [-1.70235276222229, -3.0629005432128906]
softmax log gather [-1.70235276222229, -3.0629007816314697]
raw gather [-0.19741509854793549, -0.7625571489334106]

Check 3D action_logprobs

shell
python3 -c "
import torch, reforacle as R, objectives as O
torch.manual_seed(5)
l=torch.randn(2,3,5); a=torch.randint(0,5,(2,3))
print('ref', R.action_logprobs(l,a).tolist())
print('mine', O.action_logprobs(l,a).tolist())
"
output
ref [[-0.746162474155426, -2.9050049781799316, -2.2688262462615967], [-1.4956307411193848, -2.19138503074646, -3.3165507316589355]]
mine [[-0.746162474155426, -2.9050049781799316, -2.2688262462615967], [-1.4956307411193848, -2.19138503074646, -3.3165507316589355]]

[stdout]
ref [[-0.746162474155426, -2.9050049781799316, -2.2688262462615967], [-1.4956307411193848, -2.19138503074646, -3.3165507316589355]]
mine [[-0.746162474155426, -2.9050049781799316, -2.2688262462615967], [-1.4956307411193848, -2.19138503074646, -3.3165507316589355]]

Find epsilon_filter mismatch

shell
python3 -c "
import torch, reforacle as R, objectives as O
torch.manual_seed(123)
# find a mismatch for epsilon
for t in range(2000):
    B=torch.randint(1,4,(1,)).item(); V=torch.randint(2,8,(1,)).item()
    l=torch.randn(B,V)
    for e in [0.01,0.1,0.3]:
        r=R.epsilon_filter(l,e); m=O.epsilon_filter(l,e)
        if not torch.equal(r.isinf(),m.isinf()):
            print('MISMATCH e=',e,'l=',l.tolist())
            print('ref',r.tolist()); print('mine',m.tolist())
            p=l.softmax(-1); print('probs',p.tolist())
            import sys; sys.exit()
print('no mismatch found')
"
output
MISMATCH e= 0.3 l= [[0.32390275597572327, -0.10852263122797012, 0.21033115684986115, -0.39084282517433167, 0.23497341573238373], [0.6652604341506958, 0.3528207540512085, 0.9728211164474487, -0.03861621767282486, -0.8861014246940613]]
ref [[0.32390275597572327, -inf, -inf, -inf, -inf], [-inf, -inf, 0.9728211164474487, -inf, -inf]]
mine [[-inf, -inf, -inf, -inf, -inf], [-inf, -inf, 0.9728211164474487, -inf, -inf]]
probs [[0.2534340023994446, 0.16446176171302795, 0.22622540593147278, 0.12400946021080017, 0.23186935484409332], [0.26326969265937805, 0.19262383878231049, 0.35807397961616516, 0.13023002445697784, 0.05580241233110428]]

[stdout]
MISMATCH e= 0.3 l= [[0.32390275597572327, -0.10852263122797012, 0.21033115684986115, -0.39084282517433167, 0.23497341573238373], [0.6652604341506958, 0.3528207540512085, 0.9728211164474487, -0.03861621767282486, -0.8861014246940613]]
ref [[0.32390275597572327, -inf, -inf, -inf, -inf], [-inf, -inf, 0.9728211164474487, -inf, -inf]]
mine [[-inf, -inf, -inf, -inf, -inf], [-inf, -inf, 0.9728211164474487, -inf, -inf]]
probs [[0.2534340023994446, 0.16446176171302795, 0.22622540593147278, 0.12400946021080017, 0.23186935484409332], [0.26326969265937805, 0.19262383878231049, 0.35807397961616516, 0.13023002445697784, 0.05580241233110428]]

/app/objectives.py

− removed
def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    return logits.masked_fill(probs < epsilon, NEG_INF)
+ added
def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    top = torch.topk(logits, 1, dim=-1).values[..., -1, None]
    remove = (probs < epsilon) & (logits < top)
    return logits.masked_fill(remove, NEG_INF)

/app/test_all.py

− removed
    check("action_logprobs", O.action_logprobs(logits,torch.randint(0,V,(B,))), R.action_logprobs(logits,torch.randint(0,V,(B,))))
+ added
    a1 = torch.randint(0,V,(B,))
    check("action_logprobs", O.action_logprobs(logits,a1), R.action_logprobs(logits,a1))

Rerun harness

shell
python3 test_all.py
output
========================================
0 failing functions out of 37

[stdout]
========================================
0 failing functions out of 37

Stress test ties

shell
python3 -c "
import torch, reforacle as R, objectives as O
torch.manual_seed(7)
def cmp(name,a,b):
    if a.shape!=b.shape or a.dtype!=b.dtype: print('DIFF',name,'shape/dtype',a.shape,b.shape,a.dtype,b.dtype); return
    if not torch.equal(a.isinf(),b.isinf()): print('DIFF',name,'inf'); print(a.tolist()); print(b.tolist()); return
    fin=~a.isinf()
    if not torch.allclose(a[fin].float(),b[fin].float(),atol=1e-5,rtol=1e-4): print('DIFF',name,'val',(a[fin]-b[fin]).abs().max().item())
# TIES: integer logits
for t in range(300):
    l=torch.randint(-3,3,(3,6)).float()
    for k in [1,2,3,5]:
        cmp('topk',O.top_k_filter(l,k),R.top_k_filter(l,k))
        cmp('topkidx',O.topk_indices(l,k),R.topk_indices(l,k))
    for p in [0.0,0.2,0.5,0.7,0.9,1.0]:
        cmp('topp',O.top_p_filter(l,p),R.top_p_filter(l,p))
        cmp('minp',O.min_p_filter(l,p),R.min_p_filter(l,p))
    for e in [0.05,0.1,0.2,0.5]:
        cmp('eps',O.epsilon_filter(l,e),R.epsilon_filter(l,e))
    cmp('greedy',O.greedy_token(l),R.greedy_token(l))
    m=(torch.rand(3,6)>0.4)
    for b in range(3):
        if m[b].sum()==0: m[b,0]=True
    cmp('margmax',O.masked_argmax(l,m),R.masked_argmax(l,m))
    cmp('ent',O.entropy_from_logits(l),R.entropy_from_logits(l))
print('ties done')
"
output
<persisted-output>
Output too large (303.5KB). Full output saved to: /logs/agent/sessions/projects/-app/16b7c556-e986-4dcb-a12e-8e7ecbaeef2d/tool-results/bpyre6qi4.txt

Preview (first 2KB):
DIFF topkidx val 1
DIFF topkidx val 5
DIFF topkidx val 5
DIFF topkidx val 5
DIFF topp inf
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF topp inf
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF topp inf
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, 2.0, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 1.0]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [2.0, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 1.0]]
DIFF topp inf
[[0.0, 1.0, -inf, -1.0, -inf, 0.0], [2.0, 2.0, 1.0, 2.0, 1.0, -inf], [-inf, 0.0, -1.0, 1.0, -inf, 1.0]]
[[0.0, 1.0, -inf, -1.0, -inf, 0.0], [2.0, 2.0, 1.0, 2.0, 1.0, -inf], [-1.0, 0.0, -inf, 1.0, -inf, 1.0]]
DIFF eps inf
[[-inf, 1.0, -inf, -inf, -inf, -inf], [2.0, 2.0, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 1.0]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0]]
DIFF greedy val 3
DIFF margmax val 5
DIFF topkidx val 1
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, 2.0, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, 2.0, -inf]]
...
</persisted-output>

[stdout]
DIFF topkidx val 1
DIFF topkidx val 5
DIFF topkidx val 5
DIFF topkidx val 5
DIFF topp inf
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF topp inf
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF topp inf
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, 2.0, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 1.0]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [2.0, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 1.0]]
DIFF topp inf
[[0.0, 1.0, -inf, -1.0, -inf, 0.0], [2.0, 2.0, 1.0, 2.0, 1.0, -inf], [-inf, 0.0, -1.0, 1.0, -inf, 1.0]]
[[0.0, 1.0, -inf, -1.0, -inf, 0.0], [2.0, 2.0, 1.0, 2.0, 1.0, -inf], [-1.0, 0.0, -inf, 1.0, -inf, 1.0]]
DIFF eps inf
[[-inf, 1.0, -inf, -inf, -inf, -inf], [2.0, 2.0, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 1.0]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0]]
DIFF greedy val 3
DIFF margmax val 5
DIFF topkidx val 1
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, 2.0, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, 2.0, -inf]]
[[-inf, 2.0, 2.0, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, 2.0, -inf]]
DIFF topp inf
[[-inf, 2.0, 2.0, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -1.0], [-inf, -inf, -inf, 2.0, 2.0, -inf]]
[[-inf, 2.0, 2.0, -inf, -inf, 2.0], [-inf, -1.0, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, 2.0, -inf]]
DIFF topp inf
[[-inf, 2.0, 2.0, -inf, -inf, 2.0], [-inf, -1.0, -1.0, -inf, 1.0, -1.0], [-inf, 1.0, -inf, 2.0, 2.0, 0.0]]
[[-inf, 2.0, 2.0, -inf, -inf, 2.0], [-inf, -1.0, -1.0, -inf, 1.0, -1.0], [-inf, 1.0, 0.0, 2.0, 2.0, -inf]]
DIFF eps inf
[[-inf, 2.0, 2.0, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, 2.0, -inf]]
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
DIFF greedy val 4
DIFF margmax val 4
DIFF topkidx val 2
DIFF topkidx val 2
DIFF topkidx val 2
DIFF topkidx val 2
DIFF topp inf
[[-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 2.0, -inf, -inf]]
[[2.0, -inf, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 2.0, -inf, -inf]]
[[2.0, -inf, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 2.0, 2.0, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, 2.0, -inf, -inf]]
[[2.0, 2.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, 2.0, -inf, -inf]]
DIFF topp inf
[[2.0, 2.0, 2.0, 1.0, -inf, -inf], [2.0, 1.0, -inf, -inf, -inf, 2.0], [-inf, 0.0, 2.0, 2.0, 1.0, -inf]]
[[2.0, 2.0, 2.0, 1.0, -inf, -inf], [2.0, 1.0, -inf, -inf, -inf, 2.0], [0.0, -inf, 2.0, 2.0, 1.0, -inf]]
DIFF eps inf
[[2.0, 2.0, 2.0, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, 2.0, -inf, -inf]]
[[-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF greedy val 5
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topp inf
[[1.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[1.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[1.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[1.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[1.0, -inf, -inf, -inf, -inf, 0.0], [-inf, -inf, -inf, 2.0, -inf, 1.0], [-inf, -inf, 2.0, -inf, 2.0, -inf]]
[[1.0, 0.0, -inf, -inf, -inf, -inf], [-inf, -inf, 1.0, 2.0, -inf, -inf], [-inf, -inf, 2.0, -inf, 2.0, -inf]]
DIFF eps inf
[[1.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, 2.0, -inf, 2.0, -inf]]
[[1.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
DIFF greedy val 2
DIFF margmax val 2
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topkidx val 4
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, 1.0, 2.0, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, 2.0]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [1.0, -inf, -inf, 2.0, -inf, -inf], [-inf, 2.0, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[0.0, 2.0, -inf, -inf, 1.0, -inf], [1.0, -inf, 1.0, 2.0, -inf, 0.0], [1.0, 2.0, 2.0, -inf, -inf, 2.0]]
[[0.0, 2.0, -inf, -inf, 1.0, -inf], [1.0, 0.0, 1.0, 2.0, -inf, -inf], [1.0, 2.0, 2.0, -inf, -inf, 2.0]]
DIFF eps inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, 2.0, 2.0, -inf, -inf, 2.0]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
DIFF greedy val 4
DIFF margmax val 2
DIFF topkidx val 1
DIFF topkidx val 1
DIFF topkidx val 4
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 0.0], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, -inf, 0.0, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 0.0], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, -inf, 0.0, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 0.0, 0.0], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, -inf, 0.0, -inf, 0.0, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF eps inf
[[-inf, -inf, 0.0, -inf, 0.0, 0.0], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, -inf, 0.0], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF greedy val 3
DIFF margmax val 2
DIFF topkidx val 1
DIFF topkidx val 1
DIFF topkidx val 5
DIFF topkidx val 5
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [0.0, -inf, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [0.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [0.0, -inf, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [0.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, 2.0, 2.0, -inf], [-inf, 1.0, -inf, 2.0, -inf, -inf], [0.0, -inf, -inf, -inf, -inf, -2.0]]
[[-inf, -inf, -inf, 2.0, 2.0, -inf], [-inf, 1.0, -inf, 2.0, -inf, -inf], [0.0, -inf, -inf, -2.0, -inf, -inf]]
DIFF eps inf
[[-inf, -inf, -inf, 2.0, 2.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [0.0, -inf, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [0.0, -inf, -inf, -inf, -inf, -inf]]
DIFF greedy val 1
DIFF margmax val 2
DIFF topkidx val 1
DIFF topkidx val 1
DIFF topkidx val 1
DIFF topkidx val 1
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 0.0, -inf, -inf, 1.0], [2.0, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [0.0, -inf, -inf, -inf, -inf, 1.0], [2.0, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 1.0, -inf, -inf, 2.0, -inf], [-inf, 0.0, 0.0, -inf, -inf, 1.0], [2.0, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, 1.0, -inf, -inf, 2.0, -inf], [0.0, 0.0, -inf, -inf, -inf, 1.0], [2.0, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 1.0, -inf, 0.0, 2.0, -inf], [0.0, 0.0, 0.0, -inf, -inf, 1.0], [2.0, 2.0, -inf, -inf, 0.0, -inf]]
[[-inf, 1.0, 0.0, -inf, 2.0, -inf], [0.0, 0.0, 0.0, -inf, -inf, 1.0], [2.0, 2.0, -inf, 0.0, -inf, -inf]]
DIFF eps inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0], [2.0, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF greedy val 1
DIFF topkidx val 1
DIFF topkidx val 3
DIFF topp inf
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF topp inf
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF eps inf
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 1.0, -inf, 1.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF greedy val 2
DIFF topkidx val 1
DIFF topkidx val 1
DIFF topkidx val 5
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 1.0, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, 2.0]]
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 1.0, 1.0, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, 2.0]]
DIFF topp inf
[[-inf, -inf, 1.0, -inf, -inf, 2.0], [-inf, -inf, 1.0, 1.0, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, 2.0]]
[[-inf, 1.0, -inf, -inf, -inf, 2.0], [-inf, -inf, 1.0, 1.0, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, 2.0]]
DIFF eps inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 1.0, 1.0, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, 2.0]]
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
DIFF greedy val 2
DIFF margmax val 2
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
[[-inf, -inf, 2.0, -inf, -inf, -inf], [1.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
[[-inf, -inf, 2.0, -inf, -inf, -inf], [1.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, 2.0, -inf, -inf, 2.0], [1.0, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, 1.0, 2.0, -inf, -inf]]
[[-inf, -inf, 2.0, -inf, -inf, 2.0], [1.0, -inf, 1.0, -inf, -inf, -inf], [1.0, -inf, -inf, 2.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, 2.0, -inf, -inf, 2.0], [1.0, -inf, 1.0, -inf, -inf, -inf], [-inf, 1.0, 1.0, 2.0, -inf, -inf]]
[[-inf, -inf, 2.0, -inf, -inf, 2.0], [1.0, -inf, 1.0, -inf, -inf, -inf], [1.0, 1.0, -inf, 2.0, -inf, -inf]]
DIFF eps inf
[[-inf, -inf, 2.0, -inf, -inf, 2.0], [1.0, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF greedy val 3
DIFF margmax val 2
DIFF topkidx val 2
DIFF topkidx val 2
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, 2.0, -inf, 2.0], [2.0, -inf, -inf, -inf, -inf, 1.0], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, 2.0, -inf, 2.0, -inf, -inf], [2.0, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 2.0, -inf, 2.0, -inf, 2.0], [2.0, -inf, -inf, -inf, 1.0, 1.0], [-inf, -inf, 2.0, -inf, -inf, 1.0]]
[[-inf, 2.0, -inf, 2.0, -inf, 2.0], [2.0, -inf, 1.0, 1.0, -inf, -inf], [-inf, -inf, 2.0, 1.0, -inf, -inf]]
DIFF eps inf
[[-inf, 2.0, -inf, 2.0, -inf, 2.0], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, -inf, 2.0], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF greedy val 4
DIFF margmax val 2
DIFF topkidx val 1
DIFF topkidx val 4
DIFF topkidx val 5
DIFF topkidx val 5
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 1.0, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 1.0, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 1.0, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, 2.0, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 1.0, -inf, 0.0]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, 2.0, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 1.0, 0.0, -inf]]
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -1.0, -inf], [2.0, 2.0, 1.0, 1.0, -inf, 2.0], [-inf, -inf, -inf, 1.0, 0.0, 0.0]]
[[-inf, 2.0, -1.0, -inf, -inf, -inf], [2.0, 2.0, 1.0, 1.0, -inf, 2.0], [-inf, -inf, -inf, 1.0, 0.0, 0.0]]
DIFF eps inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, 2.0, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 1.0, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF greedy val 5
DIFF margmax val 4
DIFF topkidx val 2
DIFF topkidx val 2
DIFF topkidx val 2
DIFF topkidx val 2
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF eps inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [2.0, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF greedy val 2
DIFF topkidx val 3
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [1.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [1.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, 1.0, -inf, 1.0, -inf, 1.0]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [1.0, 1.0, -inf, 1.0, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, 1.0], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, 1.0, -inf, 1.0, -inf, 1.0]]
[[-inf, -inf, -inf, -inf, 2.0, 1.0], [-inf, -inf, 1.0, -inf, -inf, -inf], [1.0, 1.0, -inf, 1.0, -inf, -inf]]
DIFF eps inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [1.0, 1.0, -inf, 1.0, -inf, 1.0]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 1.0]]
DIFF greedy val 5
DIFF margmax val 5
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topp inf
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -1.0, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -1.0, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -1.0, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -1.0, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 0.0, -inf, 2.0, -inf, 0.0], [-inf, -inf, -1.0, -3.0, -1.0, -2.0], [-inf, 1.0, 2.0, 0.0, -inf, -inf]]
[[-inf, 0.0, -inf, 2.0, -inf, 0.0], [-3.0, -inf, -1.0, -inf, -1.0, -2.0], [-inf, 1.0, 2.0, 0.0, -inf, -inf]]
DIFF eps inf
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -1.0, -inf, -1.0, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, -1.0, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF greedy val 2
DIFF topkidx val 3
DIFF topkidx val 5
DIFF topkidx val 5
DIFF topp inf
[[-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 1.0, 1.0, -inf, -inf, -inf], [-inf, -inf, 2.0, 1.0, -inf, -inf], [-inf, -inf, 1.0, 1.0, 0.0, 0.0]]
[[-inf, 1.0, 1.0, -inf, -inf, -inf], [1.0, -inf, 2.0, -inf, -inf, -inf], [0.0, 0.0, 1.0, 1.0, -inf, -inf]]
DIFF topp inf
[[-inf, 1.0, 1.0, -inf, -inf, -1.0], [1.0, -inf, 2.0, 1.0, -inf, -inf], [0.0, 0.0, 1.0, 1.0, 0.0, 0.0]]
[[-1.0, 1.0, 1.0, -inf, -inf, -inf], [1.0, -inf, 2.0, 1.0, -inf, -inf], [0.0, 0.0, 1.0, 1.0, 0.0, 0.0]]
DIFF eps inf
[[-inf, 1.0, 1.0, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, 1.0, 1.0, -inf, -inf]]
[[-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf]]
DIFF greedy val 1
DIFF margmax val 1
DIFF topkidx val 4
DIFF topkidx val 1
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF eps inf
[[2.0, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf], [2.0, -inf, -inf, 2.0, -inf, -inf]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF greedy val 4
DIFF margmax val 4
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topp inf
[[-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 0.0, 1.0, -inf, 1.0, 0.0], [-inf, 2.0, -inf, -1.0, -1.0, -inf], [2.0, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, 0.0, 1.0, -inf, 1.0, 0.0], [-1.0, 2.0, -1.0, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, 2.0, -inf]]
DIFF eps inf
[[-inf, -inf, 1.0, -inf, 1.0, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
DIFF greedy val 4
DIFF margmax val 4
DIFF topkidx val 1
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topp inf
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
[[-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
[[-inf, -inf, 2.0, -inf, -inf, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, 2.0, 2.0, -inf, -inf], [-inf, -inf, 1.0, -inf, 1.0, -inf], [-inf, -inf, 2.0, 2.0, -inf, -inf]]
[[-inf, -inf, 2.0, 2.0, -inf, -inf], [-inf, -inf, 1.0, -inf, 1.0, -inf], [-inf, 2.0, 2.0, -inf, -inf, -inf]]
DIFF eps inf
[[-inf, -inf, 2.0, 2.0, -inf, -inf], [-inf, -inf, 1.0, -inf, 1.0, -inf], [-inf, 2.0, 2.0, 2.0, -inf, -inf]]
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, 1.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf]]
DIFF greedy val 2
DIFF margmax val 1
DIFF topkidx val 2
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF eps inf
[[-inf, 2.0, -inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, 2.0, -inf, -inf], [2.0, -inf, 2.0, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, 2.0, -inf, -inf, -inf]]
DIFF greedy val 4
DIFF margmax val 4
DIFF topkidx val 1
DIFF topkidx val 2
DIFF topkidx val 2
DIFF topkidx val 2
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[2.0, -inf, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[2.0, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, 2.0, -inf, -inf, 1.0, -inf]]
[[2.0, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, 2.0, 1.0, -inf, -inf, -inf]]
DIFF eps inf
[[2.0, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF greedy val 1
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topkidx val 3
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, -inf, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, 1.0, 2.0, -inf], [-inf, -inf, 1.0, 1.0, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
[[1.0, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, 1.0, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[1.0, 0.0, -inf, 1.0, 2.0, -inf], [-inf, -1.0, 1.0, 1.0, 0.0, -inf], [2.0, -inf, -inf, -inf, -inf, -1.0]]
[[1.0, 0.0, -inf, 1.0, 2.0, -inf], [-1.0, -inf, 1.0, 1.0, 0.0, -inf], [2.0, -inf, -1.0, -inf, -inf, -inf]]
DIFF eps inf
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, 1.0, 1.0, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
[[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 1.0, -inf, -inf], [2.0, -inf, -inf, -inf, -inf, -inf]]
DIFF greedy val 1
DIFF margmax val 1
DIFF topkidx val 5
DIFF topkidx val 5
DIFF topp inf
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -1.0, -inf, -inf, 0.0]]
[[-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -1.0, -inf, -inf, -inf, 0.0]]
DIFF topp inf
[[-inf, -inf, 1.0, 2.0, -inf, -1.0], [-inf, -inf, 0.0, -inf, 2.0, 1.0], [-inf, -1.0, -1.0, -inf, -2.0, 0.0]]
[[-1.0, -inf, 1.0, 2.0, -inf, -inf], [0.0, -inf, -inf, -inf, 2.0, 1.0], [-2.0, -1.0, -1.0, -inf, -inf, 0.0]]
DIFF topkidx val 1
DIFF topkidx val 1
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 1.0], [-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 1.0], [-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, 1.0, -inf, -inf, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
DIFF topp inf
[[-inf, 1.0, 0.0, 0.0, -inf, 1.0], [-inf, 2.0, 0.0, 1.0, -inf, -inf], [-inf, -inf, -inf, -1.0, 2.0, -inf]]
[[-inf, 1.0, 0.0, 0.0, -inf, 1.0], [-inf, 2.0, 0.0, 1.0, -inf, -inf], [-inf, -1.0, -inf, -inf, 2.0, -inf]]
DIFF eps inf
[[-inf, 1.0, -inf, -inf, -inf, 1.0], [-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
[[-inf, -inf, -inf, -inf, -inf, 1.0], [-inf, 2.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
DIFF greedy val 4
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topp inf
[[-inf, -inf, -1.0, -inf, -inf, -2.0], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
[[-2.0, -inf, -1.0, -inf, -inf, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, -inf]]
DIFF topp inf
[[-inf, -inf, -1.0, -inf, -2.0, -2.0], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, 1.0]]
[[-2.0, -inf, -1.0, -inf, -2.0, -inf], [-inf, -inf, -inf, 2.0, -inf, -inf], [-inf, 2.0, -inf, -inf, -inf, 1.0]]
DIFF topp inf
[[-2.0, -inf, -1.0, -3.0, -2.0, -2.0], [-inf, -inf, -inf, 2.0, -inf, -1.0], [-inf, 2.0, 0.0, -inf, -inf, 1.0]]
[[-2.0, -3.0, -1.0, -inf, -2.0, -2.0], [-inf, -1.0, -inf, 2.0, -inf, -inf], [0.0, 2.0, -inf, -inf, -inf, 1.0]]
DIFF margmax val 5
DIFF topkidx val 1
DIFF topkidx val 4
DIFF topkidx val 4
DIFF topkidx val 5
DIFF topp inf
[[-inf, -inf, -inf, -inf, -inf, 0.0], [-inf, -inf, -inf, 0.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
[[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
DIFF topp inf
[[-inf, -inf, -inf, -inf, 0.0, 0.0], [-inf, -inf, -inf, 0.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
[[0.0, 0.0, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
DIFF topp inf
[[-inf, -inf, 0.0, -inf, 0.0, 0.0], [-inf, -inf, 0.0, 0.0, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
[[0.0, 0.0, 0.0, -inf, -inf, -inf], [-inf, 0.0, 0.0, -inf, -inf, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
DIFF topp inf
[[-inf, 0.0, 0.0, -inf, 0.0, 0.0], [-inf, 0.0, 0.0, 0.0, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 2.0]]
[[0.0, 0.0, 0.0, -inf, 0.0, -inf], [-inf, 0.0, 0.0, 0.0, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 2.0]]
DIFF eps inf
[[0.0, 0.0, 0.0, -inf, 0.0, 0.0], [-inf, 0.0, 0.0, 0.0, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 2.0]]
[[-inf, -inf, -inf, -inf, -inf, 0.0], [-inf, 0.0, 0.0, 0.0, -inf, -inf], [-inf, -inf, -inf, 1.0, -inf, 2.0]]
DIFF eps inf
[[0.0, 0.0, 0.0, -inf, 0.0, 0.0], [-inf, 0.0, 0.0, 0.0, -inf,

Probe tie-breaking for greedy/argmax/topk

shell
python3 -c "
import torch, reforacle as R, objectives as O
l=torch.tensor([[1.,3.,3.,2.,3.,0.]])
print('ref greedy', R.greedy_token(l).tolist())
print('mine greedy', O.greedy_token(l).tolist())
print('argmax', l.argmax(-1).tolist())
m=torch.tensor([[True,False,True,True,False,True]])
l2=torch.tensor([[1.,5.,3.,3.,0.,3.]])
print('ref margmax', R.masked_argmax(l2,m).tolist())
print('mine margmax', O.masked_argmax(l2,m).tolist())
print('ref topkidx', R.topk_indices(l,3).tolist())
print('mine topkidx', O.topk_indices(l,3).tolist())
"
output
ref greedy [4]
mine greedy [1]
argmax [1]
ref margmax [5]
mine margmax [2]
ref topkidx [[1, 2, 4]]
mine topkidx [[2, 4, 1]]

[stdout]
ref greedy [4]
mine greedy [1]
argmax [1]
ref margmax [5]
mine margmax [2]
ref topkidx [[1, 2, 4]]
mine topkidx [[2, 4, 1]]

Probe filter tie behavior

shell
python3 -c "
import torch, reforacle as R
l=torch.tensor([[1.,3.,3.,2.,3.,0.]])
print('topk k=2', R.top_k_filter(l,2).tolist())
print('topk k=1', R.top_k_filter(l,1).tolist())
# top_p ties
l2=torch.tensor([[2.,2.,2.,2.]])
print('topp 0.5', R.top_p_filter(l2,0.5).tolist())
print('topp 0.6', R.top_p_filter(l2,0.6).tolist())
# min_p ties
print('minp 0.5', R.min_p_filter(l2,0.5).tolist())
# flip-argmax check for greedy multi-row
l3=torch.tensor([[3.,3.,1.],[0.,5.,5.]])
print('greedy', R.greedy_token(l3).tolist())
"
output
topk k=2 [[-inf, 3.0, 3.0, -inf, 3.0, -inf]]
topk k=1 [[-inf, 3.0, 3.0, -inf, 3.0, -inf]]
topp 0.5 [[2.0, 2.0, -inf, -inf]]
topp 0.6 [[2.0, 2.0, 2.0, -inf]]
minp 0.5 [[2.0, 2.0, 2.0, 2.0]]
greedy [1, 2]

[stdout]
topk k=2 [[-inf, 3.0, 3.0, -inf, 3.0, -inf]]
topk k=1 [[-inf, 3.0, 3.0, -inf, 3.0, -inf]]
topp 0.5 [[2.0, 2.0, -inf, -inf]]
topp 0.6 [[2.0, 2.0, 2.0, -inf]]
minp 0.5 [[2.0, 2.0, 2.0, 2.0]]
greedy [1, 2]

Clean probe epsilon/greedy/margmax ties

shell
python3 -c "
import torch, reforacle as R
# epsilon with tied max, prob below epsilon
l=torch.tensor([[2.,2.,-5.,2.,-5.,-5.]])
p=l.softmax(-1); print('probs',[round(x,4) for x in p[0].tolist()])
print('eps 0.5', R.epsilon_filter(l,0.5).tolist())  # all probs<0.5? maxprob=0.33
print('eps 0.1', R.epsilon_filter(l,0.1).tolist())
# greedy last-max confirm across rows
print('greedy', R.greedy_token(torch.tensor([[3.,3.,1.],[0.,5.,5.],[7.,7.,7.]])).tolist())
# masked_argmax last max
mm=torch.tensor([[True,True,True,False]]); print('margmax', R.masked_argmax(torch.tensor([[4.,4.,4.,9.]]),mm).tolist())
"
output
probs [0.333, 0.333, 0.0003, 0.333, 0.0003, 0.0003]
eps 0.5 [[-inf, -inf, -inf, 2.0, -inf, -inf]]
eps 0.1 [[2.0, 2.0, -inf, 2.0, -inf, -inf]]
greedy [1, 2, 2]
margmax [2]

[stdout]
probs [0.333, 0.333, 0.0003, 0.333, 0.0003, 0.0003]
eps 0.5 [[-inf, -inf, -inf, 2.0, -inf, -inf]]
eps 0.1 [[2.0, 2.0, -inf, 2.0, -inf, -inf]]
greedy [1, 2, 2]
margmax [2]

/app/objectives.py

− removed
NEG_INF = float("-inf")


def greedy_token(logits):
    return logits.argmax(dim=-1)
+ added
NEG_INF = float("-inf")


def _last_argmax(x):
    n = x.size(-1)
    return n - 1 - x.flip(-1).argmax(dim=-1)


def greedy_token(logits):
    return _last_argmax(logits)

/app/objectives.py

− removed
def top_p_filter(logits, p):
    sorted_logits, sorted_indices = torch.sort(logits, descending=False, dim=-1)
    cum_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
    sorted_remove = cum_probs <= (1 - p)
    sorted_remove[..., -1:] = False
    remove = sorted_remove.scatter(-1, sorted_indices, sorted_remove)
    return logits.masked_fill(remove, NEG_INF)
+ added
def top_p_filter(logits, p):
    sorted_logits, sorted_indices = torch.sort(logits, descending=True, stable=True, dim=-1)
    probs = sorted_logits.softmax(dim=-1)
    cum = probs.cumsum(dim=-1)
    sorted_remove = (cum - probs) >= p
    sorted_remove[..., 0] = False
    remove = sorted_remove.scatter(-1, sorted_indices, sorted_remove)
    return logits.masked_fill(remove, NEG_INF)

/app/objectives.py

− removed
def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    top = torch.topk(logits, 1, dim=-1).values[..., -1, None]
    remove = (probs < epsilon) & (logits < top)
    return logits.masked_fill(remove, NEG_INF)
+ added
def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    remove = probs < epsilon
    keep = _last_argmax(logits).unsqueeze(-1)
    remove.scatter_(-1, keep, False)
    return logits.masked_fill(remove, NEG_INF)

/app/objectives.py

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


def topk_indices(logits, k):
    return torch.topk(logits, k, dim=-1).indices
+ added
def masked_argmax(logits, mask):
    return _last_argmax(logits.masked_fill(~mask.bool(), NEG_INF))


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

Rerun full and tie stress tests

shell
python3 test_all.py && echo "=== TIE STRESS ===" && python3 -c "
import torch, reforacle as R, objectives as O
torch.manual_seed(7)
bad={}
def cmp(name,a,b):
    ok = a.shape==b.shape and a.dtype==b.dtype and torch.equal(a.isinf(),b.isinf())
    if ok:
        fin=~a.isinf()
        ok = torch.allclose(a[fin].float(),b[fin].float(),atol=1e-5,rtol=1e-4)
    if not ok: bad[name]=bad.get(name,0)+1
for t in range(400):
    l=torch.randint(-3,4,(3,7)).float()
    for k in [1,2,3,7]:
        cmp('top_k_filter',O.top_k_filter(l,k),R.top_k_filter(l,k))
        cmp('topk_indices',O.topk_indices(l,k),R.topk_indices(l,k))
    for p in [0.0,0.2,0.5,0.7,0.9,1.0]:
        cmp('top_p_filter',O.top_p_filter(l,p),R.top_p_filter(l,p))
        cmp('min_p_filter',O.min_p_filter(l,p),R.min_p_filter(l,p))
    for e in [0.05,0.1,0.2,0.34,0.5]:
        cmp('epsilon_filter',O.epsilon_filter(l,e),R.epsilon_filter(l,e))
    cmp('greedy_token',O.greedy_token(l),R.greedy_token(l))
    m=(torch.rand(3,7)>0.4)
    for b in range(3):
        if m[b].sum()==0: m[b,0]=True
    cmp('masked_argmax',O.masked_argmax(l,m),R.masked_argmax(l,m))
    for T in [0.5,1.0,2.0]:
      for k in [1,3]:
        for p in [0.5,1.0]:
          cmp('logits_warper',O.logits_warper(l,T,k,p),R.logits_warper(l,T,k,p))
          cmp('sample_greedy_warped',O.sample_greedy_warped(l,T,k,p),R.sample_greedy_warped(l,T,k,p))
print('tie fails:', bad if bad else 'NONE')
"
output
========================================
0 failing functions out of 37
=== TIE STRESS ===
tie fails: NONE

[stdout]
========================================
0 failing functions out of 37
=== TIE STRESS ===
tie fails: NONE

Edge case checks

shell
python3 -c "
import torch, reforacle as R, objectives as O
def cmp(name,a,b):
    ok = a.shape==b.shape and a.dtype==b.dtype and torch.equal(a.isinf(),b.isinf())
    if ok:
        fin=~a.isinf(); ok=torch.allclose(a[fin].float(),b[fin].float(),atol=1e-5,rtol=1e-4)
    print(('OK  ' if ok else 'FAIL')+' '+name, '' if ok else (a.tolist(),b.tolist()))
# 1D filters
l=torch.randn(6)
for fn in ['top_k_filter','top_p_filter','min_p_filter','epsilon_filter']:
    pass
cmp('1d topk',O.top_k_filter(l,2),R.top_k_filter(l,2))
cmp('1d topp',O.top_p_filter(l,0.5),R.top_p_filter(l,0.5))
cmp('1d minp',O.min_p_filter(l,0.3),R.min_p_filter(l,0.3))
cmp('1d eps',O.epsilon_filter(l,0.2),R.epsilon_filter(l,0.2))
cmp('1d temp0',O.temperature_scale(l,0.0),R.temperature_scale(l,0.0))
cmp('1d renorm',O.renormalize_logits(l),R.renormalize_logits(l))
cmp('1d entropy? ',O.entropy_from_logits(l).reshape(1),R.entropy_from_logits(l).reshape(1))
# packed truncation total<sum
sl=torch.tensor([3,4,2])
cmp('packed trunc',O.build_position_ids_packed(sl,5),R.build_position_ids_packed(sl,5))
cmp('packed exact',O.build_position_ids_packed(sl,9),R.build_position_ids_packed(sl,9))
# gen stop min_len>L
ids=torch.tensor([[5,1,3,4],[1,1,1,1]])
cmp('genstop ml>L',O.generation_stop_mask(ids,1,10),R.generation_stop_mask(ids,1,10))
# pad already multiple
ids2=torch.tensor([[1,2,3,4]])
cmp('pad exact',O.pad_to_multiple(ids2,0,4),R.pad_to_multiple(ids2,0,4))
cmp('pad exact2',O.pad_to_multiple(ids2,0,2),R.pad_to_multiple(ids2,0,2))
# 3d discount / reward
x=torch.randn(2,3,4)
cmp('3d discount',O.discount_cumsum(x,0.9),R.discount_cumsum(x,0.9))
cmp('3d reward',O.reward_to_go(x,(torch.rand(2,3,4)>0.3).float()),R.reward_to_go(x,(torch.rand(2,3,4)>0.3).float()) if False else O.reward_to_go(x,(torch.rand(2,3,4)>0.3).float()))
" 2>&1 | head -40
output
OK   1d topk 
OK   1d topp 
OK   1d minp 
OK   1d eps 
OK   1d temp0 
OK   1d renorm 
OK   1d entropy?  
Traceback (most recent call last):
  File "<string>", line 21, in <module>
  File "/app/objectives.py", line 184, in build_position_ids_packed
    out[: cat.size(0)] = cat
    ~~~^^^^^^^^^^^^^^^
RuntimeError: The expanded size of the tensor (5) must match the existing size (9) at non-singleton dimension 0.  Target sizes: [5].  Tensor sizes: [9]

[stdout]
OK   1d topk 
OK   1d topp 
OK   1d minp 
OK   1d eps 
OK   1d temp0 
OK   1d renorm 
OK   1d entropy?  
Traceback (most recent call last):
  File "<string>", line 21, in <module>
  File "/app/objectives.py", line 184, in build_position_ids_packed
    out[: cat.size(0)] = cat
    ~~~^^^^^^^^^^^^^^^
RuntimeError: The expanded size of the tensor (5) must match the existing size (9) at non-singleton dimension 0.  Target sizes: [5].  Tensor sizes: [9]

Check packed truncation behavior

shell
python3 -c "
import torch, reforacle as R
sl=torch.tensor([3,4,2])
print('total5', R.build_position_ids_packed(sl,5).tolist())
print('total7', R.build_position_ids_packed(sl,7).tolist())
"
output
total5 [0, 1, 2, 0, 1]
total7 [0, 1, 2, 0, 1, 2, 3]

[stdout]
total5 [0, 1, 2, 0, 1]
total7 [0, 1, 2, 0, 1, 2, 3]

/app/objectives.py

− removed
    out = torch.zeros(total_len, dtype=torch.long, device=seq_lengths.device)
    if parts:
        cat = torch.cat(parts)
        out[: cat.size(0)] = cat
    return out
+ added
    out = torch.zeros(total_len, dtype=torch.long, device=seq_lengths.device)
    if parts:
        cat = torch.cat(parts)[:total_len]
        out[: cat.size(0)] = cat
    return out

Recheck edge cases

shell
python3 -c "
import torch, reforacle as R, objectives as O
def cmp(name,a,b):
    ok = a.shape==b.shape and a.dtype==b.dtype and torch.equal(a.isinf(),b.isinf())
    if ok:
        fin=~a.isinf(); ok=torch.allclose(a[fin].float(),b[fin].float(),atol=1e-5,rtol=1e-4)
    print(('OK  ' if ok else 'FAIL')+' '+name, '' if ok else (a.tolist(),b.tolist()))
sl=torch.tensor([3,4,2])
cmp('packed trunc',O.build_position_ids_packed(sl,5),R.build_position_ids_packed(sl,5))
cmp('packed exact',O.build_position_ids_packed(sl,9),R.build_position_ids_packed(sl,9))
cmp('packed big',O.build_position_ids_packed(sl,12),R.build_position_ids_packed(sl,12))
ids=torch.tensor([[5,1,3,4],[1,1,1,1]])
cmp('genstop ml>L',O.generation_stop_mask(ids,1,10),R.generation_stop_mask(ids,1,10))
ids2=torch.tensor([[1,2,3,4]])
cmp('pad exact',O.pad_to_multiple(ids2,0,4),R.pad_to_multiple(ids2,0,4))
cmp('pad exact2',O.pad_to_multiple(ids2,0,2),R.pad_to_multiple(ids2,0,2))
x=torch.randn(2,3,4); msk=(torch.rand(2,3,4)>0.3).float()
cmp('3d discount',O.discount_cumsum(x,0.9),R.discount_cumsum(x,0.9))
cmp('3d reward',O.reward_to_go(x,msk),R.reward_to_go(x,msk))
# temperature negative
l=torch.randn(2,5)
cmp('temp neg',O.temperature_scale(l,-2.0),R.temperature_scale(l,-2.0))
# suppress with single int list
cmp('suppress',O.suppress_tokens(l,[0]),R.suppress_tokens(l,[0]))
"
output
OK   packed trunc 
OK   packed exact 
OK   packed big 
OK   genstop ml>L 
OK   pad exact 
OK   pad exact2 
OK   3d discount 
OK   3d reward 
OK   temp neg 
OK   suppress

[stdout]
OK   packed trunc 
OK   packed exact 
OK   packed big 
OK   genstop ml>L 
OK   pad exact 
OK   pad exact2 
OK   3d discount 
OK   3d reward 
OK   temp neg 
OK   suppress

Verify standalone and final run

shell
grep -n "reforacle" objectives.py; echo "import check:"; python3 -c "import ast; ast.parse(open('objectives.py').read()); print('parses ok')"; echo "--- run full harness x3 seeds ---"; for s in 1 2 3; do python3 -c "
import torch; torch.manual_seed($s)
import importlib, test_all
" 2>/dev/null; done; python3 test_all.py
output
import check:
parses ok
--- run full harness x3 seeds ---
========================================
0 failing functions out of 37
========================================
0 failing functions out of 37
========================================
0 failing functions out of 37
========================================
0 failing functions out of 37

[stdout]
import check:
parses ok
--- run full harness x3 seeds ---
========================================
0 failing functions out of 37
========================================
0 failing functions out of 37
========================================
0 failing functions out of 37
========================================
0 failing functions out of 37

Final sweep, different seed, more trials

shell
sed 's/manual_seed(123)/manual_seed(999)/; s/range(50)/range(120)/' test_all.py > /tmp/t2.py && python3 /tmp/t2.py
output
Exit code 1
Traceback (most recent call last):
  File "/tmp/t2.py", line 2, in <module>
    import reforacle as R
ModuleNotFoundError: No module named 'reforacle'

[error] tool reported failure

Final sweep from app dir

shell
sed 's/manual_seed(123)/manual_seed(999)/; s/range(50)/range(120)/' test_all.py > /app/t2.py && python3 t2.py; rm -f /app/t2.py
output
========================================
0 failing functions out of 37

[stdout]
========================================
0 failing functions out of 37
All 37 functions pass across every test: 120 randomized trials on multiple seeds, exact tie-breaking cases, 1D/3D inputs, and boundary conditions. I reverse-engineered `reforacle` entirely by probing, recovering several non-obvious conventions that a naive implementation would miss: - **Argmax family returns the *last* max on ties** (`greedy_token`, `masked_argmax`, `sample_greedy_warped`), not torch's default first-max. - **`topk_indices` uses stable descending sort** (ties in ascending index order), not `torch.topk`. - **`top_p_filter`** keeps the smallest top-set with cumulative prob ≥ p, breaking ties toward lower indices. - **`epsilon_filter`** force-keeps the last-argmax token when everything falls below ε. - **`repetition_penalty`** applies the penalty per *occurrence count* (`penalty**count`), not once per unique token. - **`prepare_labels`** shifts then masks the prompt region (`i < prompt_len-1`), whereas **`completion_loss_labels`** masks the first `prompt_len` with no shift. - **`generation_stop_mask`** length = `max(first_eos+1, min_len)` capped at sequence length. - **`temperature_scale(T=0)`** returns a hard one-hot (`0` at argmax, `-inf` elsewhere). The file `/app/objectives.py` depends only on `torch` , no import, call, or reference to `reforacle` remains, so it runs standalone in the grading sandbox.
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_74e3520b6b40432b. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.

Trial trial_74e3520b6b40432b · verifier authoritative; classifier explanatory.