SyncVals · Trajectory

reverse-engineer-decoding

claude-code claude-opus-4-8 ✓ resolved GOOD_SUCCESS ↑ 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_SUCCESSLegitimate solve, implements the asked-for behavior; tests verify real functionality.

SubtypeLegitimate Solution - Reverse Engineering

EvidenceTest output: "FUNCTIONS PASSED: 37/37 (partial_score=1.0)\nREWARD = 1 (1 iff ALL 37 functions match the sealed reference exactly)". Agent trajectory shows: step 51 final sweep "FINAL: ALL 37 PASS (60 seeds, ties+edges)", step 52 robustness "ALL EXTRA CHECKS PASS", step 53 standalone verification "functions defined: 37, NotImplementedError remaining: 0, imports reforacle: False". Agent documented key non-obvious behaviors discovered through probing: last-argmax convention, stable sort for top_p, epsilon_filter keeping exactly one token, repetition_penalty being count-based, sequence_lengths vs last_token_indices semantics.

Root causeAgent successfully reverse-engineered the sealed oracle module by systematically probing its behavior across diverse inputs, identifying 37 functions' exact conventions (including subtle tie-breaking and edge-case behaviors), and implementing them in standalone torch code that passes comprehensive validation across 60+ random seeds and edge cases.

RecommendationN/A - task is fine

Trajectory

Tool-by-tool agent trajectory

55 tool calls · 4 tool types · 57 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 members

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

Inspect constants

shell

python3 -c "import reforacle; print('NEG_INF', repr(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

shell

python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(0)
logits=torch.randn(2,6,generator=g)
print('logits',logits.tolist())
print('greedy',R.greedy_token(logits), R.greedy_token(logits).dtype, R.greedy_token(logits).shape)
print('temp T=2',R.temperature_scale(logits,2.0).tolist())
print('temp T=0',R.temperature_scale(logits,0.0).tolist())
# 1d logits
l1=torch.randn(6,generator=g)
print('greedy1d',R.greedy_token(l1), R.greedy_token(l1).shape)
print('temp1d T=0',R.temperature_scale(l1,0.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]]
greedy tensor([0, 1]) torch.int64 torch.Size([2])
temp T=2 [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
temp T=0 [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]]
greedy1d tensor(1) torch.Size([])
temp1d T=0 [-inf, 0.0, -inf, -inf, -inf, -inf]

[stdout]
logits [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
greedy tensor([0, 1]) torch.int64 torch.Size([2])
temp T=2 [[0.7704980373382568, -0.146714448928833, -1.0893946886062622, 0.28421562910079956, -0.5422611832618713, -0.6992977261543274], [0.20167341828346252, 0.4190131723880768, -0.3596287965774536, -0.20167176425457, -0.29831767082214355, 0.09101824462413788]]
temp T=0 [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]]
greedy1d tensor(1) torch.Size([])
temp1d T=0 [-inf, 0.0, -inf, -inf, -inf, -inf]

Probe top_k_filter

shell

python3 -c "
import reforacle as R, torch
g=torch.Generator().manual_seed(1)
logits=torch.randn(2,6,generator=g)
print('logits',logits.tolist())
for k in [1,2,3]:
    print('topk',k,R.top_k_filter(logits,k).tolist())
print('topk1d',R.top_k_filter(logits[0],2).tolist())
"

output

logits [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
topk 1 [[0.6613521575927734, -inf, -inf, -inf, -inf, -inf], [-inf, 0.38168391585350037, -inf, -inf, -inf, -inf]]
topk 2 [[0.6613521575927734, -inf, -inf, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -inf, -inf, -0.05825017765164375]]
topk 3 [[0.6613521575927734, 0.266924113035202, -inf, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -inf, -0.05825017765164375]]
topk1d [0.6613521575927734, -inf, -inf, 0.6213173270225525, -inf, -inf]

[stdout]
logits [[0.6613521575927734, 0.266924113035202, 0.06167725846171379, 0.6213173270225525, -0.4519059658050537, -0.16613022983074188], [-1.522768497467041, 0.38168391585350037, -1.0276086330413818, -0.563052773475647, -0.8922905325889587, -0.05825017765164375]]
topk 1 [[0.6613521575927734, -inf, -inf, -inf, -inf, -inf], [-inf, 0.38168391585350037, -inf, -inf, -inf, -inf]]
topk 2 [[0.6613521575927734, -inf, -inf, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -inf, -inf, -0.05825017765164375]]
topk 3 [[0.6613521575927734, 0.266924113035202, -inf, 0.6213173270225525, -inf, -inf], [-inf, 0.38168391585350037, -inf, -0.563052773475647, -inf, -0.05825017765164375]]
topk1d [0.6613521575927734, -inf, -inf, 0.6213173270225525, -inf, -inf]

Probe top_p, min_p, epsilon filters

shell

python3 -c "
import reforacle as R, torch
torch.set_printoptions(precision=6, sci_mode=False)
g=torch.Generator().manual_seed(2)
logits=torch.randn(2,6,generator=g)
print('logits',logits.tolist())
for p in [0.3,0.5,0.9]:
    print('topp',p,R.top_p_filter(logits,p).tolist())
print()
for p in [0.1,0.3,0.5]:
    print('minp',p,R.min_p_filter(logits,p).tolist())
print()
for e in [0.01,0.1,0.2]:
    print('eps',e,R.epsilon_filter(logits,e).tolist())
"

output

logits [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 1.318617343902588, 0.747639000415802]]
topp 0.3 [[-inf, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, -inf]]
topp 0.5 [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, 0.747639000415802]]
topp 0.9 [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -inf, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -inf, -inf, 1.318617343902588, 0.747639000415802]]

minp 0.1 [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -0.39190584421157837, -inf, 1.318617343902588, 0.747639000415802]]
minp 0.3 [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [0.5731067657470703, 0.540947437286377, -inf, -inf, 1.318617343902588, 0.747639000415802]]
minp 0.5 [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, 0.747639000415802]]

eps 0.01 [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 1.318617343902588, 0.747639000415802]]
eps 0.1 [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -inf, 1.0444635152816772, -inf], [0.5731067657470703, 0.540947437286377, -inf, -inf, 1.318617343902588, 0.747639000415802]]
eps 0.2 [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, 0.747639000415802]]

[stdout]
logits [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 1.318617343902588, 0.747639000415802]]
topp 0.3 [[-inf, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, -inf]]
topp 0.5 [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, 0.747639000415802]]
topp 0.9 [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -inf, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -inf, -inf, 1.318617343902588, 0.747639000415802]]

minp 0.1 [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -0.39190584421157837, -inf, 1.318617343902588, 0.747639000415802]]
minp 0.3 [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [0.5731067657470703, 0.540947437286377, -inf, -inf, 1.318617343902588, 0.747639000415802]]
minp 0.5 [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, 0.747639000415802]]

eps 0.01 [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -1.2050371170043945, 1.0444635152816772, -0.6332277059555054], [0.5731067657470703, 0.540947437286377, -0.39190584421157837, -1.0426788330078125, 1.318617343902588, 0.747639000415802]]
eps 0.1 [[0.39229682087898254, -0.223564013838768, -0.31950026750564575, -inf, 1.0444635152816772, -inf], [0.5731067657470703, 0.540947437286377, -inf, -inf, 1.318617343902588, 0.747639000415802]]
eps 0.2 [[0.39229682087898254, -inf, -inf, -inf, 1.0444635152816772, -inf], [-inf, -inf, -inf, -inf, 1.318617343902588, 0.747639000415802]]

Probe penalties, suppress, renorm

shell

python3 -c "
import reforacle as R, torch
torch.set_printoptions(precision=6, sci_mode=False)
g=torch.Generator().manual_seed(3)
logits=torch.randn(2,6,generator=g)
ids=torch.tensor([[0,0,1,4],[2,2,2,5]])
print('logits',logits.tolist())
print('reppen1.2',R.repetition_penalty(logits,ids,1.2).tolist())
print('freqpen0.5',R.frequency_penalty(logits,ids,0.5).tolist())
print('prespen0.5',R.presence_penalty(logits,ids,0.5).tolist())
print('suppress',R.suppress_tokens(logits,[1,3]).tolist())
print('renorm',R.renormalize_logits(logits).tolist())
print('renorm sum exp', torch.exp(R.renormalize_logits(logits)).sum(-1))
"

output

logits [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -0.6137180328369141, 0.04618244990706444, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -1.435179352760315, 0.9774317741394043, 0.5220499634742737, 1.2379167079925537]]
reppen1.2 [[0.5578305721282959, 0.14569447934627533, 0.08897809684276581, -0.6137180328369141, 0.03848537430167198, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -2.479990243911743, 0.9774317741394043, 0.5220499634742737, 1.0315972566604614]]
freqpen0.5 [[-0.19672393798828125, -0.32516661286354065, 0.08897809684276581, -0.6137180328369141, -0.45381754636764526, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -2.9351792335510254, 0.9774317741394043, 0.5220499634742737, 0.7379167079925537]]
prespen0.5 [[0.30327606201171875, -0.32516661286354065, 0.08897809684276581, -0.6137180328369141, -0.45381754636764526, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -1.935179352760315, 0.9774317741394043, 0.5220499634742737, 0.7379167079925537]]
suppress [[0.8032760620117188, -inf, 0.08897809684276581, -inf, 0.04618244990706444, -1.3682591915130615], [0.3374950885772705, -inf, -1.435179352760315, -inf, 0.5220499634742737, 1.2379167079925537]]
renorm [[-1.0467677116394043, -1.6752103567123413, -1.7610657215118408, -2.463761806488037, -1.803861379623413, -3.2183029651641846], [-2.1622753143310547, -1.488654375076294, -3.9349498748779297, -1.5223385095596313, -1.9777202606201172, -1.261853575706482]]
renorm sum exp tensor([1.000000, 1.000000])

[stdout]
logits [[0.8032760620117188, 0.17483338713645935, 0.08897809684276581, -0.6137180328369141, 0.04618244990706444, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -1.435179352760315, 0.9774317741394043, 0.5220499634742737, 1.2379167079925537]]
reppen1.2 [[0.5578305721282959, 0.14569447934627533, 0.08897809684276581, -0.6137180328369141, 0.03848537430167198, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -2.479990243911743, 0.9774317741394043, 0.5220499634742737, 1.0315972566604614]]
freqpen0.5 [[-0.19672393798828125, -0.32516661286354065, 0.08897809684276581, -0.6137180328369141, -0.45381754636764526, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -2.9351792335510254, 0.9774317741394043, 0.5220499634742737, 0.7379167079925537]]
prespen0.5 [[0.30327606201171875, -0.32516661286354065, 0.08897809684276581, -0.6137180328369141, -0.45381754636764526, -1.3682591915130615], [0.3374950885772705, 1.0111159086227417, -1.935179352760315, 0.9774317741394043, 0.5220499634742737, 0.7379167079925537]]
suppress [[0.8032760620117188, -inf, 0.08897809684276581, -inf, 0.04618244990706444, -1.3682591915130615], [0.3374950885772705, -inf, -1.435179352760315, -inf, 0.5220499634742737, 1.2379167079925537]]
renorm [[-1.0467677116394043, -1.6752103567123413, -1.7610657215118408, -2.463761806488037, -1.803861379623413, -3.2183029651641846], [-2.1622753143310547, -1.488654375076294, -3.9349498748779297, -1.5223385095596313, -1.9777202606201172, -1.261853575706482]]
renorm sum exp tensor([1.000000, 1.000000])

Probe min_length_filter

shell

python3 -c "
import reforacle as R, torch
torch.set_printoptions(precision=6, sci_mode=False)
g=torch.Generator().manual_seed(4)
logits=torch.randn(2,6,generator=g)
print('minlen cur<min',R.min_length_filter(logits,3,5,2).tolist())
print('minlen cur>=min',R.min_length_filter(logits,5,5,2).tolist())
print('minlen cur>min',R.min_length_filter(logits,6,5,2).tolist())
"

output

minlen cur<min [[-1.605276346206665, 0.23248571157455444, -inf, 0.8472937941551208, 1.2006442546844482, -0.4015503227710724], [-1.4260196685791016, 0.903931736946106, -inf, 0.6888809204101562, 0.8849856853485107, 1.7706429958343506]]
minlen cur>=min [[-1.605276346206665, 0.23248571157455444, 2.239870071411133, 0.8472937941551208, 1.2006442546844482, -0.4015503227710724], [-1.4260196685791016, 0.903931736946106, 0.8557155728340149, 0.6888809204101562, 0.8849856853485107, 1.7706429958343506]]
minlen cur>min [[-1.605276346206665, 0.23248571157455444, 2.239870071411133, 0.8472937941551208, 1.2006442546844482, -0.4015503227710724], [-1.4260196685791016, 0.903931736946106, 0.8557155728340149, 0.6888809204101562, 0.8849856853485107, 1.7706429958343506]]

[stdout]
minlen cur<min [[-1.605276346206665, 0.23248571157455444, -inf, 0.8472937941551208, 1.2006442546844482, -0.4015503227710724], [-1.4260196685791016, 0.903931736946106, -inf, 0.6888809204101562, 0.8849856853485107, 1.7706429958343506]]
minlen cur>=min [[-1.605276346206665, 0.23248571157455444, 2.239870071411133, 0.8472937941551208, 1.2006442546844482, -0.4015503227710724], [-1.4260196685791016, 0.903931736946106, 0.8557155728340149, 0.6888809204101562, 0.8849856853485107, 1.7706429958343506]]
minlen cur>min [[-1.605276346206665, 0.23248571157455444, 2.239870071411133, 0.8472937941551208, 1.2006442546844482, -0.4015503227710724], [-1.4260196685791016, 0.903931736946106, 0.8557155728340149, 0.6888809204101562, 0.8849856853485107, 1.7706429958343506]]

Probe mask utilities

shell

python3 -c "
import reforacle as R, torch
ids=torch.tensor([[5,6,7,0,0],[3,0,4,0,0]])
print('attn_mask pad0',R.attention_mask(ids,0).tolist(), R.attention_mask(ids,0).dtype)
print('pos_ids 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)
print('mask_after_eos eos7',R.mask_after_eos(ids,7).tolist(), R.mask_after_eos(ids,7).dtype)
print('causal 4', R.causal_mask(4).tolist(), R.causal_mask(4).dtype)
"

output

attn_mask pad0 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 0.0, 1.0, 0.0, 0.0]] torch.float32
pos_ids pad0 [[0, 1, 2, 2, 2], [0, 0, 1, 1, 1]] torch.int64
seqlen [3, 2] torch.int64
mask_after_eos eos7 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]] torch.float32
causal 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_mask pad0 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 0.0, 1.0, 0.0, 0.0]] torch.float32
pos_ids pad0 [[0, 1, 2, 2, 2], [0, 0, 1, 1, 1]] torch.int64
seqlen [3, 2] torch.int64
mask_after_eos eos7 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]] torch.float32
causal 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 label functions

shell

python3 -c "
import reforacle as R, torch
ids=torch.tensor([[5,6,7,8,9],[3,4,5,6,7]])
print('shift ii=-100',R.shift_labels(ids,-100).tolist(), R.shift_labels(ids,-100).dtype, R.shift_labels(ids,-100).shape)
print('completion pl=[2,3] ii=-100',R.completion_loss_labels(ids,torch.tensor([2,3]),-100).tolist())
print('prepare pl=[2,3] ii=-100',R.prepare_labels(ids,torch.tensor([2,3]),-100).tolist())
"

output

shift ii=-100 [[6, 7, 8, 9, -100], [4, 5, 6, 7, -100]] torch.int64 torch.Size([2, 5])
completion pl=[2,3] ii=-100 [[-100, -100, 7, 8, 9], [-100, -100, -100, 6, 7]]
prepare pl=[2,3] ii=-100 [[-100, 7, 8, 9, -100], [-100, -100, 6, 7, -100]]

[stdout]
shift ii=-100 [[6, 7, 8, 9, -100], [4, 5, 6, 7, -100]] torch.int64 torch.Size([2, 5])
completion pl=[2,3] ii=-100 [[-100, -100, 7, 8, 9], [-100, -100, -100, 6, 7]]
prepare pl=[2,3] ii=-100 [[-100, 7, 8, 9, -100], [-100, -100, 6, 7, -100]]

Probe truncate/pad functions

shell

python3 -c "
import reforacle as R, torch
ids=torch.tensor([[5,6,7,8,9],[3,4,5,6,7]])
print('truncate_left max3',R.truncate_left(ids,3).tolist())
print('truncate_left max7',R.truncate_left(ids,7).tolist())
print('pad_to_multiple pad0 mult4',R.pad_to_multiple(ids,0,4).tolist())
print('pad_to_multiple pad0 mult8',R.pad_to_multiple(ids,0,8).tolist())
print('right_pad_to pad0 w7',R.right_pad_to(ids,0,7).tolist())
print('right_pad_to pad0 w3',R.right_pad_to(ids,0,3).tolist())
"

output

truncate_left max3 [[7, 8, 9], [5, 6, 7]]
truncate_left max7 [[5, 6, 7, 8, 9], [3, 4, 5, 6, 7]]
pad_to_multiple pad0 mult4 [[5, 6, 7, 8, 9, 0, 0, 0], [3, 4, 5, 6, 7, 0, 0, 0]]
pad_to_multiple pad0 mult8 [[5, 6, 7, 8, 9, 0, 0, 0], [3, 4, 5, 6, 7, 0, 0, 0]]
right_pad_to pad0 w7 [[5, 6, 7, 8, 9, 0, 0], [3, 4, 5, 6, 7, 0, 0]]
right_pad_to pad0 w3 [[5, 6, 7], [3, 4, 5]]

[stdout]
truncate_left max3 [[7, 8, 9], [5, 6, 7]]
truncate_left max7 [[5, 6, 7, 8, 9], [3, 4, 5, 6, 7]]
pad_to_multiple pad0 mult4 [[5, 6, 7, 8, 9, 0, 0, 0], [3, 4, 5, 6, 7, 0, 0, 0]]
pad_to_multiple pad0 mult8 [[5, 6, 7, 8, 9, 0, 0, 0], [3, 4, 5, 6, 7, 0, 0, 0]]
right_pad_to pad0 w7 [[5, 6, 7, 8, 9, 0, 0], [3, 4, 5, 6, 7, 0, 0]]
right_pad_to pad0 w3 [[5, 6, 7], [3, 4, 5]]

Probe last_token_indices, gather_last_logits

shell

python3 -c "
import reforacle as R, torch
am=torch.tensor([[1.,1,1,0,0],[1,0,1,0,0],[1,1,1,1,1]])
print('last_token_indices',R.last_token_indices(am).tolist(), R.last_token_indices(am).dtype)
g=torch.Generator().manual_seed(5)
logits=torch.randn(3,5,4,generator=g)
gl=R.gather_last_logits(logits,am)
print('gather_last shape',gl.shape)
print('gather_last',gl.tolist())
# check equals logits at last index
print('manual', logits[0,2].tolist(), logits[1,2].tolist(), logits[2,4].tolist())
"

output

last_token_indices [2, 2, 4] torch.int64
gather_last shape torch.Size([3, 4])
gather_last [[0.39076822996139526, -0.0189677607268095, -1.352687120437622, -0.7308173179626465], [1.4829323291778564, 0.49403807520866394, 0.24921877682209015, 1.7470314502716064], [0.28316107392311096, -1.49812912940979, -0.30913057923316956, -0.6040987968444824]]
manual [0.39076822996139526, -0.0189677607268095, -1.352687120437622, -0.7308173179626465] [1.4829323291778564, 0.49403807520866394, 0.24921877682209015, 1.7470314502716064] [0.28316107392311096, -1.49812912940979, -0.30913057923316956, -0.6040987968444824]

[stdout]
last_token_indices [2, 2, 4] torch.int64
gather_last shape torch.Size([3, 4])
gather_last [[0.39076822996139526, -0.0189677607268095, -1.352687120437622, -0.7308173179626465], [1.4829323291778564, 0.49403807520866394, 0.24921877682209015, 1.7470314502716064], [0.28316107392311096, -1.49812912940979, -0.30913057923316956, -0.6040987968444824]]
manual [0.39076822996139526, -0.0189677607268095, -1.352687120437622, -0.7308173179626465] [1.4829323291778564, 0.49403807520866394, 0.24921877682209015, 1.7470314502716064] [0.28316107392311096, -1.49812912940979, -0.30913057923316956, -0.6040987968444824]

Probe action_logprobs, entropy

shell

python3 -c "
import reforacle as R, torch
torch.set_printoptions(precision=6, sci_mode=False)
g=torch.Generator().manual_seed(6)
logits=torch.randn(2,3,5,generator=g)
actions=torch.tensor([[1,4,2],[0,3,1]])
alp=R.action_logprobs(logits,actions)
print('action_logprobs shape',alp.shape)
print('action_logprobs',alp.tolist())
# manual log_softmax gather
ls=torch.log_softmax(logits,dim=-1)
print('manual',ls.gather(-1,actions.unsqueeze(-1)).squeeze(-1).tolist())
print('entropy shape',R.entropy_from_logits(logits).shape)
print('entropy',R.entropy_from_logits(logits).tolist())
p=torch.softmax(logits,-1); print('manual ent',(-(p*torch.log_softmax(logits,-1)).sum(-1)).tolist())
"

output

action_logprobs shape torch.Size([2, 3])
action_logprobs [[-0.5579891204833984, -0.918282687664032, -3.401867151260376], [-2.8832602500915527, -0.9549267292022705, -2.012803554534912]]
manual [[-0.5579891204833984, -0.918282687664032, -3.401867151260376], [-2.8832602500915527, -0.9549267292022705, -2.012803554534912]]
entropy shape torch.Size([2, 3])
entropy [[1.238924503326416, 1.479904294013977, 0.8871920108795166], [1.2988593578338623, 1.5167714357376099, 1.4523380994796753]]
manual ent [[1.238924503326416, 1.479904294013977, 0.8871921300888062], [1.2988593578338623, 1.5167714357376099, 1.4523382186889648]]

[stdout]
action_logprobs shape torch.Size([2, 3])
action_logprobs [[-0.5579891204833984, -0.918282687664032, -3.401867151260376], [-2.8832602500915527, -0.9549267292022705, -2.012803554534912]]
manual [[-0.5579891204833984, -0.918282687664032, -3.401867151260376], [-2.8832602500915527, -0.9549267292022705, -2.012803554534912]]
entropy shape torch.Size([2, 3])
entropy [[1.238924503326416, 1.479904294013977, 0.8871920108795166], [1.2988593578338623, 1.5167714357376099, 1.4523380994796753]]
manual ent [[1.238924503326416, 1.479904294013977, 0.8871921300888062], [1.2988593578338623, 1.5167714357376099, 1.4523382186889648]]

Probe reward_to_go, discount_cumsum

shell

python3 -c "
import reforacle as R, torch
torch.set_printoptions(precision=6, sci_mode=False)
rewards=torch.tensor([[1.,2,3,4],[5,6,7,8]])
mask=torch.tensor([[1.,1,1,0],[1,1,0,0]])
print('reward_to_go',R.reward_to_go(rewards,mask).tolist())
print('reward_to_go nomask',R.reward_to_go(rewards,torch.ones_like(mask)).tolist())
x=torch.tensor([[1.,2,3,4]])
print('discount 0.9',R.discount_cumsum(x,0.9).tolist())
print('discount 1d',R.discount_cumsum(torch.tensor([1.,2,3]),0.5).tolist())
"

output

reward_to_go [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]]
reward_to_go nomask [[10.0, 9.0, 7.0, 4.0], [26.0, 21.0, 15.0, 8.0]]
discount 0.9 [[8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]]
discount 1d [2.75, 3.5, 3.0]

[stdout]
reward_to_go [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]]
reward_to_go nomask [[10.0, 9.0, 7.0, 4.0], [26.0, 21.0, 15.0, 8.0]]
discount 0.9 [[8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0]]
discount 1d [2.75, 3.5, 3.0]

Probe build_position_ids_packed

shell

python3 -c "
import reforacle as R, torch
sl=torch.tensor([2,3,1])
print('build_pos_packed total6',R.build_position_ids_packed(sl,6).tolist(), R.build_position_ids_packed(sl,6).dtype)
print('build_pos_packed total8',R.build_position_ids_packed(sl,8).tolist())
sl2=torch.tensor([3,2])
print('build_pos total5',R.build_position_ids_packed(sl2,5).tolist())
"

output

build_pos_packed total6 [0, 1, 0, 1, 2, 0] torch.int64
build_pos_packed total8 [0, 1, 0, 1, 2, 0, 0, 0]
build_pos total5 [0, 1, 2, 0, 1]

[stdout]
build_pos_packed total6 [0, 1, 0, 1, 2, 0] torch.int64
build_pos_packed total8 [0, 1, 0, 1, 2, 0, 0, 0]
build_pos total5 [0, 1, 2, 0, 1]

Probe warper, sample, topk_indices, masked_argmax

shell

python3 -c "
import reforacle as R, torch
torch.set_printoptions(precision=6, sci_mode=False)
g=torch.Generator().manual_seed(7)
logits=torch.randn(2,6,generator=g)
print('logits',logits.tolist())
print('warper T2 k3 p0.9',R.logits_warper(logits,2.0,3,0.9).tolist())
print('sample_greedy_warped',R.sample_greedy_warped(logits,2.0,3,0.9).tolist(),R.sample_greedy_warped(logits,2.0,3,0.9).dtype)
print('topk_indices k3',R.topk_indices(logits,3).tolist(),R.topk_indices(logits,3).dtype)
mask=torch.tensor([[1.,0,1,1,0,1],[0,0,1,1,1,0]])
print('masked_argmax',R.masked_argmax(logits,mask).tolist(),R.masked_argmax(logits,mask).dtype)
"

output

logits [[-0.1467950940132141, 0.7861412763595581, 0.9468216300010681, -1.1143440008163452, 1.6907901763916016, -0.8948279023170471], [-0.3556250333786011, 1.2323857545852661, 0.13817265629768372, -1.6821985244750977, 0.317678302526474, 0.13280697166919708]]
warper T2 k3 p0.9 [[-inf, 0.39307063817977905, 0.47341081500053406, -inf, 0.8453950881958008, -inf], [-inf, 0.6161928772926331, 0.06908632814884186, -inf, 0.158839151263237, -inf]]
sample_greedy_warped [4, 1] torch.int64
topk_indices k3 [[4, 2, 1], [1, 4, 2]] torch.int64
masked_argmax [2, 4] torch.int64

[stdout]
logits [[-0.1467950940132141, 0.7861412763595581, 0.9468216300010681, -1.1143440008163452, 1.6907901763916016, -0.8948279023170471], [-0.3556250333786011, 1.2323857545852661, 0.13817265629768372, -1.6821985244750977, 0.317678302526474, 0.13280697166919708]]
warper T2 k3 p0.9 [[-inf, 0.39307063817977905, 0.47341081500053406, -inf, 0.8453950881958008, -inf], [-inf, 0.6161928772926331, 0.06908632814884186, -inf, 0.158839151263237, -inf]]
sample_greedy_warped [4, 1] torch.int64
topk_indices k3 [[4, 2, 1], [1, 4, 2]] torch.int64
masked_argmax [2, 4] torch.int64

Probe rollout_last_logits, first_stop_index

shell

python3 -c "
import reforacle as R, torch
torch.set_printoptions(precision=6, sci_mode=False)
g=torch.Generator().manual_seed(8)
logits=torch.randn(2,5,4,generator=g)
ids=torch.tensor([[5,6,7,0,0],[3,4,5,6,0]])
rl=R.rollout_last_logits(logits,ids,0)
print('rollout_last shape',rl.shape)
print('rollout_last',rl.tolist())
print('manual last idx2,3', logits[0,2].tolist(), logits[1,3].tolist())
ids2=torch.tensor([[5,6,7,8,9],[3,9,5,6,7]])
print('first_stop [7,9]',R.first_stop_index(ids2,[7,9]).tolist(),R.first_stop_index(ids2,[7,9]).dtype)
print('first_stop [100]',R.first_stop_index(ids2,[100]).tolist())
"

output

Exit code 1
rollout_last shape torch.Size([2, 4])
rollout_last [[-0.6387829780578613, 1.3384251594543457, -1.1907936334609985, -0.7600940465927124], [1.6447601318359375, -1.6710339784622192, -0.35282421112060547, -0.26325100660324097]]
manual last idx2,3 [-0.6387829780578613, 1.3384251594543457, -1.1907936334609985, -0.7600940465927124] [1.6447601318359375, -1.6710339784622192, -0.35282421112060547, -0.26325100660324097]
Traceback (most recent call last):
  File "<string>", line 12, in <module>
  File "reforacle.py", line 243, in reforacle.first_stop_index
AttributeError: 'list' object has no attribute 'tolist'

[error] tool reported failure

Check first_stop_index return type, rollout interior pad

shell

python3 -c "
import reforacle as R, torch
ids2=torch.tensor([[5,6,7,8,9],[3,9,5,6,7]])
r=R.first_stop_index(ids2,[7,9])
print('type',type(r), r)
r2=R.first_stop_index(ids2,[100])
print('nomatch',type(r2),r2)
# interior pad test for rollout
logits=torch.randn(1,5,3)
ids=torch.tensor([[5,0,7,0,0]])
import torch as T
rl=R.rollout_last_logits(logits,ids,0)
print('rollout interior pad', (rl[0]==logits[0,2]).all().item(), (rl[0]==logits[0,0]).all().item())
"

output

Exit code 1
Traceback (most recent call last):
  File "<string>", line 4, in <module>
  File "reforacle.py", line 243, in reforacle.first_stop_index
AttributeError: 'list' object has no attribute 'tolist'

[error] tool reported failure

first_stop_index with tensor, rollout interior pad

shell

python3 -c "
import reforacle as R, torch
ids2=torch.tensor([[5,6,7,8,9],[3,9,5,6,7]])
r=R.first_stop_index(ids2,torch.tensor([7,9]))
print('type',type(r), r, r.dtype if hasattr(r,'dtype') else '')
r2=R.first_stop_index(ids2,torch.tensor([100]))
print('nomatch',r2)
# interior pad test for rollout
logits=torch.randn(1,5,3)
ids=torch.tensor([[5,0,7,0,0]])
rl=R.rollout_last_logits(logits,ids,0)
print('rollout interior: idx2?',(rl[0]==logits[0,2]).all().item(),'idx0?',(rl[0]==logits[0,0]).all().item())
"

output

type <class 'torch.Tensor'> tensor([2, 1]) torch.int64
nomatch tensor([-1, -1])
rollout interior: idx2? True idx0? False

[stdout]
type <class 'torch.Tensor'> tensor([2, 1]) torch.int64
nomatch tensor([-1, -1])
rollout interior: idx2? True idx0? False

Probe generation_stop_mask

shell

python3 -c "
import reforacle as R, torch
ids=torch.tensor([[5,6,2,8,2],[3,4,5,6,7],[2,4,5,2,7]])
for ml in [0,2,4]:
    r=R.generation_stop_mask(ids,2,ml)
    print('min_len',ml,r.tolist(), r.dtype)
"

output

min_len 0 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 0.0, 0.0, 0.0, 0.0]] torch.float32
min_len 2 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 0.0, 0.0, 0.0]] torch.float32
min_len 4 [[1.0, 1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 0.0]] torch.float32

[stdout]
min_len 0 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 0.0, 0.0, 0.0, 0.0]] torch.float32
min_len 2 [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 0.0, 0.0, 0.0]] torch.float32
min_len 4 [[1.0, 1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 0.0]] torch.float32

Determine repetition_penalty float method

shell

python3 -c "
import reforacle as R, torch
torch.manual_seed(0)
g=torch.Generator().manual_seed(11)
logits=torch.randn(3,8,generator=g)
ids=torch.tensor([[0,0,0,1,1,2],[3,3,3,3,4,4],[5,5,6,7,7,7]])
p=1.3
ref=R.repetition_penalty(logits,ids,p)
# my formula count-based
vocab=logits.shape[-1]
counts=torch.zeros_like(logits)
counts.scatter_add_(-1, ids, torch.ones_like(ids,dtype=logits.dtype))
pos=logits>0
mine=torch.where(pos, logits/ (p**counts), logits*(p**counts))
print('max diff pow', (ref-mine).abs().max().item())
# iterative approach
mine2=logits.clone()
score=torch.gather(mine2,-1,ids)
# hf-like but count: multiply/divide per occurrence
mine3=logits.clone()
for b in range(ids.shape[0]):
    for t in ids[b].tolist():
        v=mine3[b,t]
        mine3[b,t]= v/p if v>0 else v*p
print('max diff iter', (ref-mine3).abs().max().item())
"

output

max diff pow 1.1920928955078125e-07
max diff iter 0.0

[stdout]
max diff pow 1.1920928955078125e-07
max diff iter 0.0

write · /app/objectives.py

from __future__ import annotations

import torch

NEG_INF = float("-inf")


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


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


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


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


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


def repetition_penalty(logits, input_ids, penalty):
    out = logits.clone()
    seq = input_ids.shape[-1]
    for j in range(seq):
        tok = input_ids[..., j : j + 1]
        v = out.gather(-1, tok)
        newv = torch.where(v > 0, v / penalty, v * penalty)
        out.scatter_(-1, tok, newv)
    return out


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


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


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


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


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


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


def position_ids(input_ids, pad_id):
    mask = (input_ids != pad_id).long()
    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()
    seq = input_ids.shape[-1]
    pos = torch.arange(seq, device=input_ids.device)
    mask = pos < prompt_lengths.unsqueeze(-1)
    labels[mask] = ignore_index
    return labels


def mask_after_eos(input_ids, eos_id):
    seq = input_ids.shape[-1]
    pos = torch.arange(seq, device=input_ids.device)
    iseos = input_ids == eos_id
    cand = torch.where(iseos, pos, torch.full_like(input_ids, seq))
    first = cand.min(dim=-1).values
    valid = torch.where(first < seq, first + 1, torch.full_like(first, seq))
    return (pos < valid.unsqueeze(-1)).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.shape[-1]
    pos = torch.arange(L, device=attn_mask.device)
    masked_pos = torch.where(
        attn_mask != 0, pos, torch.full_like(pos, -1)
    )
    return masked_pos.max(dim=-1).values


def gather_last_logits(logits, attn_mask):
    idx = last_token_indices(attn_mask)
    return logits[torch.arange(logits.shape[0], 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):
    seq = input_ids.shape[-1]
    target = ((seq + multiple - 1) // multiple) * multiple
    pad = target - seq
    if pad <= 0:
        return input_ids
    shape = list(input_ids.shape)
    shape[-1] = pad
    padding = torch.full(shape, pad_id, dtype=input_ids.dtype, device=input_ids.device)
    return torch.cat([input_ids, padding], dim=-1)


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


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


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


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


def build_position_ids_packed(seq_lengths, total_len):
    result = torch.zeros(total_len, dtype=torch.long, device=seq_lengths.device)
    offset = 0
    for l in seq_lengths.tolist():
        result[offset : offset + l] = torch.arange(l, device=seq_lengths.device)
        offset += l
    return result


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


def sample_greedy_warped(logits, temperature, top_k, top_p):
    warped = logits_warper(logits, temperature, top_k, top_p)
    return warped.argmax(dim=-1)


def prepare_labels(input_ids, prompt_lengths, ignore_index):
    labels = shift_labels(input_ids, ignore_index)
    seq = input_ids.shape[-1]
    pos = torch.arange(seq, device=input_ids.device)
    mask = pos < (prompt_lengths.unsqueeze(-1) - 1)
    labels[mask] = ignore_index
    return labels


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


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


def first_stop_index(input_ids, stop_ids):
    L = input_ids.shape[-1]
    pos = torch.arange(L, device=input_ids.device)
    isstop = (input_ids.unsqueeze(-1) == stop_ids).any(dim=-1)
    cand = torch.where(isstop, pos, torch.full_like(input_ids, L))
    first = cand.min(dim=-1).values
    return torch.where(first == L, torch.full_like(first, -1), first)


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


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


def right_pad_to(input_ids, pad_id, width):
    seq = input_ids.shape[-1]
    if width <= seq:
        return input_ids[..., :width]
    shape = list(input_ids.shape)
    shape[-1] = width - seq
    padding = torch.full(shape, pad_id, dtype=input_ids.dtype, device=input_ids.device)
    return torch.cat([input_ids, padding], dim=-1)


def generation_stop_mask(input_ids, eos_id, min_len):
    seq = input_ids.shape[-1]
    pos = torch.arange(seq, device=input_ids.device)
    iseos = input_ids == eos_id
    cand = torch.where(iseos, pos, torch.full_like(input_ids, seq))
    first = cand.min(dim=-1).values
    has = first < seq
    eos_len = torch.clamp(first + 1, min=min_len)
    valid = torch.where(has, eos_len, torch.full_like(first, seq))
    return (pos < valid.unsqueeze(-1)).float()

Test primitive functions vs reference

shell

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

def cmp(name, a, b, atol=1e-5, rtol=1e-4):
    ta = torch.as_tensor(a); tb=torch.as_tensor(b)
    if ta.shape!=tb.shape:
        print('FAIL',name,'shape',ta.shape,tb.shape); return
    if ta.dtype!=tb.dtype:
        print('DTYPE',name,ta.dtype,tb.dtype)
    ok=torch.allclose(ta.float(),tb.float(),atol=atol,rtol=rtol,equal_nan=True)
    print(('ok  ' if ok else 'FAIL')+' '+name, '' if ok else (ta-tb).abs().max().item())

g=torch.Generator().manual_seed(100)
def L(*s): return torch.randn(*s,generator=g)

# primitives
for seed in range(5):
    gg=torch.Generator().manual_seed(seed)
    lo=torch.randn(3,7,generator=gg)
    cmp('greedy',O.greedy_token(lo),R.greedy_token(lo))
    cmp('temp2',O.temperature_scale(lo,2.0),R.temperature_scale(lo,2.0))
    cmp('temp0',O.temperature_scale(lo,0.0),R.temperature_scale(lo,0.0))
    for k in [1,2,3,7,10]:
        cmp('topk%d'%k,O.top_k_filter(lo,k),R.top_k_filter(lo,k))
    for p in [0.1,0.3,0.5,0.8,0.95,1.0]:
        cmp('topp%.2f'%p,O.top_p_filter(lo,p),R.top_p_filter(lo,p))
        cmp('minp%.2f'%p,O.min_p_filter(lo,p),R.min_p_filter(lo,p))
    for e in [0.01,0.05,0.1,0.2]:
        cmp('eps%.2f'%e,O.epsilon_filter(lo,e),R.epsilon_filter(lo,e))
    cmp('renorm',O.renormalize_logits(lo),R.renormalize_logits(lo))
    cmp('entropy',O.entropy_from_logits(lo),R.entropy_from_logits(lo))
print('primitives batch done')
" 2>&1 | grep -v '^ok' | head -50

output

primitives batch done

[stdout]
primitives batch done

Test penalties and masking

shell

python3 -c "
import torch
import reforacle as R
import objectives as O

def cmp(name, a, b, atol=1e-5, rtol=1e-4):
    ta = torch.as_tensor(a); tb=torch.as_tensor(b)
    if ta.shape!=tb.shape:
        print('FAIL',name,'shape',ta.shape,tb.shape); return
    if ta.dtype!=tb.dtype:
        print('DTYPE',name,ta.dtype,tb.dtype)
    ok=torch.allclose(ta.float(),tb.float(),atol=atol,rtol=rtol,equal_nan=True)
    print(('ok  ' if ok else 'FAIL')+' '+name, '' if ok else (ta-tb).abs().max().item())

for seed in range(6):
    gg=torch.Generator().manual_seed(seed+50)
    lo=torch.randn(3,8,generator=gg)
    ids=torch.randint(0,8,(3,5),generator=gg)
    for pen in [1.0,1.2,1.5,0.8]:
        cmp('reppen%.1f'%pen,O.repetition_penalty(lo,ids,pen),R.repetition_penalty(lo,ids,pen),1e-6,1e-5)
    for pen in [0.0,0.5,1.0,2.0]:
        cmp('freq%.1f'%pen,O.frequency_penalty(lo,ids,pen),R.frequency_penalty(lo,ids,pen),1e-6,1e-5)
        cmp('pres%.1f'%pen,O.presence_penalty(lo,ids,pen),R.presence_penalty(lo,ids,pen),1e-6,1e-5)
    cmp('suppress',O.suppress_tokens(lo,[0,3,7]),R.suppress_tokens(lo,[0,3,7]))
    for cl,ml in [(2,5),(5,5),(6,5),(0,3)]:
        cmp('minlen%d_%d'%(cl,ml),O.min_length_filter(lo,cl,ml,2),R.min_length_filter(lo,cl,ml,2))

# masking / collation with pads
for seed in range(6):
    gg=torch.Generator().manual_seed(seed+200)
    ids=torch.randint(1,9,(3,6),generator=gg)
    # inject pads at random
    m=torch.rand(3,6,generator=gg)<0.4
    ids=ids.masked_fill(m,0)
    cmp('attn',O.attention_mask(ids,0),R.attention_mask(ids,0))
    cmp('posid',O.position_ids(ids,0),R.position_ids(ids,0))
    cmp('seqlen',O.sequence_lengths(ids,0),R.sequence_lengths(ids,0))
    cmp('maskeos',O.mask_after_eos(ids,3),R.mask_after_eos(ids,3))
    cmp('lasttok',O.last_token_indices(O.attention_mask(ids,0)),R.last_token_indices(R.attention_mask(ids,0)))
    lo=torch.randn(3,6,4,generator=gg)
    cmp('gatherlast',O.gather_last_logits(lo,O.attention_mask(ids,0)),R.gather_last_logits(lo,R.attention_mask(ids,0)))
    cmp('rollout',O.rollout_last_logits(lo,ids,0),R.rollout_last_logits(lo,ids,0))
    cmp('genstop0',O.generation_stop_mask(ids,3,0),R.generation_stop_mask(ids,3,0))
    cmp('genstop3',O.generation_stop_mask(ids,3,3),R.generation_stop_mask(ids,3,3))
    for t in [3,5]:
        cmp('causal%d'%t,O.causal_mask(t),R.causal_mask(t))
print('done')
" 2>&1 | grep -v '^ok' | head -60

output

done

[stdout]
done

Test labels, padding, RL helpers, composites

shell

python3 -c "
import torch
import reforacle as R
import objectives as O

def cmp(name, a, b, atol=1e-5, rtol=1e-4):
    ta = torch.as_tensor(a); tb=torch.as_tensor(b)
    if ta.shape!=tb.shape:
        print('FAIL',name,'shape',ta.shape,tb.shape); return
    if ta.dtype!=tb.dtype:
        print('DTYPE',name,ta.dtype,tb.dtype)
    ok=torch.allclose(ta.float(),tb.float(),atol=atol,rtol=rtol,equal_nan=True)
    print(('ok  ' if ok else 'FAIL')+' '+name, '' if ok else (ta-tb).abs().max().item())

for seed in range(6):
    gg=torch.Generator().manual_seed(seed+300)
    ids=torch.randint(0,9,(3,7),generator=gg)
    pl=torch.randint(1,6,(3,),generator=gg)
    cmp('shift',O.shift_labels(ids,-100),R.shift_labels(ids,-100))
    cmp('complabels',O.completion_loss_labels(ids,pl,-100),R.completion_loss_labels(ids,pl,-100))
    cmp('preplabels',O.prepare_labels(ids,pl,-100),R.prepare_labels(ids,pl,-100))
    for ml in [1,3,5,7,10]:
        cmp('trunc%d'%ml,O.truncate_left(ids,ml),R.truncate_left(ids,ml))
    for mult in [2,3,4,7,8]:
        cmp('padmult%d'%mult,O.pad_to_multiple(ids,0,mult),R.pad_to_multiple(ids,0,mult))
    for w in [3,7,10]:
        cmp('rightpad%d'%w,O.right_pad_to(ids,0,w),R.right_pad_to(ids,0,w))
    stop=torch.tensor([2,5])
    cmp('firststop',O.first_stop_index(ids,stop),R.first_stop_index(ids,stop))

for seed in range(6):
    gg=torch.Generator().manual_seed(seed+400)
    lo=torch.randn(3,4,6,generator=gg)
    act=torch.randint(0,6,(3,4),generator=gg)
    cmp('actlogp',O.action_logprobs(lo,act),R.action_logprobs(lo,act))
    lo2=torch.randn(3,6,generator=gg)
    mask=(torch.rand(3,6,generator=gg)>0.3).float()
    cmp('maskargmax',O.masked_argmax(lo2,mask),R.masked_argmax(lo2,mask))
    for k in [1,2,4,6]:
        cmp('topkidx%d'%k,O.topk_indices(lo2,k),R.topk_indices(lo2,k))
    cmp('warper',O.logits_warper(lo2,1.5,3,0.9),R.logits_warper(lo2,1.5,3,0.9))
    cmp('samplegw',O.sample_greedy_warped(lo2,1.5,3,0.9),R.sample_greedy_warped(lo2,1.5,3,0.9))
    rew=torch.randn(3,5,generator=gg); rmask=(torch.rand(3,5,generator=gg)>0.3).float()
    cmp('r2go',O.reward_to_go(rew,rmask),R.reward_to_go(rew,rmask))
    for gm in [0.9,0.99,1.0,0.5]:
        cmp('disc%.2f'%gm,O.discount_cumsum(rew,gm),R.discount_cumsum(rew,gm))
    sl=torch.randint(1,4,(3,),generator=gg)
    tot=int(sl.sum())+2
    cmp('buildpos',O.build_position_ids_packed(sl,tot),R.build_position_ids_packed(sl,tot))
print('done')
" 2>&1 | grep -v '^ok' | head -60

output

FAIL r2go 4.238382339477539
FAIL r2go 1.8694974184036255
FAIL r2go 1.7696998119354248
FAIL r2go 2.6005287170410156
FAIL r2go 0.44666120409965515
FAIL r2go 1.0062823295593262
done

[stdout]
FAIL r2go 4.238382339477539
FAIL r2go 1.8694974184036255
FAIL r2go 1.7696998119354248
FAIL r2go 2.6005287170410156
FAIL r2go 0.44666120409965515
FAIL r2go 1.0062823295593262
done

Disambiguate reward_to_go with interior zeros

shell

python3 -c "
import torch
import reforacle as R
rew=torch.tensor([[1.,2,3,4,5]])
mask=torch.tensor([[1.,0,1,0,1]])
print('ref',R.reward_to_go(rew,mask).tolist())
# hypotheses
rm=rew*mask
print('revcumsum masked',rm.flip(-1).cumsum(-1).flip(-1).tolist())
print('revcumsum then *mask',(rew.flip(-1).cumsum(-1).flip(-1)*mask).tolist())
print('revcumsum unmasked',rew.flip(-1).cumsum(-1).flip(-1).tolist())
"

output

ref [[9.0, 0.0, 8.0, 0.0, 5.0]]
revcumsum masked [[9.0, 8.0, 8.0, 5.0, 5.0]]
revcumsum then *mask [[15.0, 0.0, 12.0, 0.0, 5.0]]
revcumsum unmasked [[15.0, 14.0, 12.0, 9.0, 5.0]]

[stdout]
ref [[9.0, 0.0, 8.0, 0.0, 5.0]]
revcumsum masked [[9.0, 8.0, 8.0, 5.0, 5.0]]
revcumsum then *mask [[15.0, 0.0, 12.0, 0.0, 5.0]]
revcumsum unmasked [[15.0, 14.0, 12.0, 9.0, 5.0]]

/app/objectives.py

− removed

    rm = rewards * mask
    return rm.flip(-1).cumsum(dim=-1).flip(-1)

+ added

    rm = rewards * mask
    return rm.flip(-1).cumsum(dim=-1).flip(-1) * mask

Retest reward_to_go

shell

python3 -c "
import torch
import reforacle as R
import objectives as O
def cmp(name, a, b, atol=1e-5, rtol=1e-4):
    ta = torch.as_tensor(a); tb=torch.as_tensor(b)
    if ta.shape!=tb.shape: print('FAIL',name,'shape'); return
    ok=torch.allclose(ta.float(),tb.float(),atol=atol,rtol=rtol,equal_nan=True) and ta.dtype==tb.dtype
    if not ok: print('FAIL',name)
for seed in range(8):
    gg=torch.Generator().manual_seed(seed+900)
    rew=torch.randn(4,6,generator=gg); rmask=(torch.rand(4,6,generator=gg)>0.4).float()
    cmp('r2go',O.reward_to_go(rew,rmask),R.reward_to_go(rew,rmask))
print('r2go retest done')
"

output

r2go retest done

[stdout]
r2go retest done

Full comprehensive sweep of all 37 functions

shell

python3 -c "
import torch
import reforacle as R
import objectives as O
fails=set()
def cmp(name, a, b, comp=False):
    atol,rtol=(1e-5,1e-4) if comp else (1e-6,1e-5)
    ta=torch.as_tensor(a); tb=torch.as_tensor(b)
    if ta.shape!=tb.shape or ta.dtype!=tb.dtype or not torch.allclose(ta.float(),tb.float(),atol=atol,rtol=rtol,equal_nan=True):
        fails.add(name); print('FAIL',name, ta.shape,tb.shape,ta.dtype,tb.dtype)

for seed in range(20):
    gg=torch.Generator().manual_seed(seed)
    B=int(torch.randint(1,5,(1,),generator=gg)); V=int(torch.randint(3,12,(1,),generator=gg)); S=int(torch.randint(2,8,(1,),generator=gg))
    lo=torch.randn(B,V,generator=gg)
    ids=torch.randint(0,V,(B,S),generator=gg)
    idp=ids.masked_fill(torch.rand(B,S,generator=gg)<0.4,0)
    cmp('greedy_token',O.greedy_token(lo),R.greedy_token(lo))
    for T in [0.0,0.5,1.0,2.0]: cmp('temperature_scale',O.temperature_scale(lo,T),R.temperature_scale(lo,T))
    for k in range(1,V+2): cmp('top_k_filter',O.top_k_filter(lo,k),R.top_k_filter(lo,k))
    for p in [0.0,0.15,0.4,0.7,0.99,1.0]:
        cmp('top_p_filter',O.top_p_filter(lo,p),R.top_p_filter(lo,p))
        cmp('min_p_filter',O.min_p_filter(lo,p),R.min_p_filter(lo,p))
        cmp('epsilon_filter',O.epsilon_filter(lo,p),R.epsilon_filter(lo,p))
    for pen in [0.7,1.0,1.3,1.8]: cmp('repetition_penalty',O.repetition_penalty(lo,ids,pen),R.repetition_penalty(lo,ids,pen))
    for pen in [0.0,0.5,1.5]:
        cmp('frequency_penalty',O.frequency_penalty(lo,ids,pen),R.frequency_penalty(lo,ids,pen))
        cmp('presence_penalty',O.presence_penalty(lo,ids,pen),R.presence_penalty(lo,ids,pen))
    cmp('suppress_tokens',O.suppress_tokens(lo,[0,min(1,V-1)]),R.suppress_tokens(lo,[0,min(1,V-1)]))
    for cl in [0,S,S+2]: cmp('min_length_filter',O.min_length_filter(lo,cl,S,min(2,V-1)),R.min_length_filter(lo,cl,S,min(2,V-1)))
    cmp('renormalize_logits',O.renormalize_logits(lo),R.renormalize_logits(lo))
    cmp('attention_mask',O.attention_mask(idp,0),R.attention_mask(idp,0))
    cmp('position_ids',O.position_ids(idp,0),R.position_ids(idp,0))
    cmp('causal_mask',O.causal_mask(S),R.causal_mask(S))
    cmp('shift_labels',O.shift_labels(ids,-100),R.shift_labels(ids,-100))
    pl=torch.randint(1,S+1,(B,),generator=gg)
    cmp('completion_loss_labels',O.completion_loss_labels(ids,pl,-100),R.completion_loss_labels(ids,pl,-100))
    cmp('prepare_labels',O.prepare_labels(ids,pl,-100),R.prepare_labels(ids,pl,-100))
    cmp('mask_after_eos',O.mask_after_eos(idp,3),R.mask_after_eos(idp,3))
    cmp('sequence_lengths',O.sequence_lengths(idp,0),R.sequence_lengths(idp,0))
    am=O.attention_mask(idp,0)
    cmp('last_token_indices',O.last_token_indices(am),R.last_token_indices(am))
    lo3=torch.randn(B,S,V,generator=gg)
    cmp('gather_last_logits',O.gather_last_logits(lo3,am),R.gather_last_logits(lo3,am))
    for ml in [1,S,S+3]: cmp('truncate_left',O.truncate_left(ids,ml),R.truncate_left(ids,ml))
    for mult in [2,3,4,8]: cmp('pad_to_multiple',O.pad_to_multiple(ids,0,mult),R.pad_to_multiple(ids,0,mult))
    for w in [1,S,S+4]: cmp('right_pad_to',O.right_pad_to(ids,0,w),R.right_pad_to(ids,0,w))
    act=torch.randint(0,V,(B,S),generator=gg)
    cmp('action_logprobs',O.action_logprobs(lo3,act),R.action_logprobs(lo3,act),True)
    cmp('entropy_from_logits',O.entropy_from_logits(lo),R.entropy_from_logits(lo),True)
    rew=torch.randn(B,S,generator=gg); rmask=(torch.rand(B,S,generator=gg)>0.4).float()
    cmp('reward_to_go',O.reward_to_go(rew,rmask),R.reward_to_go(rew,rmask),True)
    for gm in [0.5,0.9,1.0]: cmp('discount_cumsum',O.discount_cumsum(rew,gm),R.discount_cumsum(rew,gm),True)
    sl=torch.randint(1,4,(B,),generator=gg); tot=int(sl.sum())+int(torch.randint(0,3,(1,),generator=gg))
    cmp('build_position_ids_packed',O.build_position_ids_packed(sl,tot),R.build_position_ids_packed(sl,tot))
    for T in [0.7,1.0,2.0]:
      for k in range(1,V+1):
        for p in [0.3,0.7,1.0]:
            cmp('logits_warper',O.logits_warper(lo,T,k,p),R.logits_warper(lo,T,k,p),True)
            cmp('sample_greedy_warped',O.sample_greedy_warped(lo,T,k,p),R.sample_greedy_warped(lo,T,k,p),True)
    cmp('rollout_last_logits',O.rollout_last_logits(lo3,idp,0),R.rollout_last_logits(lo3,idp,0))
    stop=torch.tensor([2,3,min(4,V-1)])
    cmp('first_stop_index',O.first_stop_index(ids,stop),R.first_stop_index(ids,stop))
    mk=(torch.rand(B,V,generator=gg)>0.3).float()
    if mk.sum()>0: cmp('masked_argmax',O.masked_argmax(lo,mk),R.masked_argmax(lo,mk))
    for k in range(1,V+1): cmp('topk_indices',O.topk_indices(lo,k),R.topk_indices(lo,k))
    for ml in [0,2,S]: cmp('generation_stop_mask',O.generation_stop_mask(idp,3,ml),R.generation_stop_mask(idp,3,ml))
print('TOTAL FAILS:', sorted(fails) if fails else 'NONE - ALL PASS')
"

output

FAIL epsilon_filter torch.Size([1, 3]) torch.Size([1, 3]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 3]) torch.Size([1, 3]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 3]) torch.Size([1, 3]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 3]) torch.Size([1, 3]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL gather_last_logits torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL rollout_last_logits torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 9]) torch.Size([1, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 9]) torch.Size([1, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 9]) torch.Size([1, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 9]) torch.Size([1, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 8]) torch.Size([3, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 8]) torch.Size([3, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 8]) torch.Size([3, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 8]) torch.Size([3, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 6]) torch.Size([3, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 6]) torch.Size([3, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 6]) torch.Size([3, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 6]) torch.Size([3, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL gather_last_logits torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL rollout_last_logits torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 9]) torch.Size([3, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 9]) torch.Size([3, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 9]) torch.Size([3, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 9]) torch.Size([3, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL gather_last_logits torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL rollout_last_logits torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 6]) torch.Size([2, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 6]) torch.Size([2, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 6]) torch.Size([2, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL gather_last_logits torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL rollout_last_logits torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 6]) torch.Size([1, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 6]) torch.Size([1, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 6]) torch.Size([1, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 6]) torch.Size([1, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 11]) torch.Size([2, 11]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 11]) torch.Size([2, 11]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 11]) torch.Size([2, 11]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 11]) torch.Size([2, 11]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 9]) torch.Size([4, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 9]) torch.Size([4, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 9]) torch.Size([4, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 9]) torch.Size([4, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 10]) torch.Size([2, 10]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 10]) torch.Size([2, 10]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 10]) torch.Size([2, 10]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 10]) torch.Size([2, 10]) torch.float32 torch.float32
TOTAL FAILS: ['epsilon_filter', 'gather_last_logits', 'rollout_last_logits']

[stdout]
FAIL epsilon_filter torch.Size([1, 3]) torch.Size([1, 3]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 3]) torch.Size([1, 3]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 3]) torch.Size([1, 3]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 3]) torch.Size([1, 3]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL gather_last_logits torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL rollout_last_logits torch.Size([2, 8]) torch.Size([2, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 9]) torch.Size([1, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 9]) torch.Size([1, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 9]) torch.Size([1, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 9]) torch.Size([1, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 8]) torch.Size([3, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 8]) torch.Size([3, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 8]) torch.Size([3, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 8]) torch.Size([3, 8]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 6]) torch.Size([3, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 6]) torch.Size([3, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 6]) torch.Size([3, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 6]) torch.Size([3, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL gather_last_logits torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL rollout_last_logits torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 9]) torch.Size([3, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 9]) torch.Size([3, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 9]) torch.Size([3, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 9]) torch.Size([3, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL gather_last_logits torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL rollout_last_logits torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 6]) torch.Size([2, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 6]) torch.Size([2, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 6]) torch.Size([2, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL gather_last_logits torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL rollout_last_logits torch.Size([2, 9]) torch.Size([2, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 7]) torch.Size([4, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 6]) torch.Size([4, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 6]) torch.Size([1, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 6]) torch.Size([1, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 6]) torch.Size([1, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([1, 6]) torch.Size([1, 6]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 11]) torch.Size([2, 11]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 11]) torch.Size([2, 11]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 11]) torch.Size([2, 11]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 11]) torch.Size([2, 11]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 9]) torch.Size([4, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 9]) torch.Size([4, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 9]) torch.Size([4, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([4, 9]) torch.Size([4, 9]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([3, 7]) torch.Size([3, 7]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 10]) torch.Size([2, 10]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 10]) torch.Size([2, 10]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 10]) torch.Size([2, 10]) torch.float32 torch.float32
FAIL epsilon_filter torch.Size([2, 10]) torch.Size([2, 10]) torch.float32 torch.float32
TOTAL FAILS: ['epsilon_filter', 'gather_last_logits', 'rollout_last_logits']

Investigate epsilon and all-pad edge cases

shell

python3 -c "
import torch
import reforacle as R
lo=torch.randn(1,6)
for e in [0.0,0.15,0.4,0.7,0.99,1.0]:
    r=R.epsilon_filter(lo,e)
    print('eps',e,'ninf',int(torch.isinf(r).sum()),'of',lo.numel())
print('logits',lo.tolist())
# all-pad mask
lo3=torch.randn(2,4,3)
am=torch.tensor([[1.,1,0,0],[0,0,0,0]])
print('ref last_idx',R.last_token_indices(am).tolist())
print('ref gather row1', R.gather_last_logits(lo3,am)[1].tolist())
print('logits row1 pos0',lo3[1,0].tolist(),'pos3',lo3[1,3].tolist())
"

output

eps 0.0 ninf 0 of 6
eps 0.15 ninf 4 of 6
eps 0.4 ninf 5 of 6
eps 0.7 ninf 5 of 6
eps 0.99 ninf 5 of 6
eps 1.0 ninf 5 of 6
logits [[-2.2268707752227783, -1.978712558746338, -0.9505015015602112, 1.2946971654891968, 0.1624869704246521, -0.6946415901184082]]
ref last_idx [1, -1]
ref gather row1 [-0.6713274121284485, 1.022416114807129, -0.5417229533195496]
logits row1 pos0 [-0.6713274121284485, 1.022416114807129, -0.5417229533195496] pos3 [0.882072389125824, 1.5909405946731567, 1.3524227142333984]

[stdout]
eps 0.0 ninf 0 of 6
eps 0.15 ninf 4 of 6
eps 0.4 ninf 5 of 6
eps 0.7 ninf 5 of 6
eps 0.99 ninf 5 of 6
eps 1.0 ninf 5 of 6
logits [[-2.2268707752227783, -1.978712558746338, -0.9505015015602112, 1.2946971654891968, 0.1624869704246521, -0.6946415901184082]]
ref last_idx [1, -1]
ref gather row1 [-0.6713274121284485, 1.022416114807129, -0.5417229533195496]
logits row1 pos0 [-0.6713274121284485, 1.022416114807129, -0.5417229533195496] pos3 [0.882072389125824, 1.5909405946731567, 1.3524227142333984]

/app/objectives.py

− removed

    idx = last_token_indices(attn_mask)
    return logits[torch.arange(logits.shape[0], device=logits.device), idx]

+ added

    idx = last_token_indices(attn_mask).clamp(min=0)
    return logits[torch.arange(logits.shape[0], device=logits.device), idx]

/app/objectives.py

− removed

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

+ added

    probs = logits.softmax(dim=-1)
    top = probs.max(dim=-1, keepdim=True).values
    remove = (probs < epsilon) & (probs < top)
    return logits.masked_fill(remove, NEG_INF)

Full comprehensive sweep, 40 seeds

shell

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

for seed in range(40):
    gg=torch.Generator().manual_seed(seed)
    B=int(torch.randint(1,5,(1,),generator=gg)); V=int(torch.randint(3,13,(1,),generator=gg)); S=int(torch.randint(2,9,(1,),generator=gg))
    lo=torch.randn(B,V,generator=gg)
    ids=torch.randint(0,V,(B,S),generator=gg)
    idp=ids.masked_fill(torch.rand(B,S,generator=gg)<0.5,0)
    cmp('greedy_token',O.greedy_token(lo),R.greedy_token(lo))
    for T in [0.0,0.5,1.0,2.0,3.7]: cmp('temperature_scale',O.temperature_scale(lo,T),R.temperature_scale(lo,T))
    for k in range(1,V+3): cmp('top_k_filter',O.top_k_filter(lo,k),R.top_k_filter(lo,k))
    for p in [0.0,0.15,0.4,0.7,0.9,0.99,1.0]:
        cmp('top_p_filter',O.top_p_filter(lo,p),R.top_p_filter(lo,p))
        cmp('min_p_filter',O.min_p_filter(lo,p),R.min_p_filter(lo,p))
        cmp('epsilon_filter',O.epsilon_filter(lo,p),R.epsilon_filter(lo,p))
    for pen in [0.5,0.7,1.0,1.3,1.8,2.5]: cmp('repetition_penalty',O.repetition_penalty(lo,ids,pen),R.repetition_penalty(lo,ids,pen))
    for pen in [0.0,0.5,1.5,3.0]:
        cmp('frequency_penalty',O.frequency_penalty(lo,ids,pen),R.frequency_penalty(lo,ids,pen))
        cmp('presence_penalty',O.presence_penalty(lo,ids,pen),R.presence_penalty(lo,ids,pen))
    cmp('suppress_tokens',O.suppress_tokens(lo,[0,min(1,V-1)]),R.suppress_tokens(lo,[0,min(1,V-1)]))
    for cl in [0,S,S+2]: cmp('min_length_filter',O.min_length_filter(lo,cl,S,min(2,V-1)),R.min_length_filter(lo,cl,S,min(2,V-1)))
    cmp('renormalize_logits',O.renormalize_logits(lo),R.renormalize_logits(lo))
    cmp('attention_mask',O.attention_mask(idp,0),R.attention_mask(idp,0))
    cmp('position_ids',O.position_ids(idp,0),R.position_ids(idp,0))
    cmp('causal_mask',O.causal_mask(S),R.causal_mask(S))
    cmp('shift_labels',O.shift_labels(ids,-100),R.shift_labels(ids,-100))
    pl=torch.randint(1,S+1,(B,),generator=gg)
    cmp('completion_loss_labels',O.completion_loss_labels(ids,pl,-100),R.completion_loss_labels(ids,pl,-100))
    cmp('prepare_labels',O.prepare_labels(ids,pl,-100),R.prepare_labels(ids,pl,-100))
    cmp('mask_after_eos',O.mask_after_eos(idp,3),R.mask_after_eos(idp,3))
    cmp('sequence_lengths',O.sequence_lengths(idp,0),R.sequence_lengths(idp,0))
    am=O.attention_mask(idp,0)
    cmp('last_token_indices',O.last_token_indices(am),R.last_token_indices(am))
    lo3=torch.randn(B,S,V,generator=gg)
    cmp('gather_last_logits',O.gather_last_logits(lo3,am),R.gather_last_logits(lo3,am))
    for ml in [1,S,S+3]: cmp('truncate_left',O.truncate_left(ids,ml),R.truncate_left(ids,ml))
    for mult in [1,2,3,4,8]: cmp('pad_to_multiple',O.pad_to_multiple(ids,0,mult),R.pad_to_multiple(ids,0,mult))
    for w in [1,S,S+4]: cmp('right_pad_to',O.right_pad_to(ids,0,w),R.right_pad_to(ids,0,w))
    act=torch.randint(0,V,(B,S),generator=gg)
    cmp('action_logprobs',O.action_logprobs(lo3,act),R.action_logprobs(lo3,act),True)
    cmp('entropy_from_logits',O.entropy_from_logits(lo),R.entropy_from_logits(lo),True)
    rew=torch.randn(B,S,generator=gg); rmask=(torch.rand(B,S,generator=gg)>0.4).float()
    cmp('reward_to_go',O.reward_to_go(rew,rmask),R.reward_to_go(rew,rmask),True)
    for gm in [0.5,0.9,1.0]: cmp('discount_cumsum',O.discount_cumsum(rew,gm),R.discount_cumsum(rew,gm),True)
    sl=torch.randint(1,4,(B,),generator=gg); tot=int(sl.sum())+int(torch.randint(0,3,(1,),generator=gg))
    cmp('build_position_ids_packed',O.build_position_ids_packed(sl,tot),R.build_position_ids_packed(sl,tot))
    for T in [0.7,1.0,2.0]:
      for k in range(1,V+1):
        for p in [0.3,0.7,1.0]:
            cmp('logits_warper',O.logits_warper(lo,T,k,p),R.logits_warper(lo,T,k,p),True)
            cmp('sample_greedy_warped',O.sample_greedy_warped(lo,T,k,p),R.sample_greedy_warped(lo,T,k,p),True)
    cmp('rollout_last_logits',O.rollout_last_logits(lo3,idp,0),R.rollout_last_logits(lo3,idp,0))
    stop=torch.tensor([2,3,min(4,V-1)])
    cmp('first_stop_index',O.first_stop_index(ids,stop),R.first_stop_index(ids,stop))
    mk=(torch.rand(B,V,generator=gg)>0.3).float()
    cmp('masked_argmax',O.masked_argmax(lo,mk),R.masked_argmax(lo,mk))
    for k in range(1,V+1): cmp('topk_indices',O.topk_indices(lo,k),R.topk_indices(lo,k))
    for ml in [0,2,S,S+2]: cmp('generation_stop_mask',O.generation_stop_mask(idp,3,ml),R.generation_stop_mask(idp,3,ml))
print('TOTAL FAILS:', fails if fails else 'NONE - ALL 37 PASS')
"

output

TOTAL FAILS: NONE - ALL 37 PASS

[stdout]
TOTAL FAILS: NONE - ALL 37 PASS

Verify standalone in clean sandbox

shell

cd /tmp && cp /app/objectives.py /tmp/obj_test.py && python3 -c "
import sys
# simulate clean sandbox: block reforacle
import builtins
_orig=builtins.__import__
def guard(name,*a,**k):
    if name=='reforacle': raise ImportError('reforacle not available')
    return _orig(name,*a,**k)
builtins.__import__=guard
import obj_test as O
import torch
# spot checks + edge: all-zero masked_argmax
print('imports clean, no reforacle dependency')
lo=torch.randn(2,5); mk=torch.zeros(2,5)
print('masked_argmax all-zero:', O.masked_argmax(lo,mk).tolist())
print('NotImplementedError present:', 'NotImplementedError' in open('/tmp/obj_test.py').read())
"

output

imports clean, no reforacle dependency
masked_argmax all-zero: [0, 0]
NotImplementedError present: False
Shell cwd was reset to /app

[stdout]
imports clean, no reforacle dependency
masked_argmax all-zero: [0, 0]
NotImplementedError present: False
[stderr]

Shell cwd was reset to /app

Confirm all-zero masked_argmax matches ref, check budget

shell

python3 -c "
import torch, reforacle as R, objectives as O
lo=torch.randn(2,5); mk=torch.zeros(2,5)
print('ref',R.masked_argmax(lo,mk).tolist(),'mine',O.masked_argmax(lo,mk).tolist())
"; bash /app/timer.sh

output

ref [4, 4] mine [0, 0]
Remaining time (hours:minutes):
1:50

[stdout]
ref [4, 4] mine [0, 0]
Remaining time (hours:minutes):
1:50

Understand masked_argmax masking convention

shell

python3 -c "
import torch, reforacle as R
g=torch.Generator().manual_seed(0)
lo=torch.randn(3,6,generator=g)
print('logits argmax', lo.argmax(-1).tolist())
mk=torch.zeros(3,6)
print('ref all-zero', R.masked_argmax(lo,mk).tolist())
# partial masks
mk2=torch.tensor([[0.,0,1,0,0,0],[1.,1,0,0,0,0],[0.,0,0,0,0,1]])
print('ref partial', R.masked_argmax(lo,mk2).tolist())
# my additive approach test
NEG=-1e9
print('additive', (lo+(mk-1)*1e9).argmax(-1).tolist(), 'partial', (lo+(mk2-1)*1e9).argmax(-1).tolist())
print('masked_fill approach partial', lo.masked_fill(mk2==0,float('-inf')).argmax(-1).tolist())
"

output

logits argmax [3, 1, 1]
ref all-zero [5, 5, 5]
ref partial [2, 1, 5]
additive [0, 0, 0] partial [2, 1, 5]
masked_fill approach partial [2, 1, 5]

[stdout]
logits argmax [3, 1, 1]
ref all-zero [5, 5, 5]
ref partial [2, 1, 5]
additive [0, 0, 0] partial [2, 1, 5]
masked_fill approach partial [2, 1, 5]

Check argmax tie-breaking on -inf and real ties

shell

python3 -c "
import torch, reforacle as R
print('argmax all -inf 6:', torch.full((6,),float('-inf')).argmax(-1).item())
print('argmax all -inf 2x6:', torch.full((2,6),float('-inf')).argmax(-1).tolist())
# tie in real values with mask
lo=torch.tensor([[1.,3.,3.,1.,3.,0.]])
mk=torch.tensor([[1.,1.,1.,1.,1.,1.]])
print('ref tie',R.masked_argmax(lo,mk).tolist(),'torch',lo.argmax(-1).tolist())
mk2=torch.tensor([[0.,0.,0.,0.,0.,0.]])
print('ref allzero tie', R.masked_argmax(lo,mk2).tolist())
"

output

argmax all -inf 6: 0
argmax all -inf 2x6: [0, 0]
ref tie [4] torch [1]
ref allzero tie [5]

[stdout]
argmax all -inf 6: 0
argmax all -inf 2x6: [0, 0]
ref tie [4] torch [1]
ref allzero tie [5]

Check greedy/sample tie-breaking

shell

python3 -c "
import torch, reforacle as R
lo=torch.tensor([[1.,3.,3.,1.,3.,0.],[5.,5.,0.,0.,0.,5.]])
print('greedy ref',R.greedy_token(lo).tolist(),'torch',lo.argmax(-1).tolist())
# sample_greedy_warped with ties (temp keeps ties, high top_k/p keep all)
print('samplegw ref',R.sample_greedy_warped(lo,1.0,6,1.0).tolist())
# topk_indices ties
print('topk_idx ref',R.topk_indices(lo,3).tolist())
# min_length uses nothing. first_stop uses min. 
"

output

greedy ref [4, 5] torch [1, 0]
samplegw ref [4, 5]
topk_idx ref [[1, 2, 4], [0, 1, 5]]

[stdout]
greedy ref [4, 5] torch [1, 0]
samplegw ref [4, 5]
topk_idx ref [[1, 2, 4], [0, 1, 5]]

Check temp0, topk_indices ties

shell

python3 -c "
import torch, reforacle as R, objectives as O
lo=torch.tensor([[1.,3.,3.,1.,3.,0.],[5.,5.,0.,0.,0.,5.]])
print('temp0 ref',R.temperature_scale(lo,0.0).tolist())
print('temp0 mine',O.temperature_scale(lo,0.0).tolist())
print('topk_idx ref',R.topk_indices(lo,3).tolist())
print('topk_idx mine',O.topk_indices(lo,3).tolist())
# top_k_filter tie: does the kth threshold matter (keeps >= kth so ties fine)
print('topk_filter ref',R.top_k_filter(lo,2).tolist())
print('topk_filter mine',O.top_k_filter(lo,2).tolist())
"

output

temp0 ref [[-inf, -inf, -inf, -inf, 0.0, -inf], [-inf, -inf, -inf, -inf, -inf, 0.0]]
temp0 mine [[-inf, 0.0, -inf, -inf, -inf, -inf], [0.0, -inf, -inf, -inf, -inf, -inf]]
topk_idx ref [[1, 2, 4], [0, 1, 5]]
topk_idx mine [[2, 4, 1], [1, 5, 0]]
topk_filter ref [[-inf, 3.0, 3.0, -inf, 3.0, -inf], [5.0, 5.0, -inf, -inf, -inf, 5.0]]
topk_filter mine [[-inf, 3.0, 3.0, -inf, 3.0, -inf], [5.0, 5.0, -inf, -inf, -inf, 5.0]]

[stdout]
temp0 ref [[-inf, -inf, -inf, -inf, 0.0, -inf], [-inf, -inf, -inf, -inf, -inf, 0.0]]
temp0 mine [[-inf, 0.0, -inf, -inf, -inf, -inf], [0.0, -inf, -inf, -inf, -inf, -inf]]
topk_idx ref [[1, 2, 4], [0, 1, 5]]
topk_idx mine [[2, 4, 1], [1, 5, 0]]
topk_filter ref [[-inf, 3.0, 3.0, -inf, 3.0, -inf], [5.0, 5.0, -inf, -inf, -inf, 5.0]]
topk_filter mine [[-inf, 3.0, 3.0, -inf, 3.0, -inf], [5.0, 5.0, -inf, -inf, -inf, 5.0]]

/app/objectives.py

− removed

NEG_INF = float("-inf")


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


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

+ added

NEG_INF = float("-inf")


def _last_argmax(x):
    L = x.shape[-1]
    return (L - 1) - x.flip(-1).argmax(dim=-1)


def greedy_token(logits):
    return _last_argmax(logits)


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

/app/objectives.py

− removed

    warped = logits_warper(logits, temperature, top_k, top_p)
    return warped.argmax(dim=-1)

+ added

    warped = logits_warper(logits, temperature, top_k, top_p)
    return _last_argmax(warped)

/app/objectives.py

− removed

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


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

+ added

    filled = logits.masked_fill(mask == 0, NEG_INF)
    return _last_argmax(filled)


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

Tie-heavy tests for argmax/topk convention

shell

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

# tie-heavy: integer logits
for seed in range(30):
    gg=torch.Generator().manual_seed(seed+7000)
    B=int(torch.randint(1,5,(1,),generator=gg)); V=int(torch.randint(3,10,(1,),generator=gg))
    lo=torch.randint(0,4,(B,V),generator=gg).float()  # lots of ties
    cmp('greedy_token',O.greedy_token(lo),R.greedy_token(lo))
    cmp('temperature_scale0',O.temperature_scale(lo,0.0),R.temperature_scale(lo,0.0))
    for k in range(1,V+1):
        cmp('topk_indices',O.topk_indices(lo,k),R.topk_indices(lo,k))
        cmp('top_k_filter',O.top_k_filter(lo,k),R.top_k_filter(lo,k))
    mk=(torch.randint(0,2,(B,V),generator=gg)).float()
    cmp('masked_argmax',O.masked_argmax(lo,mk),R.masked_argmax(lo,mk))
    for T in [1.0,2.0]:
      for k in range(1,V+1):
        for p in [0.5,1.0]:
            cmp('sample_greedy_warped',O.sample_greedy_warped(lo,T,k,p),R.sample_greedy_warped(lo,T,k,p),True)
            cmp('logits_warper',O.logits_warper(lo,T,k,p),R.logits_warper(lo,T,k,p),True)
print('tie tests:', fails if fails else 'ALL TIE CASES PASS')
"

output

tie tests: {'sample_greedy_warped': 120, 'logits_warper': 120}

[stdout]
tie tests: {'sample_greedy_warped': 120, 'logits_warper': 120}

Isolate top_p tie handling in warper

shell

python3 -c "
import torch, reforacle as R, objectives as O
gg=torch.Generator().manual_seed(7000)
lo=torch.randint(0,4,(2,6),generator=gg).float()
print('lo',lo.tolist())
for p in [0.3,0.5,0.7,1.0]:
    r=R.top_p_filter(lo,p); m=O.top_p_filter(lo,p)
    if not torch.equal(r,m): print('top_p',p,'DIFF'); print(' ref',r.tolist()); print(' mine',m.tolist())
# step through warper T=1,k=3,p=0.5
T,k,p=1.0,3,0.5
print('warper ref',R.logits_warper(lo,T,k,p).tolist())
print('warper mine',O.logits_warper(lo,T,k,p).tolist())
"

output

lo [[1.0, 1.0, 3.0, 1.0, 1.0, 3.0], [3.0, 1.0, 3.0, 3.0, 2.0, 3.0]]
warper ref [[-inf, -inf, 3.0, -inf, -inf, 3.0], [3.0, -inf, 3.0, -inf, -inf, -inf]]
warper mine [[-inf, -inf, 3.0, -inf, -inf, 3.0], [3.0, -inf, 3.0, 3.0, -inf, -inf]]

[stdout]
lo [[1.0, 1.0, 3.0, 1.0, 1.0, 3.0], [3.0, 1.0, 3.0, 3.0, 2.0, 3.0]]
warper ref [[-inf, -inf, 3.0, -inf, -inf, 3.0], [3.0, -inf, 3.0, -inf, -inf, -inf]]
warper mine [[-inf, -inf, 3.0, -inf, -inf, 3.0], [3.0, -inf, 3.0, 3.0, -inf, -inf]]

Determine top_p boundary > vs >= with equal probs

shell

python3 -c "
import torch, reforacle as R, objectives as O
# exact boundary: 4 equal logits
x=torch.tensor([[3.,float('-inf'),3.,3.,float('-inf'),3.]])
for p in [0.25,0.5,0.75,1.0]:
    print('p',p,'ref',R.top_p_filter(x,p).tolist())
print()
# simple 2-equal
y=torch.tensor([[0.,0.]])
for p in [0.5,0.4,0.6]:
    print('2eq p',p,'ref',R.top_p_filter(y,p).tolist())
"

output

p 0.25 ref [[3.0, -inf, -inf, -inf, -inf, -inf]]
p 0.5 ref [[3.0, -inf, 3.0, -inf, -inf, -inf]]
p 0.75 ref [[3.0, -inf, 3.0, 3.0, -inf, -inf]]
p 1.0 ref [[3.0, -inf, 3.0, 3.0, -inf, 3.0]]

2eq p 0.5 ref [[0.0, -inf]]
2eq p 0.4 ref [[0.0, -inf]]
2eq p 0.6 ref [[0.0, 0.0]]

[stdout]
p 0.25 ref [[3.0, -inf, -inf, -inf, -inf, -inf]]
p 0.5 ref [[3.0, -inf, 3.0, -inf, -inf, -inf]]
p 0.75 ref [[3.0, -inf, 3.0, 3.0, -inf, -inf]]
p 1.0 ref [[3.0, -inf, 3.0, 3.0, -inf, 3.0]]

2eq p 0.5 ref [[0.0, -inf]]
2eq p 0.4 ref [[0.0, -inf]]
2eq p 0.6 ref [[0.0, 0.0]]

/app/objectives.py

− removed

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

+ added

    sorted_logits, sorted_idx = torch.sort(logits, descending=True, dim=-1, stable=True)
    probs = sorted_logits.softmax(dim=-1)
    cum = probs.cumsum(dim=-1)
    remove = (cum - probs) >= p
    remove[..., 0] = False
    remove_orig = torch.zeros_like(remove).scatter(-1, sorted_idx, remove)
    return logits.masked_fill(remove_orig, NEG_INF)

Retest ties for top_p/warper

shell

python3 -c "
import torch, reforacle as R, objectives as O
fails={}
def cmp(name,a,b,comp=False):
    atol,rtol=(1e-5,1e-4) if comp else (1e-6,1e-5)
    ta=torch.as_tensor(a);tb=torch.as_tensor(b)
    if ta.shape!=tb.shape or ta.dtype!=tb.dtype or not torch.allclose(ta.float(),tb.float(),atol=atol,rtol=rtol,equal_nan=True):
        fails[name]=fails.get(name,0)+1
# tie-heavy + continuous mix
for seed in range(40):
    gg=torch.Generator().manual_seed(seed+7000)
    B=int(torch.randint(1,5,(1,),generator=gg)); V=int(torch.randint(2,10,(1,),generator=gg))
    loi=torch.randint(0,4,(B,V),generator=gg).float()
    loc=torch.randn(B,V,generator=gg)
    for lo in [loi,loc]:
        for p in [0.0,0.1,0.25,0.5,0.75,0.9,1.0]:
            cmp('top_p_filter',O.top_p_filter(lo,p),R.top_p_filter(lo,p))
            cmp('min_p_filter',O.min_p_filter(lo,p),R.min_p_filter(lo,p))
            cmp('epsilon_filter',O.epsilon_filter(lo,p),R.epsilon_filter(lo,p))
        for T in [1.0,2.0]:
          for k in range(1,V+1):
            for p in [0.3,0.5,0.75,1.0]:
                cmp('logits_warper',O.logits_warper(lo,T,k,p),R.logits_warper(lo,T,k,p),True)
                cmp('sample_greedy_warped',O.sample_greedy_warped(lo,T,k,p),R.sample_greedy_warped(lo,T,k,p),True)
print('tie+warper:', fails if fails else 'ALL PASS')
"

output

tie+warper: {'epsilon_filter': 133}

[stdout]
tie+warper: {'epsilon_filter': 133}

Find epsilon tie mismatch

shell

python3 -c "
import torch, reforacle as R, objectives as O
gg=torch.Generator().manual_seed(7000)
for _ in range(200):
    B=int(torch.randint(1,4,(1,),generator=gg)); V=int(torch.randint(2,8,(1,),generator=gg))
    lo=torch.randint(0,4,(B,V),generator=gg).float()
    for p in [0.1,0.25,0.5,0.9]:
        r=R.epsilon_filter(lo,p);m=O.epsilon_filter(lo,p)
        if not torch.equal(r,m):
            print('p',p,'lo',lo.tolist()); print(' ref',r.tolist()); print(' mine',m.tolist()); raise SystemExit
"

output

p 0.25 lo [[3.0, 1.0, 1.0, 3.0, 3.0, 1.0, 3.0], [3.0, 2.0, 3.0, 3.0, 1.0, 3.0, 0.0]]
 ref [[-inf, -inf, -inf, -inf, -inf, -inf, 3.0], [-inf, -inf, -inf, -inf, -inf, 3.0, -inf]]
 mine [[3.0, -inf, -inf, 3.0, 3.0, -inf, 3.0], [3.0, -inf, 3.0, 3.0, -inf, 3.0, -inf]]

[stdout]
p 0.25 lo [[3.0, 1.0, 1.0, 3.0, 3.0, 1.0, 3.0], [3.0, 2.0, 3.0, 3.0, 1.0, 3.0, 0.0]]
 ref [[-inf, -inf, -inf, -inf, -inf, -inf, 3.0], [-inf, -inf, -inf, -inf, -inf, 3.0, -inf]]
 mine [[3.0, -inf, -inf, 3.0, 3.0, -inf, 3.0], [3.0, -inf, 3.0, 3.0, -inf, 3.0, -inf]]

/app/objectives.py

− removed

    probs = logits.softmax(dim=-1)
    top = probs.max(dim=-1, keepdim=True).values
    remove = (probs < epsilon) & (probs < top)
    return logits.masked_fill(remove, NEG_INF)

+ added

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

Final comprehensive sweep with ties and edges

shell

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

for seed in range(60):
    gg=torch.Generator().manual_seed(seed)
    B=int(torch.randint(1,5,(1,),generator=gg)); V=int(torch.randint(2,13,(1,),generator=gg)); S=int(torch.randint(2,9,(1,),generator=gg))
    loc=torch.randn(B,V,generator=gg)
    loi=torch.randint(0,4,(B,V),generator=gg).float()
    for lo in [loc,loi]:
        cmp('greedy_token',O.greedy_token(lo),R.greedy_token(lo))
        for T in [0.0,0.5,1.0,2.0]: cmp('temperature_scale',O.temperature_scale(lo,T),R.temperature_scale(lo,T))
        for k in range(1,V+2):
            cmp('top_k_filter',O.top_k_filter(lo,k),R.top_k_filter(lo,k))
            cmp('topk_indices',O.topk_indices(lo,min(k,V)),R.topk_indices(lo,min(k,V)))
        for p in [0.0,0.15,0.4,0.7,0.9,0.99,1.0]:
            cmp('top_p_filter',O.top_p_filter(lo,p),R.top_p_filter(lo,p))
            cmp('min_p_filter',O.min_p_filter(lo,p),R.min_p_filter(lo,p))
            cmp('epsilon_filter',O.epsilon_filter(lo,p),R.epsilon_filter(lo,p))
        cmp('renormalize_logits',O.renormalize_logits(lo),R.renormalize_logits(lo))
        cmp('entropy_from_logits',O.entropy_from_logits(lo),R.entropy_from_logits(lo),True)
        mk=(torch.randint(0,2,(B,V),generator=gg)).float()
        cmp('masked_argmax',O.masked_argmax(lo,mk),R.masked_argmax(lo,mk))
        for T in [0.7,1.0,2.0]:
          for k in range(1,V+1):
            for p in [0.3,0.7,1.0]:
                cmp('logits_warper',O.logits_warper(lo,T,k,p),R.logits_warper(lo,T,k,p),True)
                cmp('sample_greedy_warped',O.sample_greedy_warped(lo,T,k,p),R.sample_greedy_warped(lo,T,k,p),True)
    ids=torch.randint(0,V,(B,S),generator=gg)
    idp=ids.masked_fill(torch.rand(B,S,generator=gg)<0.5,0)
    for pen in [0.5,0.7,1.0,1.3,1.8,2.5]: cmp('repetition_penalty',O.repetition_penalty(loc,ids,pen),R.repetition_penalty(loc,ids,pen))
    for pen in [0.0,0.5,1.5,3.0]:
        cmp('frequency_penalty',O.frequency_penalty(loc,ids,pen),R.frequency_penalty(loc,ids,pen))
        cmp('presence_penalty',O.presence_penalty(loc,ids,pen),R.presence_penalty(loc,ids,pen))
    cmp('suppress_tokens',O.suppress_tokens(loc,[0,min(1,V-1)]),R.suppress_tokens(loc,[0,min(1,V-1)]))
    for cl in [0,S,S+2]: cmp('min_length_filter',O.min_length_filter(loc,cl,S,min(2,V-1)),R.min_length_filter(loc,cl,S,min(2,V-1)))
    cmp('attention_mask',O.attention_mask(idp,0),R.attention_mask(idp,0))
    cmp('position_ids',O.position_ids(idp,0),R.position_ids(idp,0))
    cmp('causal_mask',O.causal_mask(S),R.causal_mask(S))
    cmp('shift_labels',O.shift_labels(ids,-100),R.shift_labels(ids,-100))
    pl=torch.randint(1,S+1,(B,),generator=gg)
    cmp('completion_loss_labels',O.completion_loss_labels(ids,pl,-100),R.completion_loss_labels(ids,pl,-100))
    cmp('prepare_labels',O.prepare_labels(ids,pl,-100),R.prepare_labels(ids,pl,-100))
    cmp('mask_after_eos',O.mask_after_eos(idp,3),R.mask_after_eos(idp,3))
    cmp('sequence_lengths',O.sequence_lengths(idp,0),R.sequence_lengths(idp,0))
    am=O.attention_mask(idp,0)
    cmp('last_token_indices',O.last_token_indices(am),R.last_token_indices(am))
    lo3=torch.randn(B,S,V,generator=gg)
    cmp('gather_last_logits',O.gather_last_logits(lo3,am),R.gather_last_logits(lo3,am))
    cmp('rollout_last_logits',O.rollout_last_logits(lo3,idp,0),R.rollout_last_logits(lo3,idp,0))
    for ml in [1,S,S+3]: cmp('truncate_left',O.truncate_left(ids,ml),R.truncate_left(ids,ml))
    for mult in [1,2,3,4,8]: cmp('pad_to_multiple',O.pad_to_multiple(ids,0,mult),R.pad_to_multiple(ids,0,mult))
    for w in [1,S,S+4]: cmp('right_pad_to',O.right_pad_to(ids,0,w),R.right_pad_to(ids,0,w))
    act=torch.randint(0,V,(B,S),generator=gg)
    cmp('action_logprobs',O.action_logprobs(lo3,act),R.action_logprobs(lo3,act),True)
    rew=torch.randn(B,S,generator=gg); rmask=(torch.rand(B,S,generator=gg)>0.4).float()
    cmp('reward_to_go',O.reward_to_go(rew,rmask),R.reward_to_go(rew,rmask),True)
    for gm in [0.5,0.9,1.0]: cmp('discount_cumsum',O.discount_cumsum(rew,gm),R.discount_cumsum(rew,gm),True)
    sl=torch.randint(1,4,(B,),generator=gg); tot=int(sl.sum())+int(torch.randint(0,3,(1,),generator=gg))
    cmp('build_position_ids_packed',O.build_position_ids_packed(sl,tot),R.build_position_ids_packed(sl,tot))
    stop=torch.tensor([2,3,min(4,V-1)])
    cmp('first_stop_index',O.first_stop_index(ids,stop),R.first_stop_index(ids,stop))
    for ml in [0,2,S,S+2]: cmp('generation_stop_mask',O.generation_stop_mask(idp,3,ml),R.generation_stop_mask(idp,3,ml))
print('FINAL:', fails if fails else 'ALL 37 PASS (60 seeds, ties+edges)')
"

output

FINAL: ALL 37 PASS (60 seeds, ties+edges)

[stdout]
FINAL: ALL 37 PASS (60 seeds, ties+edges)

Robustness: large vocab, nonzero pad, custom ignore_index

shell

python3 -c "
import torch, reforacle as R, objectives as O
fails={}
def cmp(name,a,b,comp=False):
    atol,rtol=(1e-5,1e-4) if comp else (1e-6,1e-5)
    ta=torch.as_tensor(a);tb=torch.as_tensor(b)
    if ta.shape!=tb.shape or ta.dtype!=tb.dtype or not torch.allclose(ta.float(),tb.float(),atol=atol,rtol=rtol,equal_nan=True):
        fails[name]=fails.get(name,0)+1;print('FAIL',name)
gg=torch.Generator().manual_seed(555)
# large realistic
lo=torch.randn(4,32000,generator=gg)
cmp('greedy_big',O.greedy_token(lo),R.greedy_token(lo))
cmp('topp_big',O.top_p_filter(lo,0.9),R.top_p_filter(lo,0.9))
cmp('topk_big',O.top_k_filter(lo,50),R.top_k_filter(lo,50))
cmp('entropy_big',O.entropy_from_logits(lo),R.entropy_from_logits(lo),True)
# nonzero pad, custom ignore
ids=torch.randint(0,50,(3,10),generator=gg)
idp=ids.masked_fill(torch.rand(3,10,generator=gg)<0.4, 7)
cmp('attn_pad7',O.attention_mask(idp,7),R.attention_mask(idp,7))
cmp('posid_pad7',O.position_ids(idp,7),R.position_ids(idp,7))
cmp('seqlen_pad7',O.sequence_lengths(idp,7),R.sequence_lengths(idp,7))
cmp('rollout_pad7',O.rollout_last_logits(torch.randn(3,10,50,generator=gg),idp,7),R.rollout_last_logits(torch.randn(3,10,50,generator=gg),idp,7)) if False else None
lo3=torch.randn(3,10,50,generator=gg)
cmp('rollout_pad7',O.rollout_last_logits(lo3,idp,7),R.rollout_last_logits(lo3,idp,7))
cmp('maskeos5',O.mask_after_eos(idp,5),R.mask_after_eos(idp,5))
cmp('genstop_pad7',O.generation_stop_mask(idp,5,3),R.generation_stop_mask(idp,5,3))
for ii in [-100,-1,0,999]:
    cmp('shift_ii',O.shift_labels(ids,ii),R.shift_labels(ids,ii))
    pl=torch.randint(1,11,(3,),generator=gg)
    cmp('prep_ii',O.prepare_labels(ids,pl,ii),R.prepare_labels(ids,pl,ii))
    cmp('comp_ii',O.completion_loss_labels(ids,pl,ii),R.completion_loss_labels(ids,pl,ii))
cmp('padmult_pad7',O.pad_to_multiple(ids,7,16),R.pad_to_multiple(ids,7,16))
cmp('rightpad_pad7',O.right_pad_to(ids,7,20),R.right_pad_to(ids,7,20))
# min_length various eos
cmp('minlen',O.min_length_filter(torch.randn(2,50,generator=gg),3,5,49),R.min_length_filter(torch.randn(2,50,generator=gg),3,5,49)) if False else None
print('EXTRA:', fails if fails else 'ALL EXTRA CHECKS PASS')
"

output

EXTRA: ALL EXTRA CHECKS PASS

[stdout]
EXTRA: ALL EXTRA CHECKS PASS

Final standalone + completeness verification

shell

cd /tmp && cp /app/objectives.py /tmp/final.py && python3 -c "
import builtins
_o=builtins.__import__
builtins.__import__=lambda n,*a,**k:(_ for _ in ()).throw(ImportError('blocked')) if n=='reforacle' else _o(n,*a,**k)
import final as O, torch
# confirm all 37 names implemented, none raise NotImplementedError
import inspect
names=[n for n,f in inspect.getmembers(O,inspect.isfunction) if not n.startswith('_')]
print('functions defined:', len(names))
src=open('/tmp/final.py').read()
print('NotImplementedError remaining:', src.count('NotImplementedError'))
print('imports reforacle:', 'reforacle' in src)
print('standalone import OK (reforacle blocked)')
"
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(9)
lo=torch.randn(2,50,generator=g)
print('minlen match:', torch.equal(O.min_length_filter(lo,3,5,49),R.min_length_filter(lo,3,5,49)))
print('suppress match:', torch.equal(O.suppress_tokens(lo,[3,7,20]),R.suppress_tokens(lo,[3,7,20])))
"

output

Exit code 1
functions defined: 37
NotImplementedError remaining: 0
imports reforacle: False
standalone import OK (reforacle blocked)
Traceback (most recent call last):
  File "<string>", line 2, in <module>
ModuleNotFoundError: No module named 'reforacle'

[error] tool reported failure

Confirm trivial funcs and check remaining budget

shell

python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(9)
lo=torch.randn(2,50,generator=g)
print('minlen match:', torch.equal(O.min_length_filter(lo,3,5,49),R.min_length_filter(lo,3,5,49)))
print('suppress match:', torch.equal(O.suppress_tokens(lo,[3,7,20]),R.suppress_tokens(lo,[3,7,20])))
"; bash /app/timer.sh

output

minlen match: True
suppress match: True
Remaining time (hours:minutes):
1:44

[stdout]
minlen match: True
suppress match: True
Remaining time (hours:minutes):
1:44

All 37 functions are implemented and verified. The implementation is complete. ## Summary I reverse-engineered the sealed `reforacle` module by probing it, then reimplemented all 37 functions in `/app/objectives.py` as a standalone `torch`-only module (verified it imports and runs with `reforacle` blocked, no `NotImplementedError` remaining). **Validation:** all 37 functions match the reference exactly across 60+ random seeds with varied shapes (batch sizes, vocab up to 32000), plus dedicated tie-heavy, boundary, and edge-case inputs (all-pad rows, interior pads, nonzero `pad_id`, custom `ignore_index`). Key non-obvious conventions I recovered by probing (not assuming standard definitions): - **Argmax picks the *last* max, not the first** , affects `greedy_token`, `masked_argmax`, `sample_greedy_warped`, and `temperature_scale(T=0)`. Implemented via a `_last_argmax` helper. - **`topk_indices`** breaks ties by *ascending index* (stable descending sort), unlike `torch.topk`. - **`top_p_filter`** uses exclusive-cumsum `>= p` with a *stable* sort , the shift-based `>` variant disagrees on tied/exact-boundary probabilities. - **`epsilon_filter`** keeps exactly one token (the last-argmax) when all probs fall below epsilon, not all tied maxes. - **`repetition_penalty`** is *count-based* (applies the penalty once per occurrence), not presence-based. - **`reward_to_go`** re-masks after the reversed cumsum (`(rewards*mask).flip.cumsum.flip * mask`). - **`sequence_lengths`** counts non-pad tokens, while **`last_token_indices`/`rollout_last_logits`** find the *last* non-pad position (differ under interior pads); `gather_last_logits` clamps the all-pad `-1` index to `0`. - **`generation_stop_mask`**: valid length = `max(min_len, first_eos+1)`.

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: 37/37  (partial_score=1.0)
REWARD = 1   (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_d145e9ebc152479d. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.

Trial trial_d145e9ebc152479d · verifier authoritative; classifier explanatory.