SyncValsverifier → artifact → classifier → verdict
SyncVals · Trajectory

reverse-engineer-decoding

claude-code claude-opus-4-8 ✗ failed GOOD_FAILURE ↑ View task
Solved from the instruction alone, tests/ and solution/ were withheld from the agent's workspace and restored only for grading.
Reward = tests/test.sh exit code (0 → resolved); the classification below is post-hoc and cannot change it.
Classification , post-hoc; cannot change the reward
GOOD_FAILUREHonest miss, the agent ran correctly but couldn't solve it. Expected for a hard task; the task is sound.
SubtypeIncomplete Understanding
EvidenceTest output: "FUNCTIONS PASSED: 36/37 (partial_score=0.973)" with reward=0.0. Agent's trajectory shows methodical probing, discovery of all 37 functions' core behavior, and identification of tie-breaking as the root cause of remaining failures. Agent narrowed failures to 3 functions (top_p_filter, topk_indices, epsilon_filter) and verified these are tie-only edge cases (400 trials without ties passed: bad_stable=0, bad_p=0), but could not fully debug the exact convention on hidden grading inputs before task timeout.
Root causeAgent correctly reverse-engineered 36/37 functions via probing but failed to fully determine the exact tie-breaking/stable-sort convention for one remaining function on the hidden grading inputs. This is a measure-zero edge case (ties) that is difficult to discover without direct oracle feedback on the exact failing input.
RecommendationN/A - task is fine. The specification is complete (all 37 functions discoverable via probing) and grading rules are clear. Agent performance was excellent (97.3%) but the final 1-3 functions require exact tie-breaking knowledge that can only be reliably obtained with more probing budget or by seeing the grading failures directly. This is expected difficulty for a reverse-engineering task with 37 functions and subtle conventions.
Trajectory
Tool-by-tool agent trajectory
75 tool calls · 4 tool types · 77 steps
Reimplement the post-training decoding-and-collation library in `/app/objectives.py` so that every function reproduces the sealed reference module `reforacle` exactly. `reforacle` is importable and callable (`import reforacle`) but its source cannot be read. It is the only ground truth; there is no written specification of what any function computes or of the conventions it follows. Recover its behavior by calling it. `/app/objectives.py` contains thirty-seven function stubs , logit processors and sampling filters, sequence collation and masking utilities, and rollout helpers used to generate and batch data for RLHF and SFT. Keep every signature exactly as given and implement every function. The names indicate the standard quantity each is meant to produce, but the reference's exact definition is whatever `reforacle` returns , not what you assume. Determine each definition by probing. Grading compares your implementation to `reforacle` on held-out inputs you do not see, function by function, requiring exact numerical agreement (primitives within `atol=1e-6, rtol=1e-5`; composite functions within `atol=1e-5, rtol=1e-4`). Each function passes only if it agrees on every graded input, and the task passes only if all thirty-seven functions pass. Agreement on the inputs you happen to try does not imply agreement on the inputs you are graded on. Grading runs your `/app/objectives.py` in a clean sandbox that does not contain `reforacle`, so any submission that imports, wraps, calls, loads, or otherwise depends on the reference at run time fails outright. Your file must be a standalone implementation whose only dependency is `torch`. Do not read, decompile, disassemble, or otherwise extract the compiled reference; recover its behavior only by calling it and comparing. `python3 /app/probe.py` shows the calling pattern; `bash /app/timer.sh` reports remaining budget.

/app/objectives.py

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

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

Remaining time (hours:minutes):
1:59

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

/app/probe.py

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

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

Remaining time (hours:minutes):
1:59

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

Check remaining budget

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

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

Remaining time (hours:minutes):
1:59

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

Run the sample probe

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

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

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

Probe any reference function this way to discover its exact behavior,
then implement it in /app/objectives.py to match.
write · /app/p.py
from __future__ import annotations
import torch
import reforacle as R

g = torch.Generator().manual_seed(0)

def show(name, val):
    try:
        print(name, "->", val.tolist() if torch.is_tensor(val) else val, "| dtype", val.dtype if torch.is_tensor(val) else type(val))
    except Exception as e:
        print(name, "EXC", e)

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

# greedy_token
show("greedy_token", R.greedy_token(logits))
# temperature_scale
show("temp T=2", R.temperature_scale(logits, 2.0))
show("temp T=0", R.temperature_scale(logits, 0.0))
show("temp T=1", R.temperature_scale(logits, 1.0))
# top_k
show("top_k k=2", R.top_k_filter(logits, 2))
show("top_k k=1", R.top_k_filter(logits, 1))
show("top_k k=100", R.top_k_filter(logits, 100))
# top_p
probs = torch.softmax(logits, dim=-1)
print("probs", probs.tolist())
show("top_p 0.5", R.top_p_filter(logits, 0.5))
show("top_p 0.9", R.top_p_filter(logits, 0.9))
show("top_p 1.0", R.top_p_filter(logits, 1.0))
# min_p
show("min_p 0.1", R.min_p_filter(logits, 0.1))
show("min_p 0.5", R.min_p_filter(logits, 0.5))

Probe first group

shell
python3 p.py
output
logits [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
greedy_token -> [0, 1] | dtype torch.int64
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]] | dtype torch.float32
temp T=0 -> [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]] | dtype torch.float32
temp T=1 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
top_k k=2 -> [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]] | dtype torch.float32
top_k k=1 -> [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]] | dtype torch.float32
top_k k=100 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
probs [[0.5926466584205627, 0.0946488156914711, 0.014365261420607567, 0.22408682107925415, 0.0429089330136776, 0.031343501061201096], [0.22293631732463837, 0.3443180024623871, 0.07255042344331741, 0.09950380027294159, 0.08201515674591064, 0.1786762923002243]]
top_p 0.5 -> [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]] | dtype torch.float32
top_p 0.9 -> [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
top_p 1.0 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
min_p 0.1 -> [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
min_p 0.5 -> [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]] | dtype torch.float32

[stdout]
logits [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
greedy_token -> [0, 1] | dtype torch.int64
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]] | dtype torch.float32
temp T=0 -> [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]] | dtype torch.float32
temp T=1 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
top_k k=2 -> [[1.5409960746765137, -inf, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]] | dtype torch.float32
top_k k=1 -> [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [-inf, 0.8380263447761536, -inf, -inf, -inf, -inf]] | dtype torch.float32
top_k k=100 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
probs [[0.5926466584205627, 0.0946488156914711, 0.014365261420607567, 0.22408682107925415, 0.0429089330136776, 0.031343501061201096], [0.22293631732463837, 0.3443180024623871, 0.07255042344331741, 0.09950380027294159, 0.08201515674591064, 0.1786762923002243]]
top_p 0.5 -> [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, -inf]] | dtype torch.float32
top_p 0.9 -> [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
top_p 1.0 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
min_p 0.1 -> [[1.5409960746765137, -0.293428897857666, -inf, 0.5684312582015991, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | dtype torch.float32
min_p 0.5 -> [[1.5409960746765137, -inf, -inf, -inf, -inf, -inf], [0.40334683656692505, 0.8380263447761536, -inf, -inf, -inf, 0.18203648924827576]] | dtype torch.float32
write · /app/p.py
from __future__ import annotations
import torch
import reforacle as R

def show(name, val):
    try:
        print(name, "->", val.tolist() if torch.is_tensor(val) else val, "|", val.dtype if torch.is_tensor(val) else type(val))
    except Exception as e:
        print(name, "EXC", repr(e))

g = torch.Generator().manual_seed(0)
logits = torch.randn(2, 6, generator=g)
input_ids = torch.tensor([[0, 1, 1, 3, 0], [2, 2, 5, 5, 5]])
print("logits", logits.tolist())

show("rep_pen 1.5", R.repetition_penalty(logits, input_ids, 1.5))
show("freq_pen 0.5", R.frequency_penalty(logits, input_ids, 0.5))
show("pres_pen 0.5", R.presence_penalty(logits, input_ids, 0.5))
show("suppress [1,3]", R.suppress_tokens(logits, [1, 3]))
show("suppress tensor", R.suppress_tokens(logits, torch.tensor([0, 5])))
show("min_length cur=2,min=5,eos=0", R.min_length_filter(logits, 2, 5, 0))
show("min_length cur=5,min=5,eos=0", R.min_length_filter(logits, 5, 5, 0))
show("min_length cur=6,min=5,eos=0", R.min_length_filter(logits, 6, 5, 0))
show("renorm", R.renormalize_logits(logits))
print("logsoftmax", torch.log_softmax(logits, -1).tolist())

Probe penalties group

shell
python3 p.py
output
logits [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]]
rep_pen 1.5 -> [[0.684887170791626, -0.6602150201797485, -2.1787893772125244, 0.3789541721343994, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -1.6183295249938965, -0.40334352850914, -0.5966353416442871, 0.053936734795570374]] | torch.float32
freq_pen 0.5 -> [[0.5409960746765137, -1.293428897857666, -2.1787893772125244, 0.06843125820159912, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -1.7192575931549072, -0.40334352850914, -0.5966353416442871, -1.3179634809494019]] | torch.float32
pres_pen 0.5 -> [[1.0409960746765137, -0.793428897857666, -2.1787893772125244, 0.06843125820159912, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -1.2192575931549072, -0.40334352850914, -0.5966353416442871, -0.31796351075172424]] | torch.float32
suppress [1,3] -> [[1.5409960746765137, -inf, -2.1787893772125244, -inf, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, -inf, -0.7192575931549072, -inf, -0.5966353416442871, 0.18203648924827576]] | torch.float32
suppress tensor -> [[-inf, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -inf], [-inf, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, -inf]] | torch.float32
min_length cur=2,min=5,eos=0 -> [[-inf, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [-inf, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | torch.float32
min_length cur=5,min=5,eos=0 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | torch.float32
min_length cur=6,min=5,eos=0 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | torch.float32
renorm -> [[-0.5231569409370422, -2.357581853866577, -4.2429423332214355, -1.4957218170166016, -3.1486752033233643, -3.4627482891082764], [-1.5008691549301147, -1.0661896467208862, -2.623473644256592, -2.3075594902038574, -2.5008511543273926, -1.7221795320510864]] | torch.float32
logsoftmax [[-0.5231569409370422, -2.357581853866577, -4.2429423332214355, -1.4957218170166016, -3.1486752033233643, -3.4627482891082764], [-1.5008691549301147, -1.0661896467208862, -2.623473644256592, -2.3075594902038574, -2.5008511543273926, -1.7221795320510864]]

[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]]
rep_pen 1.5 -> [[0.684887170791626, -0.6602150201797485, -2.1787893772125244, 0.3789541721343994, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -1.6183295249938965, -0.40334352850914, -0.5966353416442871, 0.053936734795570374]] | torch.float32
freq_pen 0.5 -> [[0.5409960746765137, -1.293428897857666, -2.1787893772125244, 0.06843125820159912, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -1.7192575931549072, -0.40334352850914, -0.5966353416442871, -1.3179634809494019]] | torch.float32
pres_pen 0.5 -> [[1.0409960746765137, -0.793428897857666, -2.1787893772125244, 0.06843125820159912, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -1.2192575931549072, -0.40334352850914, -0.5966353416442871, -0.31796351075172424]] | torch.float32
suppress [1,3] -> [[1.5409960746765137, -inf, -2.1787893772125244, -inf, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, -inf, -0.7192575931549072, -inf, -0.5966353416442871, 0.18203648924827576]] | torch.float32
suppress tensor -> [[-inf, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -inf], [-inf, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, -inf]] | torch.float32
min_length cur=2,min=5,eos=0 -> [[-inf, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [-inf, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | torch.float32
min_length cur=5,min=5,eos=0 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | torch.float32
min_length cur=6,min=5,eos=0 -> [[1.5409960746765137, -0.293428897857666, -2.1787893772125244, 0.5684312582015991, -1.0845223665237427, -1.3985954523086548], [0.40334683656692505, 0.8380263447761536, -0.7192575931549072, -0.40334352850914, -0.5966353416442871, 0.18203648924827576]] | torch.float32
renorm -> [[-0.5231569409370422, -2.357581853866577, -4.2429423332214355, -1.4957218170166016, -3.1486752033233643, -3.4627482891082764], [-1.5008691549301147, -1.0661896467208862, -2.623473644256592, -2.3075594902038574, -2.5008511543273926, -1.7221795320510864]] | torch.float32
logsoftmax [[-0.5231569409370422, -2.357581853866577, -4.2429423332214355, -1.4957218170166016, -3.1486752033233643, -3.4627482891082764], [-1.5008691549301147, -1.0661896467208862, -2.623473644256592, -2.3075594902038574, -2.5008511543273926, -1.7221795320510864]]
write · /app/p.py
from __future__ import annotations
import torch
import reforacle as R

def show(name, val):
    try:
        print(name, "->", val.tolist() if torch.is_tensor(val) else val, "|", val.dtype if torch.is_tensor(val) else type(val))
    except Exception as e:
        print(name, "EXC", repr(e))

# pad_id = 0 assumption? probe
ids = torch.tensor([[5, 6, 7, 0, 0], [8, 9, 0, 0, 0], [1, 2, 3, 4, 5]])
show("attention_mask pad0", R.attention_mask(ids, 0))
show("position_ids pad0", R.position_ids(ids, 0))
# left padded?
ids2 = torch.tensor([[0, 0, 5, 6, 7], [0, 8, 9, 1, 2]])
show("attention_mask left pad0", R.attention_mask(ids2, 0))
show("position_ids left pad0", R.position_ids(ids2, 0))
show("causal_mask 4", R.causal_mask(4))
show("shift_labels", R.shift_labels(ids, -100))
seq = torch.tensor([[10, 11, 12, 13, 14], [20, 21, 22, 23, 24]])
show("completion_loss_labels pl=[2,3]", R.completion_loss_labels(seq, torch.tensor([2, 3]), -100))
eids = torch.tensor([[5, 6, 2, 7, 8], [1, 2, 3, 2, 4]])
show("mask_after_eos eos2", R.mask_after_eos(eids, 2))
show("sequence_lengths pad0", R.sequence_lengths(ids, 0))
show("sequence_lengths left pad0", R.sequence_lengths(ids2, 0))
am = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 0, 0, 0]])
show("last_token_indices", R.last_token_indices(am))
amL = torch.tensor([[0, 0, 1, 1, 1], [0, 1, 1, 1, 1]])
show("last_token_indices left", R.last_token_indices(amL))
lg = torch.randn(2, 5, 4, torch.Generator().manual_seed(1) and None) if False else torch.randn(2, 5, 4)

Probe collation group

shell
python3 p.py
output
attention_mask pad0 -> [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]] | torch.float32
position_ids pad0 -> [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1], [0, 1, 2, 3, 4]] | torch.int64
attention_mask left pad0 -> [[0.0, 0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0, 1.0]] | torch.float32
position_ids left pad0 -> [[0, 0, 0, 1, 2], [0, 0, 1, 2, 3]] | torch.int64
causal_mask 4 -> [[0.0, -inf, -inf, -inf], [0.0, 0.0, -inf, -inf], [0.0, 0.0, 0.0, -inf], [0.0, 0.0, 0.0, 0.0]] | torch.float32
shift_labels -> [[6, 7, 0, 0, -100], [9, 0, 0, 0, -100], [2, 3, 4, 5, -100]] | torch.int64
completion_loss_labels pl=[2,3] -> [[-100, -100, 12, 13, 14], [-100, -100, -100, 23, 24]] | torch.int64
mask_after_eos eos2 -> [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] | torch.float32
sequence_lengths pad0 -> [3, 2, 5] | torch.int64
sequence_lengths left pad0 -> [3, 4] | torch.int64
last_token_indices -> [2, 1] | torch.int64
last_token_indices left -> [4, 4] | torch.int64

[stdout]
attention_mask pad0 -> [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0], [1.0, 1.0, 1.0, 1.0, 1.0]] | torch.float32
position_ids pad0 -> [[0, 1, 2, 2, 2], [0, 1, 1, 1, 1], [0, 1, 2, 3, 4]] | torch.int64
attention_mask left pad0 -> [[0.0, 0.0, 1.0, 1.0, 1.0], [0.0, 1.0, 1.0, 1.0, 1.0]] | torch.float32
position_ids left pad0 -> [[0, 0, 0, 1, 2], [0, 0, 1, 2, 3]] | torch.int64
causal_mask 4 -> [[0.0, -inf, -inf, -inf], [0.0, 0.0, -inf, -inf], [0.0, 0.0, 0.0, -inf], [0.0, 0.0, 0.0, 0.0]] | torch.float32
shift_labels -> [[6, 7, 0, 0, -100], [9, 0, 0, 0, -100], [2, 3, 4, 5, -100]] | torch.int64
completion_loss_labels pl=[2,3] -> [[-100, -100, 12, 13, 14], [-100, -100, -100, 23, 24]] | torch.int64
mask_after_eos eos2 -> [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] | torch.float32
sequence_lengths pad0 -> [3, 2, 5] | torch.int64
sequence_lengths left pad0 -> [3, 4] | torch.int64
last_token_indices -> [2, 1] | torch.int64
last_token_indices left -> [4, 4] | torch.int64
write · /app/p.py
from __future__ import annotations
import torch
import reforacle as R

def show(name, val):
    try:
        print(name, "->", val.tolist() if torch.is_tensor(val) else val, "|", val.dtype if torch.is_tensor(val) else type(val), (val.shape if torch.is_tensor(val) else ""))
    except Exception as e:
        print(name, "EXC", repr(e))

g = torch.Generator().manual_seed(3)
logits = torch.randn(2, 5, 4, generator=g)
am = torch.tensor([[1, 1, 1, 0, 0], [1, 1, 0, 0, 0]])
show("gather_last_logits", R.gather_last_logits(logits, am))
print("expected r0 idx2", logits[0,2].tolist(), "r1 idx1", logits[1,1].tolist())
amL = torch.tensor([[0, 0, 1, 1, 1], [0, 1, 1, 1, 1]])
show("gather_last left", R.gather_last_logits(logits, amL))

ids = torch.tensor([[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]])
show("truncate_left max3", R.truncate_left(ids, 3))
show("truncate_left max10", R.truncate_left(ids, 10))
ids2 = torch.tensor([[1, 2, 3], [4, 5, 6]])
show("pad_to_multiple m4 pad0", R.pad_to_multiple(ids2, 0, 4))
show("pad_to_multiple m2 pad0", R.pad_to_multiple(ids2, 0, 2))
show("pad_to_multiple m3 pad0", R.pad_to_multiple(ids2, 0, 3))

# action_logprobs
lg = torch.randn(2, 3, 5, generator=g)
actions = torch.tensor([[1, 4, 2], [0, 3, 1]])
show("action_logprobs", R.action_logprobs(lg, actions))
print("expected", torch.log_softmax(lg,-1).gather(-1, actions.unsqueeze(-1)).squeeze(-1).tolist())
show("entropy_from_logits", R.entropy_from_logits(lg))
p = torch.softmax(lg, -1)
print("expected entropy", (-(p*torch.log_softmax(lg,-1)).sum(-1)).tolist())

Probe gather/truncate/pad/rl

shell
python3 p.py
output
gather_last_logits -> [[-0.05168094485998154, -0.30597570538520813, 0.2485019415616989, -0.2225608378648758], [-1.1963194608688354, 1.0279964208602905, 0.07189108431339264, -0.18447615206241608]] | torch.float32 torch.Size([2, 4])
expected r0 idx2 [-0.05168094485998154, -0.30597570538520813, 0.2485019415616989, -0.2225608378648758] r1 idx1 [-1.1963194608688354, 1.0279964208602905, 0.07189108431339264, -0.18447615206241608]
gather_last left -> [[0.2615012228488922, 0.9311244487762451, -0.5144620537757874, -1.6516656875610352], [0.7242393493652344, -0.349177747964859, -0.5289880037307739, -0.6894744634628296]] | torch.float32 torch.Size([2, 4])
truncate_left max3 -> [[4, 5, 6], [10, 11, 12]] | torch.int64 torch.Size([2, 3])
truncate_left max10 -> [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]] | torch.int64 torch.Size([2, 6])
pad_to_multiple m4 pad0 -> [[1, 2, 3, 0], [4, 5, 6, 0]] | torch.int64 torch.Size([2, 4])
pad_to_multiple m2 pad0 -> [[1, 2, 3, 0], [4, 5, 6, 0]] | torch.int64 torch.Size([2, 4])
pad_to_multiple m3 pad0 -> [[1, 2, 3], [4, 5, 6]] | torch.int64 torch.Size([2, 3])
action_logprobs -> [[-0.5500109195709229, -0.9929203391075134, -3.190166711807251], [-1.7501140832901, -1.707658052444458, -1.421061396598816]] | torch.float32 torch.Size([2, 3])
expected [[-0.5500109195709229, -0.9929203391075134, -3.190166711807251], [-1.7501140832901, -1.707658052444458, -1.421061396598816]]
entropy_from_logits -> [[1.1941139698028564, 1.4659098386764526, 1.2124929428100586], [1.5657809972763062, 1.4768693447113037, 1.0967161655426025]] | torch.float32 torch.Size([2, 3])
expected entropy [[1.1941139698028564, 1.4659098386764526, 1.2124929428100586], [1.5657809972763062, 1.4768692255020142, 1.096716046333313]]

[stdout]
gather_last_logits -> [[-0.05168094485998154, -0.30597570538520813, 0.2485019415616989, -0.2225608378648758], [-1.1963194608688354, 1.0279964208602905, 0.07189108431339264, -0.18447615206241608]] | torch.float32 torch.Size([2, 4])
expected r0 idx2 [-0.05168094485998154, -0.30597570538520813, 0.2485019415616989, -0.2225608378648758] r1 idx1 [-1.1963194608688354, 1.0279964208602905, 0.07189108431339264, -0.18447615206241608]
gather_last left -> [[0.2615012228488922, 0.9311244487762451, -0.5144620537757874, -1.6516656875610352], [0.7242393493652344, -0.349177747964859, -0.5289880037307739, -0.6894744634628296]] | torch.float32 torch.Size([2, 4])
truncate_left max3 -> [[4, 5, 6], [10, 11, 12]] | torch.int64 torch.Size([2, 3])
truncate_left max10 -> [[1, 2, 3, 4, 5, 6], [7, 8, 9, 10, 11, 12]] | torch.int64 torch.Size([2, 6])
pad_to_multiple m4 pad0 -> [[1, 2, 3, 0], [4, 5, 6, 0]] | torch.int64 torch.Size([2, 4])
pad_to_multiple m2 pad0 -> [[1, 2, 3, 0], [4, 5, 6, 0]] | torch.int64 torch.Size([2, 4])
pad_to_multiple m3 pad0 -> [[1, 2, 3], [4, 5, 6]] | torch.int64 torch.Size([2, 3])
action_logprobs -> [[-0.5500109195709229, -0.9929203391075134, -3.190166711807251], [-1.7501140832901, -1.707658052444458, -1.421061396598816]] | torch.float32 torch.Size([2, 3])
expected [[-0.5500109195709229, -0.9929203391075134, -3.190166711807251], [-1.7501140832901, -1.707658052444458, -1.421061396598816]]
entropy_from_logits -> [[1.1941139698028564, 1.4659098386764526, 1.2124929428100586], [1.5657809972763062, 1.4768693447113037, 1.0967161655426025]] | torch.float32 torch.Size([2, 3])
expected entropy [[1.1941139698028564, 1.4659098386764526, 1.2124929428100586], [1.5657809972763062, 1.4768692255020142, 1.096716046333313]]
write · /app/p.py
from __future__ import annotations
import torch
import reforacle as R

def show(name, val):
    try:
        print(name, "->", val.tolist() if torch.is_tensor(val) else val, "|", val.dtype if torch.is_tensor(val) else type(val), (list(val.shape) if torch.is_tensor(val) else ""))
    except Exception as e:
        print(name, "EXC", repr(e))

rewards = torch.tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]])
mask = torch.tensor([[1.0, 1.0, 1.0, 0.0], [1.0, 1.0, 0.0, 0.0]])
show("reward_to_go", R.reward_to_go(rewards, mask))
show("reward_to_go nomask", R.reward_to_go(rewards, torch.ones_like(mask)))

x = torch.tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]])
show("discount_cumsum g0.9", R.discount_cumsum(x, 0.9))
x1 = torch.tensor([1.0, 2.0, 3.0])
show("discount_cumsum 1d g0.5", R.discount_cumsum(x1, 0.5))

show("build_pos_packed [3,2] tot6", R.build_position_ids_packed(torch.tensor([3, 2]), 6))
show("build_pos_packed [2,3,1] tot6", R.build_position_ids_packed(torch.tensor([2, 3, 1]), 6))
show("build_pos_packed [2,2] tot5", R.build_position_ids_packed(torch.tensor([2, 2]), 5))

g = torch.Generator().manual_seed(5)
logits = torch.randn(2, 6, generator=g)
print("logits", logits.tolist())
show("logits_warper T2,k3,p0.9", R.logits_warper(logits, 2.0, 3, 0.9))
show("sample_greedy_warped T2,k3,p0.9", R.sample_greedy_warped(logits, 2.0, 3, 0.9))

Probe rl functions

shell
python3 p.py
output
reward_to_go -> [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]] | torch.float32 [2, 4]
reward_to_go nomask -> [[10.0, 9.0, 7.0, 4.0], [26.0, 21.0, 15.0, 8.0]] | torch.float32 [2, 4]
discount_cumsum g0.9 -> [[8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0], [21.90199851989746, 18.779998779296875, 14.199999809265137, 8.0]] | torch.float32 [2, 4]
discount_cumsum 1d g0.5 -> [2.75, 3.5, 3.0] | torch.float32 [3]
build_pos_packed [3,2] tot6 -> [0, 1, 2, 0, 1, 0] | torch.int64 [6]
build_pos_packed [2,3,1] tot6 -> [0, 1, 0, 1, 2, 0] | torch.int64 [6]
build_pos_packed [2,2] tot5 -> [0, 1, 0, 1, 0] | torch.int64 [5]
logits [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
logits_warper T2,k3,p0.9 -> [[-inf, -inf, -inf, 0.33376213908195496, -0.09870754927396774, 0.9713917374610901], [-inf, -inf, 0.3156067132949829, -inf, -0.2788967192173004, 0.3453593850135803]] | torch.float32 [2, 6]
sample_greedy_warped T2,k3,p0.9 -> [5, 5] | torch.int64 [2]

[stdout]
reward_to_go -> [[6.0, 5.0, 3.0, 0.0], [11.0, 6.0, 0.0, 0.0]] | torch.float32 [2, 4]
reward_to_go nomask -> [[10.0, 9.0, 7.0, 4.0], [26.0, 21.0, 15.0, 8.0]] | torch.float32 [2, 4]
discount_cumsum g0.9 -> [[8.145999908447266, 7.939999580383301, 6.599999904632568, 4.0], [21.90199851989746, 18.779998779296875, 14.199999809265137, 8.0]] | torch.float32 [2, 4]
discount_cumsum 1d g0.5 -> [2.75, 3.5, 3.0] | torch.float32 [3]
build_pos_packed [3,2] tot6 -> [0, 1, 2, 0, 1, 0] | torch.int64 [6]
build_pos_packed [2,3,1] tot6 -> [0, 1, 0, 1, 2, 0] | torch.int64 [6]
build_pos_packed [2,2] tot5 -> [0, 1, 0, 1, 0] | torch.int64 [5]
logits [[-0.4867807626724243, -0.6038215756416321, -0.5580962300300598, 0.6675242781639099, -0.19741509854793549, 1.9427834749221802], [-1.4017020463943481, -0.7625571489334106, 0.6312134265899658, -0.899135172367096, -0.5577934384346008, 0.6907187700271606]]
logits_warper T2,k3,p0.9 -> [[-inf, -inf, -inf, 0.33376213908195496, -0.09870754927396774, 0.9713917374610901], [-inf, -inf, 0.3156067132949829, -inf, -0.2788967192173004, 0.3453593850135803]] | torch.float32 [2, 6]
sample_greedy_warped T2,k3,p0.9 -> [5, 5] | torch.int64 [2]
write · /app/p.py
from __future__ import annotations
import torch
import reforacle as R

def show(name, val):
    try:
        print(name, "->", val.tolist() if torch.is_tensor(val) else val, "|", val.dtype if torch.is_tensor(val) else type(val), (list(val.shape) if torch.is_tensor(val) else ""))
    except Exception as e:
        print(name, "EXC", repr(e))

# reward_to_go middle mask
rewards = torch.tensor([[1.0, 2.0, 3.0, 4.0]])
mask = torch.tensor([[1.0, 0.0, 1.0, 1.0]])
show("rtg midmask", R.reward_to_go(rewards, mask))

# prepare_labels
seq = torch.tensor([[10, 11, 12, 13, 14], [20, 21, 22, 23, 24]])
show("prepare_labels pl=[2,3]", R.prepare_labels(seq, torch.tensor([2, 3]), -100))

# rollout_last_logits
g = torch.Generator().manual_seed(7)
logits = torch.randn(2, 5, 4, generator=g)
ids = torch.tensor([[1, 2, 3, 0, 0], [4, 5, 0, 0, 0]])
show("rollout_last_logits pad0", R.rollout_last_logits(logits, ids, 0))
print("exp r0 idx2", logits[0,2].tolist(), "r1 idx1", logits[1,1].tolist())

# epsilon_filter
lg = torch.randn(2, 6, generator=torch.Generator().manual_seed(9))
print("lg", lg.tolist())
print("probs", torch.softmax(lg,-1).tolist())
show("epsilon_filter 0.1", R.epsilon_filter(lg, 0.1))
show("epsilon_filter 0.15", R.epsilon_filter(lg, 0.15))

# first_stop_index
ids2 = torch.tensor([[1, 2, 3, 4, 5], [7, 8, 9, 2, 3], [5, 5, 5, 5, 5]])
show("first_stop_index [3,9]", R.first_stop_index(ids2, [3, 9]))
show("first_stop_index tensor[2]", R.first_stop_index(ids2, torch.tensor([2])))
show("first_stop_index [100]", R.first_stop_index(ids2, [100]))

# masked_argmax
lg2 = torch.randn(2, 5, generator=torch.Generator().manual_seed(11))
m = torch.tensor([[1, 1, 0, 1, 0], [0, 1, 1, 0, 1]])
print("lg2", lg2.tolist())
show("masked_argmax", R.masked_argmax(lg2, m))

# topk_indices
show("topk_indices k3", R.topk_indices(lg2, 3))

# right_pad_to
ids3 = torch.tensor([[1, 2, 3], [4, 5, 6]])
show("right_pad_to w5 pad0", R.right_pad_to(ids3, 0, 5))
show("right_pad_to w2 pad0", R.right_pad_to(ids3, 0, 2))

# generation_stop_mask
ids4 = torch.tensor([[5, 6, 2, 7, 8], [1, 2, 3, 4, 2]])
show("gen_stop_mask eos2 minlen0", R.generation_stop_mask(ids4, 2, 0))
show("gen_stop_mask eos2 minlen3", R.generation_stop_mask(ids4, 2, 3))

Probe remaining functions

shell
python3 p.py
output
Exit code 1
rtg midmask -> [[8.0, 0.0, 7.0, 4.0]] | torch.float32 [1, 4]
prepare_labels pl=[2,3] -> [[-100, 12, 13, 14, -100], [-100, -100, 23, 24, -100]] | torch.int64 [2, 5]
rollout_last_logits pad0 -> [[0.9265430569648743, -0.535512387752533, -1.1597206592559814, -0.46015679836273193], [-0.13543689250946045, 0.7471182346343994, 0.6118157505989075, 1.867783546447754]] | torch.float32 [2, 4]
exp r0 idx2 [0.9265430569648743, -0.535512387752533, -1.1597206592559814, -0.46015679836273193] r1 idx1 [-0.13543689250946045, 0.7471182346343994, 0.6118157505989075, 1.867783546447754]
lg [[0.04472753778100014, 1.9112390279769897, -0.2310308963060379, 0.3459281623363495, 1.318028211593628, 0.36963701248168945], [0.38411155343055725, 0.2970382571220398, 0.7472865581512451, 1.0873746871948242, -0.771757185459137, -0.2942127287387848]]
probs [[0.06881029903888702, 0.4449068307876587, 0.05222675949335098, 0.0929957702755928, 0.2458334118127823, 0.0952269434928894], [0.16136623919010162, 0.14790989458560944, 0.23202650249004364, 0.32601380348205566, 0.050795506685972214, 0.08188799023628235]]
epsilon_filter 0.1 -> [[-inf, 1.9112390279769897, -inf, -inf, 1.318028211593628, -inf], [0.38411155343055725, 0.2970382571220398, 0.7472865581512451, 1.0873746871948242, -inf, -inf]] | torch.float32 [2, 6]
epsilon_filter 0.15 -> [[-inf, 1.9112390279769897, -inf, -inf, 1.318028211593628, -inf], [0.38411155343055725, -inf, 0.7472865581512451, 1.0873746871948242, -inf, -inf]] | torch.float32 [2, 6]
Traceback (most recent call last):
  File "/app/p.py", line 36, in <module>
    show("first_stop_index [3,9]", R.first_stop_index(ids2, [3, 9]))
                                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "reforacle.py", line 243, in reforacle.first_stop_index
AttributeError: 'list' object has no attribute 'tolist'

[error] tool reported failure
write · /app/p.py
from __future__ import annotations
import torch
import reforacle as R

def show(name, val):
    try:
        print(name, "->", val.tolist() if torch.is_tensor(val) else val, "|", val.dtype if torch.is_tensor(val) else type(val), (list(val.shape) if torch.is_tensor(val) else ""))
    except Exception as e:
        print(name, "EXC", repr(e))

# first_stop_index with tensor
ids2 = torch.tensor([[1, 2, 3, 4, 5], [7, 8, 9, 2, 3], [5, 5, 5, 5, 5]])
show("first_stop_index [3,9]", R.first_stop_index(ids2, torch.tensor([3, 9])))
show("first_stop_index [2]", R.first_stop_index(ids2, torch.tensor([2])))
show("first_stop_index [100]", R.first_stop_index(ids2, torch.tensor([100])))

lg2 = torch.randn(2, 5, generator=torch.Generator().manual_seed(11))
m = torch.tensor([[1, 1, 0, 1, 0], [0, 1, 1, 0, 1]])
print("lg2", lg2.tolist())
show("masked_argmax", R.masked_argmax(lg2, m))
mf = m.float()
show("masked_argmax float", R.masked_argmax(lg2, mf))
show("topk_indices k3", R.topk_indices(lg2, 3))
show("topk_indices k1", R.topk_indices(lg2, 1))

ids3 = torch.tensor([[1, 2, 3], [4, 5, 6]])
show("right_pad_to w5 pad0", R.right_pad_to(ids3, 0, 5))
show("right_pad_to w2 pad0", R.right_pad_to(ids3, 0, 2))
show("right_pad_to w3 pad0", R.right_pad_to(ids3, 0, 3))

ids4 = torch.tensor([[5, 6, 2, 7, 8], [1, 2, 3, 4, 2]])
show("gen_stop_mask eos2 minlen0", R.generation_stop_mask(ids4, 2, 0))
show("gen_stop_mask eos2 minlen3", R.generation_stop_mask(ids4, 2, 3))
show("gen_stop_mask eos2 minlen5", R.generation_stop_mask(ids4, 2, 5))
ids5 = torch.tensor([[1, 2, 3, 4, 5]])
show("gen_stop_mask no-eos", R.generation_stop_mask(ids5, 9, 0))

# epsilon high
lg3 = torch.randn(2, 6, generator=torch.Generator().manual_seed(9))
show("epsilon_filter 0.9", R.epsilon_filter(lg3, 0.9))
show("epsilon_filter 0.99", R.epsilon_filter(lg3, 0.99))

Probe remaining functions v2

shell
python3 p.py
output
first_stop_index [3,9] -> [2, 2, -1] | torch.int64 [3]
first_stop_index [2] -> [1, 3, -1] | torch.int64 [3]
first_stop_index [100] -> [-1, -1, -1] | torch.int64 [3]
lg2 [[0.7376294732093811, 1.9459291696548462, -0.6995416283607483, -1.3022973537445068, -0.5132622122764587], [-0.26961955428123474, 0.24618466198444366, 0.4839252829551697, 0.45041778683662415, -0.9568334221839905]]
masked_argmax -> [1, 2] | torch.int64 [2]
masked_argmax float -> [1, 2] | torch.int64 [2]
topk_indices k3 -> [[1, 0, 4], [2, 3, 1]] | torch.int64 [2, 3]
topk_indices k1 -> [[1], [2]] | torch.int64 [2, 1]
right_pad_to w5 pad0 -> [[1, 2, 3, 0, 0], [4, 5, 6, 0, 0]] | torch.int64 [2, 5]
right_pad_to w2 pad0 -> [[1, 2], [4, 5]] | torch.int64 [2, 2]
right_pad_to w3 pad0 -> [[1, 2, 3], [4, 5, 6]] | torch.int64 [2, 3]
gen_stop_mask eos2 minlen0 -> [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] | torch.float32 [2, 5]
gen_stop_mask eos2 minlen3 -> [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0, 0.0]] | torch.float32 [2, 5]
gen_stop_mask eos2 minlen5 -> [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]] | torch.float32 [2, 5]
gen_stop_mask no-eos -> [[1.0, 1.0, 1.0, 1.0, 1.0]] | torch.float32 [1, 5]
epsilon_filter 0.9 -> [[-inf, 1.9112390279769897, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0873746871948242, -inf, -inf]] | torch.float32 [2, 6]
epsilon_filter 0.99 -> [[-inf, 1.9112390279769897, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0873746871948242, -inf, -inf]] | torch.float32 [2, 6]

[stdout]
first_stop_index [3,9] -> [2, 2, -1] | torch.int64 [3]
first_stop_index [2] -> [1, 3, -1] | torch.int64 [3]
first_stop_index [100] -> [-1, -1, -1] | torch.int64 [3]
lg2 [[0.7376294732093811, 1.9459291696548462, -0.6995416283607483, -1.3022973537445068, -0.5132622122764587], [-0.26961955428123474, 0.24618466198444366, 0.4839252829551697, 0.45041778683662415, -0.9568334221839905]]
masked_argmax -> [1, 2] | torch.int64 [2]
masked_argmax float -> [1, 2] | torch.int64 [2]
topk_indices k3 -> [[1, 0, 4], [2, 3, 1]] | torch.int64 [2, 3]
topk_indices k1 -> [[1], [2]] | torch.int64 [2, 1]
right_pad_to w5 pad0 -> [[1, 2, 3, 0, 0], [4, 5, 6, 0, 0]] | torch.int64 [2, 5]
right_pad_to w2 pad0 -> [[1, 2], [4, 5]] | torch.int64 [2, 2]
right_pad_to w3 pad0 -> [[1, 2, 3], [4, 5, 6]] | torch.int64 [2, 3]
gen_stop_mask eos2 minlen0 -> [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 0.0, 0.0, 0.0]] | torch.float32 [2, 5]
gen_stop_mask eos2 minlen3 -> [[1.0, 1.0, 1.0, 0.0, 0.0], [1.0, 1.0, 1.0, 0.0, 0.0]] | torch.float32 [2, 5]
gen_stop_mask eos2 minlen5 -> [[1.0, 1.0, 1.0, 1.0, 1.0], [1.0, 1.0, 1.0, 1.0, 1.0]] | torch.float32 [2, 5]
gen_stop_mask no-eos -> [[1.0, 1.0, 1.0, 1.0, 1.0]] | torch.float32 [1, 5]
epsilon_filter 0.9 -> [[-inf, 1.9112390279769897, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0873746871948242, -inf, -inf]] | torch.float32 [2, 6]
epsilon_filter 0.99 -> [[-inf, 1.9112390279769897, -inf, -inf, -inf, -inf], [-inf, -inf, -inf, 1.0873746871948242, -inf, -inf]] | torch.float32 [2, 6]
write · /app/objectives.py
from __future__ import annotations

import torch

NEG_INF = float("-inf")


def _last_valid_index(mask):
    # last position where mask is truthy, per row
    m = (mask != 0)
    length = m.shape[-1]
    flipped = m.flip(-1)
    return length - 1 - flipped.to(torch.int64).argmax(dim=-1)


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


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


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


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


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


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


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


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


def presence_penalty(logits, input_ids, penalty):
    counts = _token_counts(logits, input_ids)
    return logits - penalty * (counts > 0).to(logits.dtype)


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


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).to(torch.float32)


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


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


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


def mask_after_eos(input_ids, eos_id):
    T = input_ids.shape[-1]
    idxs = torch.arange(T, device=input_ids.device)
    eos_pos = torch.where(input_ids == eos_id, idxs, torch.full_like(input_ids, T))
    first = eos_pos.min(dim=-1).values
    cutoff = torch.clamp(first + 1, max=T)
    return (idxs.unsqueeze(0) < cutoff.unsqueeze(-1)).to(torch.float32)


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


def last_token_indices(attn_mask):
    return _last_valid_index(attn_mask)


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


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


def pad_to_multiple(input_ids, pad_id, multiple):
    T = input_ids.shape[-1]
    rem = T % multiple
    if rem == 0:
        return input_ids.clone()
    pad = multiple - rem
    padding = torch.full((input_ids.shape[0], pad), 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):
    masked = rewards * mask
    rtg = masked.flip(-1).cumsum(dim=-1).flip(-1)
    return rtg * mask


def discount_cumsum(x, gamma):
    out = torch.zeros_like(x)
    running = torch.zeros_like(x[..., 0])
    for 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):
    parts = []
    total = 0
    for l in seq_lengths.tolist():
        parts.append(torch.arange(l, dtype=torch.int64))
        total += l
    if total < total_len:
        parts.append(torch.arange(total_len - total, dtype=torch.int64))
    return torch.cat(parts)[:total_len]


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


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


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


def rollout_last_logits(logits, input_ids, pad_id):
    idx = _last_valid_index(input_ids != pad_id)
    b = torch.arange(logits.shape[0], device=logits.device)
    return logits[b, idx]


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


def first_stop_index(input_ids, stop_ids):
    stop = torch.as_tensor(stop_ids, device=input_ids.device).reshape(-1)
    is_stop = torch.isin(input_ids, stop)
    T = input_ids.shape[-1]
    idxs = torch.arange(T, device=input_ids.device)
    pos = torch.where(is_stop, idxs, torch.full_like(input_ids, T))
    first = pos.min(dim=-1).values
    return torch.where(first < T, first, torch.full_like(first, -1))


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


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


def right_pad_to(input_ids, pad_id, width):
    T = input_ids.shape[-1]
    if T >= width:
        return input_ids[..., :width].clone()
    padding = torch.full((input_ids.shape[0], width - T), 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):
    T = input_ids.shape[-1]
    idxs = torch.arange(T, device=input_ids.device)
    eos_pos = torch.where(input_ids == eos_id, idxs, torch.full_like(input_ids, T))
    first = eos_pos.min(dim=-1).values
    has = first < T
    base = torch.where(has, first + 1, torch.full_like(first, T))
    cutoff = torch.clamp(torch.clamp(base, min=min_len), max=T)
    return (idxs.unsqueeze(0) < cutoff.unsqueeze(-1)).to(torch.float32)
write · /app/verify.py
from __future__ import annotations
import torch
import reforacle as R
import importlib
import objectives as O
importlib.reload(O)

fails = []

def cmp(name, a, b, atol=1e-5, rtol=1e-4):
    try:
        if torch.is_tensor(a):
            if a.dtype != b.dtype:
                fails.append((name, f"dtype {a.dtype} vs {b.dtype}"))
                return
            if a.shape != b.shape:
                fails.append((name, f"shape {a.shape} vs {b.shape}"))
                return
            if a.is_floating_point():
                if not torch.allclose(a, b, atol=atol, rtol=rtol, equal_nan=True):
                    fails.append((name, f"vals maxdiff {(a-b).abs().max().item() if torch.isfinite(a-b).all() else 'inf/nan'}"))
            else:
                if not torch.equal(a, b):
                    fails.append((name, "int mismatch"))
        else:
            if a != b:
                fails.append((name, f"{a} vs {b}"))
    except Exception as e:
        fails.append((name, f"CMP EXC {e!r}"))

def run(name, fn, *args):
    try:
        mine = fn(*args)
    except Exception as e:
        fails.append((name, f"MINE EXC {e!r}"))
        return
    try:
        ref = getattr(R, name)(*[a.clone() if torch.is_tensor(a) else a for a in args])
    except Exception as e:
        fails.append((name, f"REF EXC {e!r}"))
        return
    cmp(name, mine, ref)

seed = 0
for trial in range(200):
    g = torch.Generator().manual_seed(seed); seed += 1
    B = int(torch.randint(1, 5, (1,), generator=g).item())
    V = int(torch.randint(2, 10, (1,), generator=g).item())
    T = int(torch.randint(2, 8, (1,), generator=g).item())
    logits = torch.randn(B, V, generator=g)
    run("greedy_token", O.greedy_token, logits)
    run("temperature_scale", O.temperature_scale, logits, float(torch.rand(1, generator=g).item()*3+0.1))
    run("temperature_scale", O.temperature_scale, logits, 0.0)
    run("top_k_filter", O.top_k_filter, logits, int(torch.randint(1, V+2, (1,), generator=g).item()))
    run("top_p_filter", O.top_p_filter, logits, float(torch.rand(1, generator=g).item()))
    run("min_p_filter", O.min_p_filter, logits, float(torch.rand(1, generator=g).item()))
    ids = torch.randint(0, V, (B, T), generator=g)
    run("repetition_penalty", O.repetition_penalty, logits, ids, float(torch.rand(1, generator=g).item()*2+0.5))
    run("frequency_penalty", O.frequency_penalty, logits, ids, float(torch.rand(1, generator=g).item()))
    run("presence_penalty", O.presence_penalty, logits, ids, float(torch.rand(1, generator=g).item()))
    tk = torch.randint(0, V, (int(torch.randint(1,3,(1,),generator=g).item()),), generator=g)
    run("suppress_tokens", O.suppress_tokens, logits, tk)
    run("min_length_filter", O.min_length_filter, logits, int(torch.randint(0,6,(1,),generator=g).item()), 4, int(torch.randint(0,V,(1,),generator=g).item()))
    run("renormalize_logits", O.renormalize_logits, logits)
    pid = 0
    ids2 = torch.randint(0, V, (B, T), generator=g)
    # inject some pads
    ids2 = ids2 * (torch.rand(B, T, generator=g) > 0.3).long()
    run("attention_mask", O.attention_mask, ids2, pid)
    run("position_ids", O.position_ids, ids2, pid)
    run("causal_mask", O.causal_mask, T)
    run("shift_labels", O.shift_labels, ids2, -100)
    pl = torch.randint(1, T+1, (B,), generator=g)
    run("completion_loss_labels", O.completion_loss_labels, ids2, pl, -100)
    run("prepare_labels", O.prepare_labels, ids2, pl, -100)
    run("mask_after_eos", O.mask_after_eos, ids2, int(torch.randint(0,V,(1,),generator=g).item()))
    run("sequence_lengths", O.sequence_lengths, ids2, pid)
    am = (torch.rand(B, T, generator=g) > 0.4).long()
    am[:, 0] = 1
    run("last_token_indices", O.last_token_indices, am)
    logits3 = torch.randn(B, T, V, generator=g)
    run("gather_last_logits", O.gather_last_logits, logits3, am)
    run("truncate_left", O.truncate_left, ids2, int(torch.randint(1, T+2,(1,),generator=g).item()))
    run("pad_to_multiple", O.pad_to_multiple, ids2, pid, int(torch.randint(1,5,(1,),generator=g).item()))
    acts = torch.randint(0, V, (B, T), generator=g)
    run("action_logprobs", O.action_logprobs, logits3, acts)
    run("entropy_from_logits", O.entropy_from_logits, logits3)
    rew = torch.randn(B, T, generator=g)
    mk = (torch.rand(B, T, generator=g) > 0.3).float()
    run("reward_to_go", O.reward_to_go, rew, mk)
    run("discount_cumsum", O.discount_cumsum, rew, float(torch.rand(1, generator=g).item()))
    sl = torch.randint(1, 4, (int(torch.randint(1,4,(1,),generator=g).item()),), generator=g)
    tot = int(sl.sum().item()) + int(torch.randint(0,3,(1,),generator=g).item())
    run("build_position_ids_packed", O.build_position_ids_packed, sl, tot)
    run("logits_warper", O.logits_warper, logits, float(torch.rand(1,generator=g).item()*2+0.1), int(torch.randint(1,V+1,(1,),generator=g).item()), float(torch.rand(1,generator=g).item()*0.5+0.5))
    run("sample_greedy_warped", O.sample_greedy_warped, logits, float(torch.rand(1,generator=g).item()*2+0.1), int(torch.randint(1,V+1,(1,),generator=g).item()), float(torch.rand(1,generator=g).item()*0.5+0.5))
    run("rollout_last_logits", O.rollout_last_logits, logits3, ids2, pid)
    run("epsilon_filter", O.epsilon_filter, logits, float(torch.rand(1,generator=g).item()*0.4))
    stp = torch.randint(0, V, (int(torch.randint(1,3,(1,),generator=g).item()),), generator=g)
    run("first_stop_index", O.first_stop_index, ids2, stp)
    run("masked_argmax", O.masked_argmax, logits, am[:, :V] if am.shape[1]>=V else torch.cat([am, torch.ones(B, V-am.shape[1]).long()],1))
    run("topk_indices", O.topk_indices, logits, int(torch.randint(1,V+1,(1,),generator=g).item()))
    run("right_pad_to", O.right_pad_to, ids2, pid, int(torch.randint(1,T+3,(1,),generator=g).item()))
    run("generation_stop_mask", O.generation_stop_mask, ids2, int(torch.randint(0,V,(1,),generator=g).item()), int(torch.randint(0,T+2,(1,),generator=g).item()))

from collections import Counter
c = Counter(n for n, _ in fails)
print("total fails:", len(fails))
for n in sorted(c):
    ex = next(m for nn, m in fails if nn == n)
    print(f"  {n}: {c[n]} e.g. {ex}")
if not fails:
    print("ALL PASS")

Run verification harness

shell
python3 verify.py
output
total fails: 104
  build_position_ids_packed: 66 e.g. int mismatch
  rollout_last_logits: 38 e.g. vals maxdiff 2.545407295227051

[stdout]
total fails: 104
  build_position_ids_packed: 66 e.g. int mismatch
  rollout_last_logits: 38 e.g. vals maxdiff 2.545407295227051
write · /app/p.py
from __future__ import annotations
import torch
import reforacle as R

def show(name, val):
    print(name, "->", val.tolist() if torch.is_tensor(val) else val)

# rollout_last_logits with interior pads
g = torch.Generator().manual_seed(2)
logits = torch.randn(2, 5, 3, generator=g)
ids = torch.tensor([[1, 0, 3, 0, 0], [0, 5, 0, 0, 0]])
print("ids", ids.tolist())
show("rollout ref", R.rollout_last_logits(logits, ids, 0))
print("last-nonpad r0=idx2", logits[0,2].tolist())
print("count-1: r0 nonpad=2 ->idx1", logits[0,1].tolist())
print("r1 last-nonpad idx1", logits[1,1].tolist(), "count-1 idx0", logits[1,0].tolist())

# build_position_ids_packed edge cases
show("bpp [3] tot3", R.build_position_ids_packed(torch.tensor([3]), 3))
show("bpp [2] tot5", R.build_position_ids_packed(torch.tensor([2]), 5))
show("bpp [2,2] tot7", R.build_position_ids_packed(torch.tensor([2,2]), 7))
show("bpp [4] tot2", R.build_position_ids_packed(torch.tensor([4]), 2))
show("bpp [1,1,1] tot3", R.build_position_ids_packed(torch.tensor([1,1,1]), 3))
show("bpp [3,3] tot4", R.build_position_ids_packed(torch.tensor([3,3]), 4))

Probe two failing functions

shell
python3 p.py
output
ids [[1, 0, 3, 0, 0], [0, 5, 0, 0, 0]]
rollout ref -> [[-1.0573921203613281, -0.11883937567472458, -0.9078398942947388], [-0.27840104699134827, -0.2719452679157257, -1.2944108247756958]]
last-nonpad r0=idx2 [-1.0573921203613281, -0.11883937567472458, -0.9078398942947388]
count-1: r0 nonpad=2 ->idx1 [-1.1096783876419067, -1.2187780141830444, 1.1675924062728882]
r1 last-nonpad idx1 [-0.27840104699134827, -0.2719452679157257, -1.2944108247756958] count-1 idx0 [0.09373074769973755, -0.7649639248847961, 1.8298852443695068]
bpp [3] tot3 -> [0, 1, 2]
bpp [2] tot5 -> [0, 1, 0, 0, 0]
bpp [2,2] tot7 -> [0, 1, 0, 1, 0, 0, 0]
bpp [4] tot2 -> [0, 1]
bpp [1,1,1] tot3 -> [0, 0, 0]
bpp [3,3] tot4 -> [0, 1, 2, 0]

[stdout]
ids [[1, 0, 3, 0, 0], [0, 5, 0, 0, 0]]
rollout ref -> [[-1.0573921203613281, -0.11883937567472458, -0.9078398942947388], [-0.27840104699134827, -0.2719452679157257, -1.2944108247756958]]
last-nonpad r0=idx2 [-1.0573921203613281, -0.11883937567472458, -0.9078398942947388]
count-1: r0 nonpad=2 ->idx1 [-1.1096783876419067, -1.2187780141830444, 1.1675924062728882]
r1 last-nonpad idx1 [-0.27840104699134827, -0.2719452679157257, -1.2944108247756958] count-1 idx0 [0.09373074769973755, -0.7649639248847961, 1.8298852443695068]
bpp [3] tot3 -> [0, 1, 2]
bpp [2] tot5 -> [0, 1, 0, 0, 0]
bpp [2,2] tot7 -> [0, 1, 0, 1, 0, 0, 0]
bpp [4] tot2 -> [0, 1]
bpp [1,1,1] tot3 -> [0, 0, 0]
bpp [3,3] tot4 -> [0, 1, 2, 0]

/app/objectives.py

− removed
    if total < total_len:
        parts.append(torch.arange(total_len - total, dtype=torch.int64))
    return torch.cat(parts)[:total_len]
+ added
    if total < total_len:
        parts.append(torch.zeros(total_len - total, dtype=torch.int64))
    return torch.cat(parts)[:total_len]
write · /app/p.py
from __future__ import annotations
import torch
import reforacle as R
import objectives as O

# find rollout failing cases
g = torch.Generator().manual_seed(0)
for i in range(500):
    B = int(torch.randint(1,4,(1,),generator=g).item())
    T = int(torch.randint(2,6,(1,),generator=g).item())
    V = int(torch.randint(2,5,(1,),generator=g).item())
    logits = torch.randn(B,T,V,generator=g)
    ids = (torch.randint(0,V,(B,T),generator=g) * (torch.rand(B,T,generator=g)>0.3).long())
    ref = R.rollout_last_logits(logits.clone(), ids.clone(), 0)
    mine = O.rollout_last_logits(logits.clone(), ids.clone(), 0)
    if not torch.allclose(ref, mine):
        print("ids", ids.tolist())
        print("ref", ref.tolist())
        print("mine", mine.tolist())
        # show which index ref picked
        for r in range(B):
            for t in range(T):
                if torch.allclose(logits[r,t], ref[r]):
                    print(f"row{r} ref idx {t}")
        break

Find rollout failing case

shell
python3 p.py
output
ids [[0, 0, 0], [1, 0, 0], [1, 0, 0]]
ref [[-0.01595325581729412, 1.2756527662277222], [0.3093254268169403, 0.36422988772392273], [-0.4551093578338623, -2.245525360107422]]
mine [[0.3423732817173004, 0.008177380077540874], [0.3093254268169403, 0.36422988772392273], [-0.4551093578338623, -2.245525360107422]]
row0 ref idx 0
row1 ref idx 0
row2 ref idx 0

[stdout]
ids [[0, 0, 0], [1, 0, 0], [1, 0, 0]]
ref [[-0.01595325581729412, 1.2756527662277222], [0.3093254268169403, 0.36422988772392273], [-0.4551093578338623, -2.245525360107422]]
mine [[0.3423732817173004, 0.008177380077540874], [0.3093254268169403, 0.36422988772392273], [-0.4551093578338623, -2.245525360107422]]
row0 ref idx 0
row1 ref idx 0
row2 ref idx 0

/app/objectives.py

− removed
def _last_valid_index(mask):
    # last position where mask is truthy, per row
    m = (mask != 0)
    length = m.shape[-1]
    flipped = m.flip(-1)
    return length - 1 - flipped.to(torch.int64).argmax(dim=-1)
+ added
def _last_valid_index(mask):
    # last position where mask is truthy, per row; 0 if none are truthy
    m = (mask != 0)
    idxs = torch.arange(m.shape[-1], device=m.device)
    positions = torch.where(m, idxs, torch.full_like(idxs, -1))
    return positions.max(dim=-1).values.clamp(min=0)

Check all-pad convention for last_token_indices

shell
python3 -c "
import torch, reforacle as R, objectives as O
# all-pad for last_token_indices
am = torch.tensor([[0,0,0],[1,0,1],[0,0,0]])
print('ref lti', R.last_token_indices(am.clone()).tolist())
print('mine   ', O.last_token_indices(am.clone()).tolist())
logits=torch.randn(3,3,4)
print('ref gll', R.gather_last_logits(logits.clone(), am.clone()).tolist())
print('mine   ', O.gather_last_logits(logits.clone(), am.clone()).tolist())
"
output
ref lti [-1, 2, -1]
mine    [0, 2, 0]
ref gll [[-1.0656205415725708, -0.2859868109226227, -0.004184220917522907, -0.11185699701309204], [0.0979437455534935, -0.34311434626579285, -0.5564369559288025, -0.8415773510932922], [0.7099857926368713, 0.13131089508533478, 0.10408353060483932, 1.1693272590637207]]
mine    [[-1.0656205415725708, -0.2859868109226227, -0.004184220917522907, -0.11185699701309204], [0.0979437455534935, -0.34311434626579285, -0.5564369559288025, -0.8415773510932922], [0.7099857926368713, 0.13131089508533478, 0.10408353060483932, 1.1693272590637207]]

[stdout]
ref lti [-1, 2, -1]
mine    [0, 2, 0]
ref gll [[-1.0656205415725708, -0.2859868109226227, -0.004184220917522907, -0.11185699701309204], [0.0979437455534935, -0.34311434626579285, -0.5564369559288025, -0.8415773510932922], [0.7099857926368713, 0.13131089508533478, 0.10408353060483932, 1.1693272590637207]]
mine    [[-1.0656205415725708, -0.2859868109226227, -0.004184220917522907, -0.11185699701309204], [0.0979437455534935, -0.34311434626579285, -0.5564369559288025, -0.8415773510932922], [0.7099857926368713, 0.13131089508533478, 0.10408353060483932, 1.1693272590637207]]

/app/objectives.py

− removed
def last_token_indices(attn_mask):
    return _last_valid_index(attn_mask)
+ added
def last_token_indices(attn_mask):
    m = (attn_mask != 0)
    idxs = torch.arange(m.shape[-1], device=m.device)
    positions = torch.where(m, idxs, torch.full_like(idxs, -1))
    return positions.max(dim=-1).values

Explicit all-pad gather test

shell
python3 -c "
import torch, reforacle as R, objectives as O
# explicit all-pad gather test
logits=torch.arange(3*3*2).float().reshape(3,3,2)
am=torch.tensor([[0,0,0],[1,0,0],[1,1,1]])
print('logits row0', logits[0].tolist())
print('ref gll', R.gather_last_logits(logits.clone(), am.clone()).tolist())
print('mine   ', O.gather_last_logits(logits.clone(), am.clone()).tolist())
print('ref roll', R.rollout_last_logits(logits.clone(), torch.tensor([[0,0,0],[5,0,0],[5,5,5]]), 0).tolist())
print('mine roll', O.rollout_last_logits(logits.clone(), torch.tensor([[0,0,0],[5,0,0],[5,5,5]]), 0).tolist())
"
output
logits row0 [[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]]
ref gll [[0.0, 1.0], [6.0, 7.0], [16.0, 17.0]]
mine    [[0.0, 1.0], [6.0, 7.0], [16.0, 17.0]]
ref roll [[0.0, 1.0], [6.0, 7.0], [16.0, 17.0]]
mine roll [[0.0, 1.0], [6.0, 7.0], [16.0, 17.0]]

[stdout]
logits row0 [[0.0, 1.0], [2.0, 3.0], [4.0, 5.0]]
ref gll [[0.0, 1.0], [6.0, 7.0], [16.0, 17.0]]
mine    [[0.0, 1.0], [6.0, 7.0], [16.0, 17.0]]
ref roll [[0.0, 1.0], [6.0, 7.0], [16.0, 17.0]]
mine roll [[0.0, 1.0], [6.0, 7.0], [16.0, 17.0]]

Re-run full verification

shell
python3 verify.py
output
total fails: 0
ALL PASS

[stdout]
total fails: 0
ALL PASS
write · /app/verify2.py
from __future__ import annotations
import torch
import reforacle as R
import objectives as O
from collections import Counter

fails = []

def cmp(name, a, b):
    if torch.is_tensor(a):
        if a.dtype != b.dtype:
            fails.append((name, f"dtype {a.dtype} vs {b.dtype}")); return
        if a.shape != b.shape:
            fails.append((name, f"shape {tuple(a.shape)} vs {tuple(b.shape)}")); return
        if a.is_floating_point():
            if not torch.allclose(a, b, atol=1e-5, rtol=1e-4, equal_nan=True):
                fails.append((name, "vals")); return
        else:
            if not torch.equal(a, b):
                fails.append((name, f"int {a.tolist()} vs {b.tolist()}")); return
    else:
        if a != b: fails.append((name, f"{a} vs {b}"))

def run(name, *args):
    try: mine = getattr(O, name)(*[a.clone() if torch.is_tensor(a) else a for a in args])
    except Exception as e: fails.append((name, f"MINE {e!r}")); return
    try: ref = getattr(R, name)(*[a.clone() if torch.is_tensor(a) else a for a in args])
    except Exception as e: fails.append((name, f"REF {e!r}")); return
    cmp(name, mine, ref)

g = torch.Generator().manual_seed(123)
for trial in range(400):
    B = int(torch.randint(1,6,(1,),generator=g).item())
    V = int(torch.randint(1,12,(1,),generator=g).item())
    T = int(torch.randint(1,10,(1,),generator=g).item())
    scale = float(torch.rand(1,generator=g).item()*10+0.01)
    logits = torch.randn(B, V, generator=g) * scale
    # inject ties and extremes occasionally
    if trial % 5 == 0 and V >= 2:
        logits[:, 0] = logits[:, 1]
    run("greedy_token", logits)
    run("temperature_scale", logits, float(torch.rand(1,generator=g).item()*5))
    run("temperature_scale", logits, 0.0)
    run("top_k_filter", logits, int(torch.randint(1,V+3,(1,),generator=g).item()))
    run("top_p_filter", logits, float(torch.rand(1,generator=g).item()))
    run("top_p_filter", logits, 1.0)
    run("min_p_filter", logits, float(torch.rand(1,generator=g).item()))
    ids = torch.randint(0, V, (B, T), generator=g)
    run("repetition_penalty", logits, ids, float(torch.rand(1,generator=g).item()*3+0.3))
    run("frequency_penalty", logits, ids, float(torch.rand(1,generator=g).item()*2))
    run("presence_penalty", logits, ids, float(torch.rand(1,generator=g).item()*2))
    ntk = int(torch.randint(1,4,(1,),generator=g).item())
    run("suppress_tokens", logits, torch.randint(0,V,(ntk,),generator=g))
    run("suppress_tokens", logits, list(torch.randint(0,V,(ntk,),generator=g).tolist()))
    run("min_length_filter", logits, int(torch.randint(0,8,(1,),generator=g).item()), int(torch.randint(0,8,(1,),generator=g).item()), int(torch.randint(0,V,(1,),generator=g).item()))
    run("renormalize_logits", logits)
    # heavy pad injection including all-pad rows
    keep = (torch.rand(B, T, generator=g) > 0.5).long()
    if trial % 7 == 0: keep[0] = 0
    ids2 = torch.randint(1, V+1, (B, T), generator=g) * keep  # pad_id=0
    run("attention_mask", ids2, 0)
    run("position_ids", ids2, 0)
    run("causal_mask", T)
    run("shift_labels", ids2, -100)
    pl = torch.randint(0, T+1, (B,), generator=g)
    run("completion_loss_labels", ids2, pl, -100)
    run("prepare_labels", ids2, pl, -100)
    run("mask_after_eos", ids2, int(torch.randint(0,V+1,(1,),generator=g).item()))
    run("sequence_lengths", ids2, 0)
    am = keep.clone()
    run("last_token_indices", am)
    logits3 = torch.randn(B, T, V, generator=g)
    run("gather_last_logits", logits3, am)
    run("truncate_left", ids2, int(torch.randint(1,T+3,(1,),generator=g).item()))
    run("pad_to_multiple", ids2, 0, int(torch.randint(1,6,(1,),generator=g).item()))
    acts = torch.randint(0, V, (B, T), generator=g)
    run("action_logprobs", logits3, acts)
    run("entropy_from_logits", logits3)
    rew = torch.randn(B, T, generator=g)
    mk = (torch.rand(B, T, generator=g) > 0.4).float()
    run("reward_to_go", rew, mk)
    run("discount_cumsum", rew, float(torch.rand(1,generator=g).item()))
    run("discount_cumsum", rew[0], float(torch.rand(1,generator=g).item()))
    nseg = int(torch.randint(1,5,(1,),generator=g).item())
    sl = torch.randint(0,4,(nseg,),generator=g)
    tot = int(sl.sum().item()) + int(torch.randint(-2,4,(1,),generator=g).item())
    tot = max(1, tot)
    run("build_position_ids_packed", sl, tot)
    run("logits_warper", logits, float(torch.rand(1,generator=g).item()*3+0.1), int(torch.randint(1,V+2,(1,),generator=g).item()), float(torch.rand(1,generator=g).item()))
    run("sample_greedy_warped", logits, float(torch.rand(1,generator=g).item()*3+0.1), int(torch.randint(1,V+2,(1,),generator=g).item()), float(torch.rand(1,generator=g).item()))
    run("rollout_last_logits", logits3, ids2, 0)
    run("epsilon_filter", logits, float(torch.rand(1,generator=g).item()))
    run("first_stop_index", ids2, torch.randint(0,V+1,(ntk,),generator=g))
    m2 = (torch.rand(B, V, generator=g) > 0.4).long()
    m2[:, 0] = 1
    run("masked_argmax", logits, m2)
    run("masked_argmax", logits, m2.float())
    run("topk_indices", logits, int(torch.randint(1,V+1,(1,),generator=g).item()))
    run("right_pad_to", ids2, 0, int(torch.randint(1,T+4,(1,),generator=g).item()))
    run("generation_stop_mask", ids2, int(torch.randint(0,V+1,(1,),generator=g).item()), int(torch.randint(0,T+3,(1,),generator=g).item()))

c = Counter(n for n,_ in fails)
print("total fails:", len(fails))
for n in sorted(c):
    print(f"  {n}: {c[n]} e.g. {next(m for nn,m in fails if nn==n)}")
if not fails: print("ALL PASS")

Run aggressive verification

shell
python3 verify2.py
output
total fails: 340
  epsilon_filter: 17 e.g. vals
  greedy_token: 31 e.g. int [2, 2, 0, 0] vs [2, 2, 1, 1]
  masked_argmax: 76 e.g. int [2, 2, 0, 0] vs [2, 2, 1, 1]
  sample_greedy_warped: 20 e.g. int [2, 2, 0, 0] vs [2, 2, 1, 1]
  temperature_scale: 31 e.g. vals
  top_p_filter: 121 e.g. vals
  topk_indices: 44 e.g. int [[3, 1, 0], [2, 3, 1], [3, 2, 0]] vs [[3, 0, 1], [2, 3, 0], [3, 2, 0]]

[stdout]
total fails: 340
  epsilon_filter: 17 e.g. vals
  greedy_token: 31 e.g. int [2, 2, 0, 0] vs [2, 2, 1, 1]
  masked_argmax: 76 e.g. int [2, 2, 0, 0] vs [2, 2, 1, 1]
  sample_greedy_warped: 20 e.g. int [2, 2, 0, 0] vs [2, 2, 1, 1]
  temperature_scale: 31 e.g. vals
  top_p_filter: 121 e.g. vals
  topk_indices: 44 e.g. int [[3, 1, 0], [2, 3, 1], [3, 2, 0]] vs [[3, 0, 1], [2, 3, 0], [3, 2, 0]]

Probe tie-breaking

shell
python3 -c "
import torch, reforacle as R
# argmax tie-breaking
x = torch.tensor([[5.0,5.0,1.0],[1.0,5.0,5.0],[5.0,1.0,5.0],[5.0,5.0,5.0]])
print('greedy', R.greedy_token(x).tolist())
print('torch.argmax', x.argmax(-1).tolist())
print('torch.max idx', x.max(-1).indices.tolist())
# temperature T=0 tie
print('tempT0', R.temperature_scale(x.clone(),0.0).tolist())
# masked_argmax tie
m=torch.tensor([[1,1,1],[1,1,1],[1,1,1],[1,1,1]])
print('masked_argmax', R.masked_argmax(x.clone(), m).tolist())
# topk tie ordering
print('topk_indices k2', R.topk_indices(x.clone(),2).tolist())
print('torch.topk', torch.topk(x,2,-1).indices.tolist())
"
output
greedy [1, 2, 2, 2]
torch.argmax [0, 1, 0, 0]
torch.max idx [0, 1, 0, 0]
tempT0 [[-inf, 0.0, -inf], [-inf, -inf, 0.0], [-inf, -inf, 0.0], [-inf, -inf, 0.0]]
masked_argmax [1, 2, 2, 2]
topk_indices k2 [[0, 1], [1, 2], [0, 2], [0, 1]]
torch.topk [[0, 1], [1, 2], [0, 2], [0, 1]]

[stdout]
greedy [1, 2, 2, 2]
torch.argmax [0, 1, 0, 0]
torch.max idx [0, 1, 0, 0]
tempT0 [[-inf, 0.0, -inf], [-inf, -inf, 0.0], [-inf, -inf, 0.0], [-inf, -inf, 0.0]]
masked_argmax [1, 2, 2, 2]
topk_indices k2 [[0, 1], [1, 2], [0, 2], [0, 1]]
torch.topk [[0, 1], [1, 2], [0, 2], [0, 1]]

Isolate non-tie top_p/epsilon failures

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(7)
# no ties, but extreme scale
bad=0
for i in range(2000):
    B=int(torch.randint(1,5,(1,),generator=g).item()); V=int(torch.randint(2,12,(1,),generator=g).item())
    scale=float(torch.rand(1,generator=g).item()*20+0.01)
    x=torch.randn(B,V,generator=g)*scale
    for fn,args in [('top_p_filter',(float(torch.rand(1,generator=g).item()),)),
                    ('epsilon_filter',(float(torch.rand(1,generator=g).item()),)),
                    ('temperature_scale',(0.0,)),
                    ('min_p_filter',(float(torch.rand(1,generator=g).item()),))]:
        a=getattr(O,fn)(x.clone(),*args); b=getattr(R,fn)(x.clone(),*args)
        if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True):
            bad+=1
            if bad<=5: print(fn,'scale',round(scale,2),'args',args); print(' mine',a.tolist()); print(' ref ',b.tolist())
print('bad',bad)
"
output
bad 0

[stdout]
bad 0

/app/objectives.py

− removed
def greedy_token(logits):
    return logits.argmax(dim=-1)
+ added
def _last_argmax(x):
    # argmax that breaks ties toward the highest index
    length = x.shape[-1]
    return length - 1 - x.flip(-1).argmax(dim=-1)


def greedy_token(logits):
    return _last_argmax(logits)

/app/objectives.py

− removed
    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
    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
def sample_greedy_warped(logits, temperature, top_k, top_p):
    return logits_warper(logits, temperature, top_k, top_p).argmax(dim=-1)
+ added
def sample_greedy_warped(logits, temperature, top_k, top_p):
    return _last_argmax(logits_warper(logits, temperature, top_k, top_p))

/app/objectives.py

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

Re-run aggressive verification

shell
python3 verify2.py
output
total fails: 182
  epsilon_filter: 17 e.g. vals
  top_p_filter: 121 e.g. vals
  topk_indices: 44 e.g. int [[3, 1, 0], [2, 3, 1], [3, 2, 0]] vs [[3, 0, 1], [2, 3, 0], [3, 2, 0]]

[stdout]
total fails: 182
  epsilon_filter: 17 e.g. vals
  top_p_filter: 121 e.g. vals
  topk_indices: 44 e.g. int [[3, 1, 0], [2, 3, 1], [3, 2, 0]] vs [[3, 0, 1], [2, 3, 0], [3, 2, 0]]

Test stable-sort conventions for topk and top_p

shell
python3 -c "
import torch, reforacle as R, objectives as O
# topk: stable sort?
g=torch.Generator().manual_seed(4)
bad_topk=0; bad_stable=0
for i in range(3000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(2,7,(1,),generator=g).item())
    x=torch.randint(0,4,(B,V),generator=g).float()  # lots of ties
    k=int(torch.randint(1,V+1,(1,),generator=g).item())
    ref=R.topk_indices(x.clone(),k)
    stable=torch.sort(x,dim=-1,descending=True,stable=True).indices[...,:k]
    if not torch.equal(ref,stable): bad_stable+=1
print('topk vs stable-sort mismatches:', bad_stable)
# top_p stable?
bad=0
for i in range(3000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(2,7,(1,),generator=g).item())
    x=torch.randint(0,3,(B,V),generator=g).float()
    p=float(torch.rand(1,generator=g).item())
    ref=R.top_p_filter(x.clone(),p)
    # stable version
    sl,si=torch.sort(x,descending=True,dim=-1,stable=True)
    cum=sl.softmax(-1).cumsum(-1); rem=cum>p; rem[...,1:]=rem[...,:-1].clone(); rem[...,0]=False
    rem=rem.scatter(-1,si,rem); out=x.masked_fill(rem,float('-inf'))
    if not torch.allclose(ref,out,equal_nan=True): bad+=1
print('top_p stable mismatches:', bad)
"
output
topk vs stable-sort mismatches: 0
top_p stable mismatches: 0

[stdout]
topk vs stable-sort mismatches: 0
top_p stable mismatches: 0

/app/objectives.py

− removed
    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+ added
    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1, stable=True)

/app/objectives.py

− removed
def topk_indices(logits, k):
    return torch.topk(logits, int(k), dim=-1).indices
+ added
def topk_indices(logits, k):
    return torch.sort(logits, dim=-1, descending=True, stable=True).indices[..., :int(k)]

Investigate epsilon tie behavior

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(4)
bad=0; ex=0
for i in range(5000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(2,7,(1,),generator=g).item())
    x=torch.randint(0,3,(B,V),generator=g).float()
    eps=float(torch.rand(1,generator=g).item())
    ref=R.epsilon_filter(x.clone(),eps)
    mine=O.epsilon_filter(x.clone(),eps)
    if not torch.allclose(ref,mine,equal_nan=True):
        bad+=1
        if ex<4:
            ex+=1; print('x',x.tolist(),'eps',round(eps,3)); print(' ref',ref.tolist()); print(' mine',mine.tolist())
print('bad',bad)
"
output
x [[1.0, 2.0, 0.0, 2.0, 2.0, 1.0], [0.0, 0.0, 1.0, 2.0, 2.0, 1.0]] eps 0.595
 ref [[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
 mine [[-inf, 2.0, -inf, 2.0, 2.0, -inf], [-inf, -inf, -inf, 2.0, 2.0, -inf]]
x [[1.0, 0.0, 2.0, 1.0], [2.0, 2.0, 1.0, 2.0]] eps 0.543
 ref [[-inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 2.0]]
 mine [[-inf, -inf, 2.0, -inf], [2.0, 2.0, -inf, 2.0]]
x [[0.0, 1.0, 0.0, 2.0], [0.0, 0.0, 2.0, 0.0], [1.0, 1.0, 0.0, 0.0]] eps 0.654
 ref [[-inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, -inf], [-inf, 1.0, -inf, -inf]]
 mine [[-inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, -inf], [1.0, 1.0, -inf, -inf]]
x [[2.0, 2.0, 2.0, 0.0, 1.0, 2.0], [2.0, 1.0, 0.0, 2.0, 2.0, 1.0], [1.0, 0.0, 2.0, 0.0, 1.0, 2.0]] eps 0.817
 ref [[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
 mine [[2.0, 2.0, 2.0, -inf, -inf, 2.0], [2.0, -inf, -inf, 2.0, 2.0, -inf], [-inf, -inf, 2.0, -inf, -inf, 2.0]]
bad 2353

[stdout]
x [[1.0, 2.0, 0.0, 2.0, 2.0, 1.0], [0.0, 0.0, 1.0, 2.0, 2.0, 1.0]] eps 0.595
 ref [[-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, 2.0, -inf]]
 mine [[-inf, 2.0, -inf, 2.0, 2.0, -inf], [-inf, -inf, -inf, 2.0, 2.0, -inf]]
x [[1.0, 0.0, 2.0, 1.0], [2.0, 2.0, 1.0, 2.0]] eps 0.543
 ref [[-inf, -inf, 2.0, -inf], [-inf, -inf, -inf, 2.0]]
 mine [[-inf, -inf, 2.0, -inf], [2.0, 2.0, -inf, 2.0]]
x [[0.0, 1.0, 0.0, 2.0], [0.0, 0.0, 2.0, 0.0], [1.0, 1.0, 0.0, 0.0]] eps 0.654
 ref [[-inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, -inf], [-inf, 1.0, -inf, -inf]]
 mine [[-inf, -inf, -inf, 2.0], [-inf, -inf, 2.0, -inf], [1.0, 1.0, -inf, -inf]]
x [[2.0, 2.0, 2.0, 0.0, 1.0, 2.0], [2.0, 1.0, 0.0, 2.0, 2.0, 1.0], [1.0, 0.0, 2.0, 0.0, 1.0, 2.0]] eps 0.817
 ref [[-inf, -inf, -inf, -inf, -inf, 2.0], [-inf, -inf, -inf, -inf, 2.0, -inf], [-inf, -inf, -inf, -inf, -inf, 2.0]]
 mine [[2.0, 2.0, 2.0, -inf, -inf, 2.0], [2.0, -inf, -inf, 2.0, 2.0, -inf], [-inf, -inf, 2.0, -inf, -inf, 2.0]]
bad 2353

/app/objectives.py

− removed
def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    remove = probs < epsilon
    top1 = logits.max(dim=-1, keepdim=True).values
    remove = remove & (logits < top1)
    return logits.masked_fill(remove, NEG_INF)
+ added
def epsilon_filter(logits, epsilon):
    probs = logits.softmax(dim=-1)
    remove = probs < epsilon
    keep_idx = _last_argmax(logits).unsqueeze(-1)
    remove.scatter_(-1, keep_idx, False)
    return logits.masked_fill(remove, NEG_INF)

Test filter tie behaviors

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(4)
for fn,mk in [('epsilon_filter',lambda:(float(torch.rand(1,generator=g).item()),)),
              ('top_k_filter',lambda:(int(torch.randint(1,6,(1,),generator=g).item()),)),
              ('min_p_filter',lambda:(float(torch.rand(1,generator=g).item()),)),
              ('top_p_filter',lambda:(float(torch.rand(1,generator=g).item()),))]:
    bad=0
    for i in range(4000):
        B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(2,7,(1,),generator=g).item())
        x=torch.randint(0,3,(B,V),generator=g).float()
        a=mk()
        if not torch.allclose(getattr(O,fn)(x.clone(),*a),getattr(R,fn)(x.clone(),*a),equal_nan=True): bad+=1
    print(fn,'tie-bad',bad)
"
output
epsilon_filter tie-bad 0
top_k_filter tie-bad 0
min_p_filter tie-bad 0
top_p_filter tie-bad 0

[stdout]
epsilon_filter tie-bad 0
top_k_filter tie-bad 0
min_p_filter tie-bad 0
top_p_filter tie-bad 0

Run both verification harnesses

shell
python3 verify2.py && echo "---" && python3 verify.py
output
total fails: 121
  top_p_filter: 121 e.g. vals
---
total fails: 0
ALL PASS

[stdout]
total fails: 121
  top_p_filter: 121 e.g. vals
---
total fails: 0
ALL PASS

Reproduce top_p continuous tie failure

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(999)
bad=0;ex=0
for i in range(20000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(2,8,(1,),generator=g).item())
    scale=float(torch.rand(1,generator=g).item()*10+0.01)
    x=torch.randn(B,V,generator=g)*scale
    if i%2==0 and V>=2: x[:,0]=x[:,1]
    p=float(torch.rand(1,generator=g).item())
    a=O.top_p_filter(x.clone(),p); b=R.top_p_filter(x.clone(),p)
    if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True):
        bad+=1
        if ex<3:
            ex+=1; print('x',x.tolist(),'p',round(p,4)); print(' mine',a.tolist()); print(' ref ',b.tolist())
print('bad',bad)
"
output
bad 0

[stdout]
bad 0

Reproduce top_p within verify2 rng partial

shell
python3 -c "
import torch, reforacle as R, objectives as O
g = torch.Generator().manual_seed(123)
ex=0
for trial in range(400):
    B = int(torch.randint(1,6,(1,),generator=g).item())
    V = int(torch.randint(1,12,(1,),generator=g).item())
    T = int(torch.randint(1,10,(1,),generator=g).item())
    scale = float(torch.rand(1,generator=g).item()*10+0.01)
    logits = torch.randn(B, V, generator=g) * scale
    if trial % 5 == 0 and V >= 2:
        logits[:, 0] = logits[:, 1]
    p=float(torch.rand(1,generator=g).item())
    a=O.top_p_filter(logits.clone(),p); b=R.top_p_filter(logits.clone(),p)
    if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True):
        if ex<3:
            ex+=1; print('trial',trial,'V',V,'p',round(p,4)); print('x',logits.tolist()); print('mine',a.tolist()); print('ref ',b.tolist())
    # advance rng identically to verify2? no , verify2 draws many more randoms. skip
" 2>&1 | head -40
output
(Bash completed with no output)

Capture verify2 top_p failing inputs

shell
python3 -c "
import torch, reforacle as R, objectives as O
import re
src=open('verify2.py').read()
# monkeypatch: capture first top_p fail
orig=O.top_p_filter
caught=[]
def wrap(x,p):
    m=orig(x,p); r=R.top_p_filter(x.clone(),p)
    if not torch.allclose(m,r,atol=1e-5,rtol=1e-4,equal_nan=True) and len(caught)<3:
        caught.append((x.clone(),p,m,r))
    return m
O.top_p_filter=wrap
# reduce to just run generation loop by exec but strip final print
exec(src.replace('run(\"','_run(\"').replace('def run','def _run'))
for x,p,m,r in caught:
    print('p',p); print('x',x.tolist()); print('mine',m.tolist()); print('ref ',r.tolist()); print()
" 2>&1 | head -40
output
total fails: 121
  top_p_filter: 121 e.g. vals
p 1.0
x [[1.7758086919784546, 3.581681251525879, -5.410852909088135, 9.188485145568848, 6.68272066116333, -8.961294174194336, 9.625248908996582, 1.405346155166626, 5.925773620605469, 1.8858153820037842, 7.367067813873291], [-13.575481414794922, 11.884575843811035, 9.38491153717041, -0.5591022968292236, -2.1097443103790283, 8.02276611328125, -15.76061725616455, 7.115726947784424, -3.300771951675415, 11.162632942199707, -10.598807334899902], [-7.092728137969971, -4.810391902923584, 10.17324161529541, 1.045422077178955, -4.885246753692627, -5.265064239501953, 8.953729629516602, -5.946301460266113, 3.420771598815918, -17.17274284362793, -13.560656547546387], [1.1115657091140747, 2.958920478820801, -8.193439483642578, -4.31288480758667, 8.248448371887207, -2.89162278175354, 15.123235702514648, 9.90704345703125, -5.3740081787109375, 14.98278522491455, -1.7375669479370117], [-5.868556022644043, 12.217329025268555, -2.7847061157226562, -1.9306726455688477, 0.5153402090072632, -5.1682515144348145, 4.159594535827637, 10.011043548583984, 1.8727649450302124, -8.167160034179688, 10.277349472045898]]
mine [[1.7758086919784546, 3.581681251525879, -5.410852909088135, 9.188485145568848, 6.68272066116333, -8.961294174194336, 9.625248908996582, 1.405346155166626, 5.925773620605469, 1.8858153820037842, 7.367067813873291], [-13.575481414794922, 11.884575843811035, 9.38491153717041, -0.5591022968292236, -2.1097443103790283, 8.02276611328125, -15.76061725616455, 7.115726947784424, -3.300771951675415, 11.162632942199707, -10.598807334899902], [-7.092728137969971, -4.810391902923584, 10.17324161529541, 1.045422077178955, -4.885246753692627, -5.265064239501953, 8.953729629516602, -5.946301460266113, 3.420771598815918, -17.17274284362793, -13.560656547546387], [1.1115657091140747, 2.958920478820801, -inf, -4.31288480758667, 8.248448371887207, -2.89162278175354, 15.123235702514648, 9.90704345703125, -5.3740081787109375, 14.98278522491455, -1.7375669479370117], [-5.868556022644043, 12.217329025268555, -2.7847061157226562, -1.9306726455688477, 0.5153402090072632, -5.1682515144348145, 4.159594535827637, 10.011043548583984, 1.8727649450302124, -8.167160034179688, 10.277349472045898]]
ref  [[1.7758086919784546, 3.581681251525879, -5.410852909088135, 9.188485145568848, 6.68272066116333, -8.961294174194336, 9.625248908996582, 1.405346155166626, 5.925773620605469, 1.8858153820037842, 7.367067813873291], [-13.575481414794922, 11.884575843811035, 9.38491153717041, -0.5591022968292236, -2.1097443103790283, 8.02276611328125, -15.76061725616455, 7.115726947784424, -3.300771951675415, 11.162632942199707, -10.598807334899902], [-inf, -4.810391902923584, 10.17324161529541, 1.045422077178955, -4.885246753692627, -5.265064239501953, 8.953729629516602, -5.946301460266113, 3.420771598815918, -inf, -inf], [1.1115657091140747, 2.958920478820801, -inf, -inf, 8.248448371887207, -inf, 15.123235702514648, 9.90704345703125, -inf, 14.98278522491455, -inf], [-inf, 12.217329025268555, -2.7847061157226562, -1.9306726455688477, 0.5153402090072632, -inf, 4.159594535827637, 10.011043548583984, 1.8727649450302124, -inf, 10.277349472045898]]

p 1.0
x [[10.21911907196045, -4.102693557739258, -2.8474466800689697, -2.4570770263671875, 5.141626358032227, 13.388644218444824, 0.31344836950302124, -8.709283828735352, -2.512850522994995, 1.867435097694397], [4.369866847991943, 2.3859152793884277, -8.550373077392578, 1.0145007371902466, -7.1879167556762695, -13.052201271057129, 3.3485517501831055, -10.453784942626953, -3.4476804733276367, 5.702380657196045]]
mine [[10.21911907196045, -4.102693557739258, -2.8474466800689697, -2.4570770263671875, 5.141626358032227, 13.388644218444824, 0.31344836950302124, -8.709283828735352, -2.512850522994995, 1.867435097694397], [4.369866847991943, 2.3859152793884277, -8.550373077392578, 1.0145007371902466, -7.1879167556762695, -13.052201271057129, 3.3485517501831055, -10.453784942626953, -3.4476804733276367, 5.702380657196045]]
ref  [[10.21911907196045, -inf, -2.8474466800689697, -2.4570770263671875, 5.141626358032227, 13.388644218444824, 0.31344836950302124, -inf, -2.512850522994995, 1.867435097694397], [4.369866847991943, 2.3859152793884277, -8.550373077392578, 1.0145007371902466, -7.1879167556762695, -inf, 3.3485517501831055, -10.453784942626953, -3.4476804733276367, 5.702380657196045]]

p 1.0
x [[5.413155555725098, -1.6563454866409302, -4.4255499839782715, 5.050666332244873, -6.673373699188232, 1.029617428779602, 4.334243297576904, -2.0456416606903076, -4.013360500335693, -6.2872161865234375, -1.7475675344467163], [-0.1707330346107483, 9.259326934814453, -12.252002716064453, -1.0978566408157349, -7.9540019035339355, 10.626262664794922, -0.8984688520431519, 7.163742542266846, 2.3576760292053223, -2.268575429916382, -1.7430943250656128]]
mine [[5.413155555725098, -1.6563454866409302, -4.4255499839782715, 5.050666332244873, -6.673373699188232, 1.029617428779602, 4.334243297576904, -2.0456416606903076, -4.013360500335693, -6.2872161865234375, -1.7475675344467163], [-0.1707330346107483, 9.259326934814453, -12.252002716064453, -1.0978566408157349, -7.9540019035339355, 10.626262664794922, -0.8984688520431519, 7.163742542266846, 2.3576760292053223, -2.268575429916382, -1.7430943250656128]]
ref  [[5.413155555725098, -1.6563454866409302, -4.4255499839782715, 5.050666332244873, -6.673373699188232, 1.029617428779602, 4.334243297576904, -2.0456416606903076, -4.013360500335693, -6.2872161865234375, -1.7475675344467163], [-0.1707330346107483, 9.259326934814453, -inf, -1.0978566408157349, -inf, 10.626262664794922, -0.8984688520431519, 7.163742542266846, 2.3576760292053223, -2.268575429916382, -1.7430943250656128]]

[stdout]
total fails: 121
  top_p_filter: 121 e.g. vals
p 1.0
x [[1.7758086919784546, 3.581681251525879, -5.410852909088135, 9.188485145568848, 6.68272066116333, -8.961294174194336, 9.625248908996582, 1.405346155166626, 5.925773620605469, 1.8858153820037842, 7.367067813873291], [-13.575481414794922, 11.884575843811035, 9.38491153717041, -0.5591022968292236, -2.1097443103790283, 8.02276611328125, -15.76061725616455, 7.115726947784424, -3.300771951675415, 11.162632942199707, -10.598807334899902], [-7.092728137969971, -4.810391902923584, 10.17324161529541, 1.045422077178955, -4.885246753692627, -5.265064239501953, 8.953729629516602, -5.946301460266113, 3.420771598815918, -17.17274284362793, -13.560656547546387], [1.1115657091140747, 2.958920478820801, -8.193439483642578, -4.31288480758667, 8.248448371887207, -2.89162278175354, 15.123235702514648, 9.90704345703125, -5.3740081787109375, 14.98278522491455, -1.7375669479370117], [-5.868556022644043, 12.217329025268555, -2.7847061157226562, -1.9306726455688477, 0.5153402090072632, -5.1682515144348145, 4.159594535827637, 10.011043548583984, 1.8727649450302124, -8.167160034179688, 10.277349472045898]]
mine [[1.7758086919784546, 3.581681251525879, -5.410852909088135, 9.188485145568848, 6.68272066116333, -8.961294174194336, 9.625248908996582, 1.405346155166626, 5.925773620605469, 1.8858153820037842, 7.367067813873291], [-13.575481414794922, 11.884575843811035, 9.38491153717041, -0.5591022968292236, -2.1097443103790283, 8.02276611328125, -15.76061725616455, 7.115726947784424, -3.300771951675415, 11.162632942199707, -10.598807334899902], [-7.092728137969971, -4.810391902923584, 10.17324161529541, 1.045422077178955, -4.885246753692627, -5.265064239501953, 8.953729629516602, -5.946301460266113, 3.420771598815918, -17.17274284362793, -13.560656547546387], [1.1115657091140747, 2.958920478820801, -inf, -4.31288480758667, 8.248448371887207, -2.89162278175354, 15.123235702514648, 9.90704345703125, -5.3740081787109375, 14.98278522491455, -1.7375669479370117], [-5.868556022644043, 12.217329025268555, -2.7847061157226562, -1.9306726455688477, 0.5153402090072632, -5.1682515144348145, 4.159594535827637, 10.011043548583984, 1.8727649450302124, -8.167160034179688, 10.277349472045898]]
ref  [[1.7758086919784546, 3.581681251525879, -5.410852909088135, 9.188485145568848, 6.68272066116333, -8.961294174194336, 9.625248908996582, 1.405346155166626, 5.925773620605469, 1.8858153820037842, 7.367067813873291], [-13.575481414794922, 11.884575843811035, 9.38491153717041, -0.5591022968292236, -2.1097443103790283, 8.02276611328125, -15.76061725616455, 7.115726947784424, -3.300771951675415, 11.162632942199707, -10.598807334899902], [-inf, -4.810391902923584, 10.17324161529541, 1.045422077178955, -4.885246753692627, -5.265064239501953, 8.953729629516602, -5.946301460266113, 3.420771598815918, -inf, -inf], [1.1115657091140747, 2.958920478820801, -inf, -inf, 8.248448371887207, -inf, 15.123235702514648, 9.90704345703125, -inf, 14.98278522491455, -inf], [-inf, 12.217329025268555, -2.7847061157226562, -1.9306726455688477, 0.5153402090072632, -inf, 4.159594535827637, 10.011043548583984, 1.8727649450302124, -inf, 10.277349472045898]]

p 1.0
x [[10.21911907196045, -4.102693557739258, -2.8474466800689697, -2.4570770263671875, 5.141626358032227, 13.388644218444824, 0.31344836950302124, -8.709283828735352, -2.512850522994995, 1.867435097694397], [4.369866847991943, 2.3859152793884277, -8.550373077392578, 1.0145007371902466, -7.1879167556762695, -13.052201271057129, 3.3485517501831055, -10.453784942626953, -3.4476804733276367, 5.702380657196045]]
mine [[10.21911907196045, -4.102693557739258, -2.8474466800689697, -2.4570770263671875, 5.141626358032227, 13.388644218444824, 0.31344836950302124, -8.709283828735352, -2.512850522994995, 1.867435097694397], [4.369866847991943, 2.3859152793884277, -8.550373077392578, 1.0145007371902466, -7.1879167556762695, -13.052201271057129, 3.3485517501831055, -10.453784942626953, -3.4476804733276367, 5.702380657196045]]
ref  [[10.21911907196045, -inf, -2.8474466800689697, -2.4570770263671875, 5.141626358032227, 13.388644218444824, 0.31344836950302124, -inf, -2.512850522994995, 1.867435097694397], [4.369866847991943, 2.3859152793884277, -8.550373077392578, 1.0145007371902466, -7.1879167556762695, -inf, 3.3485517501831055, -10.453784942626953, -3.4476804733276367, 5.702380657196045]]

p 1.0
x [[5.413155555725098, -1.6563454866409302, -4.4255499839782715, 5.050666332244873, -6.673373699188232, 1.029617428779602, 4.334243297576904, -2.0456416606903076, -4.013360500335693, -6.2872161865234375, -1.7475675344467163], [-0.1707330346107483, 9.259326934814453, -12.252002716064453, -1.0978566408157349, -7.9540019035339355, 10.626262664794922, -0.8984688520431519, 7.163742542266846, 2.3576760292053223, -2.268575429916382, -1.7430943250656128]]
mine [[5.413155555725098, -1.6563454866409302, -4.4255499839782715, 5.050666332244873, -6.673373699188232, 1.029617428779602, 4.334243297576904, -2.0456416606903076, -4.013360500335693, -6.2872161865234375, -1.7475675344467163], [-0.1707330346107483, 9.259326934814453, -12.252002716064453, -1.0978566408157349, -7.9540019035339355, 10.626262664794922, -0.8984688520431519, 7.163742542266846, 2.3576760292053223, -2.268575429916382, -1.7430943250656128]]
ref  [[5.413155555725098, -1.6563454866409302, -4.4255499839782715, 5.050666332244873, -6.673373699188232, 1.029617428779602, 4.334243297576904, -2.0456416606903076, -4.013360500335693, -6.2872161865234375, -1.7475675344467163], [-0.1707330346107483, 9.259326934814453, -inf, -1.0978566408157349, -inf, 10.626262664794922, -0.8984688520431519, 7.163742542266846, 2.3576760292053223, -2.268575429916382, -1.7430943250656128]]

/app/objectives.py

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

Exhaustive top_p test with underflow, ties, p=0/1

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(55)
bad=0;ex=0
for i in range(30000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(1,10,(1,),generator=g).item())
    scale=float(torch.rand(1,generator=g).item()*15+0.01)
    x=torch.randn(B,V,generator=g)*scale
    r=torch.rand(1,generator=g).item()
    if r<0.3 and V>=2: x[:,0]=x[:,1]
    if r<0.15 and V>=3: x[:,2]=x[:,0]
    p=float(torch.rand(1,generator=g).item())
    if i%4==0: p=1.0
    if i%9==0: p=0.0
    a=O.top_p_filter(x.clone(),p); b=R.top_p_filter(x.clone(),p)
    if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True):
        bad+=1
        if ex<3: ex+=1;print('p',p,'x',x.tolist());print('mine',a.tolist());print('ref',b.tolist())
print('top_p bad',bad)
"
output
p 1.0 x [[-12.809303283691406, 11.811380386352539, -3.440556764602661, -6.228688716888428, 1.8893862962722778], [-0.6871057748794556, 3.42156982421875, -11.575715065002441, 8.580294609069824, 12.03497314453125]]
mine [[-12.809303283691406, 11.811380386352539, -3.440556764602661, -6.228688716888428, 1.8893862962722778], [-0.6871057748794556, 3.42156982421875, -11.575715065002441, 8.580294609069824, 12.03497314453125]]
ref [[-inf, 11.811380386352539, -3.440556764602661, -inf, 1.8893862962722778], [-0.6871057748794556, 3.42156982421875, -11.575715065002441, 8.580294609069824, 12.03497314453125]]
p 1.0 x [[12.75774097442627, 22.052257537841797, 12.83862590789795, 30.20471954345703, 16.393339157104492, 6.460799217224121], [25.852802276611328, 2.817809581756592, -2.2025952339172363, 4.290494918823242, -10.848952293395996, 3.1950926780700684]]
mine [[12.75774097442627, 22.052257537841797, 12.83862590789795, 30.20471954345703, 16.393339157104492, 6.460799217224121], [25.852802276611328, 2.817809581756592, -2.2025952339172363, 4.290494918823242, -10.848952293395996, 3.1950926780700684]]
ref [[-inf, 22.052257537841797, -inf, 30.20471954345703, 16.393339157104492, -inf], [25.852802276611328, -inf, -inf, -inf, -inf, -inf]]
p 1.0 x [[-3.730949640274048, 33.86341857910156, -7.1269941329956055, -9.73001766204834], [9.902432441711426, -3.8288395404815674, 3.016354560852051, 4.782057285308838]]
mine [[-3.730949640274048, 33.86341857910156, -7.1269941329956055, -9.73001766204834], [9.902432441711426, -3.8288395404815674, 3.016354560852051, 4.782057285308838]]
ref [[-inf, 33.86341857910156, -inf, -inf], [9.902432441711426, -3.8288395404815674, 3.016354560852051, 4.782057285308838]]
top_p bad 4681

[stdout]
p 1.0 x [[-12.809303283691406, 11.811380386352539, -3.440556764602661, -6.228688716888428, 1.8893862962722778], [-0.6871057748794556, 3.42156982421875, -11.575715065002441, 8.580294609069824, 12.03497314453125]]
mine [[-12.809303283691406, 11.811380386352539, -3.440556764602661, -6.228688716888428, 1.8893862962722778], [-0.6871057748794556, 3.42156982421875, -11.575715065002441, 8.580294609069824, 12.03497314453125]]
ref [[-inf, 11.811380386352539, -3.440556764602661, -inf, 1.8893862962722778], [-0.6871057748794556, 3.42156982421875, -11.575715065002441, 8.580294609069824, 12.03497314453125]]
p 1.0 x [[12.75774097442627, 22.052257537841797, 12.83862590789795, 30.20471954345703, 16.393339157104492, 6.460799217224121], [25.852802276611328, 2.817809581756592, -2.2025952339172363, 4.290494918823242, -10.848952293395996, 3.1950926780700684]]
mine [[12.75774097442627, 22.052257537841797, 12.83862590789795, 30.20471954345703, 16.393339157104492, 6.460799217224121], [25.852802276611328, 2.817809581756592, -2.2025952339172363, 4.290494918823242, -10.848952293395996, 3.1950926780700684]]
ref [[-inf, 22.052257537841797, -inf, 30.20471954345703, 16.393339157104492, -inf], [25.852802276611328, -inf, -inf, -inf, -inf, -inf]]
p 1.0 x [[-3.730949640274048, 33.86341857910156, -7.1269941329956055, -9.73001766204834], [9.902432441711426, -3.8288395404815674, 3.016354560852051, 4.782057285308838]]
mine [[-3.730949640274048, 33.86341857910156, -7.1269941329956055, -9.73001766204834], [9.902432441711426, -3.8288395404815674, 3.016354560852051, 4.782057285308838]]
ref [[-inf, 33.86341857910156, -inf, -inf], [9.902432441711426, -3.8288395404815674, 3.016354560852051, 4.782057285308838]]
top_p bad 4681

Extract exact cumsum and removed mask

shell
python3 -c "
import torch, reforacle as R
torch.set_printoptions(precision=12, sci_mode=True)
def analyze(x,p):
    x=x.float()
    sl,si=torch.sort(x,descending=False,dim=-1,stable=True)
    cum=sl.softmax(-1).cumsum(-1)
    ref=R.top_p_filter(x.clone(),p)
    removed=torch.isinf(ref)&(ref<0)
    # map removed back to sorted order
    rem_sorted=removed.gather(-1,si)
    for r in range(x.shape[0]):
        print('row',r,'p',p)
        print(' cum   ',[f'{v:.3e}' for v in cum[r].tolist()])
        print(' removed(sorted)',rem_sorted[r].tolist())
analyze(torch.tensor([[-12.809303,11.811380,-3.440557,-6.228689,1.889386]]),1.0)
analyze(torch.tensor([[-0.6871,3.4216,-11.5757,8.5803,12.0350]]),1.0)
analyze(torch.tensor([[25.8528,2.8178,-2.2026,4.2905,-10.849,3.1951]]),1.0)
"
output
row 0 p 1.0
 cum    ['2.029e-11', '1.465e-08', '2.524e-07', '4.933e-05', '1.000e+00']
 removed(sorted) [True, True, False, False, False]
row 0 p 1.0
 cum    ['5.400e-11', '2.893e-06', '1.790e-04', '3.080e-02', '1.000e+00']
 removed(sorted) [True, False, False, False, False]
row 0 p 1.0
 cum    ['1.150e-16', '6.543e-13', '9.974e-11', '2.442e-10', '6.764e-10', '1.000e+00']
 removed(sorted) [True, True, True, True, True, False]

[stdout]
row 0 p 1.0
 cum    ['2.029e-11', '1.465e-08', '2.524e-07', '4.933e-05', '1.000e+00']
 removed(sorted) [True, True, False, False, False]
row 0 p 1.0
 cum    ['5.400e-11', '2.893e-06', '1.790e-04', '3.080e-02', '1.000e+00']
 removed(sorted) [True, False, False, False, False]
row 0 p 1.0
 cum    ['1.150e-16', '6.543e-13', '9.974e-11', '2.442e-10', '6.764e-10', '1.000e+00']
 removed(sorted) [True, True, True, True, True, False]

/app/objectives.py

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

Exhaustive top_p test

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(2024)
bad=0;ex=0
for i in range(60000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(1,10,(1,),generator=g).item())
    scale=float(torch.rand(1,generator=g).item()*15+0.01)
    x=torch.randn(B,V,generator=g)*scale
    r=torch.rand(1,generator=g).item()
    if r<0.3 and V>=2: x[:,0]=x[:,1]
    if r<0.15 and V>=3: x[:,2]=x[:,0]
    p=float(torch.rand(1,generator=g).item())
    m=i%5
    if m==0: p=1.0
    elif m==1: p=0.0
    elif m==2: x=torch.randint(0,3,(B,V),generator=g).float(); p=float(torch.rand(1,generator=g).item())
    a=O.top_p_filter(x.clone(),p); b=R.top_p_filter(x.clone(),p)
    if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True):
        bad+=1
        if ex<3: ex+=1;print('p',p,'x',x.tolist());print('mine',a.tolist());print('ref',b.tolist())
print('top_p bad',bad)
"
output
p 1.0 x [[-3.4699440002441406, -3.4699440002441406, 20.04443359375, -11.132471084594727, 8.823887825012207, 22.87409210205078, -6.133819103240967, -9.832380294799805], [-7.481536388397217, -7.481536388397217, -15.246474266052246, 1.7407971620559692, -5.562899112701416, 18.737783432006836, -10.415597915649414, -5.630403995513916], [-5.827956676483154, -5.827956676483154, -8.039402961730957, 14.168121337890625, 12.172632217407227, 15.070635795593262, -0.9326332807540894, 4.661966800689697]]
mine [[-3.4699440002441406, -3.4699440002441406, 20.04443359375, -11.132471084594727, 8.823887825012207, 22.87409210205078, -6.133819103240967, -9.832380294799805], [-inf, -inf, -inf, -inf, -inf, 18.737783432006836, -inf, -inf], [-5.827956676483154, -5.827956676483154, -8.039402961730957, 14.168121337890625, 12.172632217407227, 15.070635795593262, -0.9326332807540894, 4.661966800689697]]
ref [[-3.4699440002441406, -3.4699440002441406, 20.04443359375, -11.132471084594727, 8.823887825012207, 22.87409210205078, -6.133819103240967, -9.832380294799805], [-inf, -inf, -inf, 1.7407971620559692, -inf, 18.737783432006836, -inf, -inf], [-5.827956676483154, -5.827956676483154, -8.039402961730957, 14.168121337890625, 12.172632217407227, 15.070635795593262, -0.9326332807540894, 4.661966800689697]]
p 1.0 x [[-6.110450267791748, 7.84336519241333, 0.27275004982948303, 1.65316641330719, -22.277545928955078, -2.1182637214660645, -1.6175609827041626], [3.504721164703369, -13.221860885620117, 1.855076789855957, -6.913834095001221, -8.8451566696167, -2.79672908782959, 20.72626304626465], [6.73052453994751, -2.1677324771881104, -7.739095211029053, -3.981194019317627, 6.361721515655518, -1.4869351387023926, 0.12118364870548248]]
mine [[-6.110450267791748, 7.84336519241333, 0.27275004982948303, 1.65316641330719, -inf, -2.1182637214660645, -1.6175609827041626], [-inf, -inf, -inf, -inf, -inf, -inf, 20.72626304626465], [6.73052453994751, -2.1677324771881104, -7.739095211029053, -3.981194019317627, 6.361721515655518, -1.4869351387023926, 0.12118364870548248]]
ref [[-6.110450267791748, 7.84336519241333, 0.27275004982948303, 1.65316641330719, -inf, -2.1182637214660645, -1.6175609827041626], [3.504721164703369, -inf, -inf, -inf, -inf, -inf, 20.72626304626465], [6.73052453994751, -2.1677324771881104, -7.739095211029053, -3.981194019317627, 6.361721515655518, -1.4869351387023926, 0.12118364870548248]]
p 1.0 x [[-0.9820225238800049, -15.955976486206055, -10.931145668029785, -0.35602664947509766, -3.17645001411438, 3.0744223594665527, -6.200700759887695, -16.941640853881836]]
mine [[-0.9820225238800049, -15.955976486206055, -10.931145668029785, -0.35602664947509766, -3.17645001411438, 3.0744223594665527, -6.200700759887695, -inf]]
ref [[-0.9820225238800049, -inf, -10.931145668029785, -0.35602664947509766, -3.17645001411438, 3.0744223594665527, -6.200700759887695, -inf]]
top_p bad 805

[stdout]
p 1.0 x [[-3.4699440002441406, -3.4699440002441406, 20.04443359375, -11.132471084594727, 8.823887825012207, 22.87409210205078, -6.133819103240967, -9.832380294799805], [-7.481536388397217, -7.481536388397217, -15.246474266052246, 1.7407971620559692, -5.562899112701416, 18.737783432006836, -10.415597915649414, -5.630403995513916], [-5.827956676483154, -5.827956676483154, -8.039402961730957, 14.168121337890625, 12.172632217407227, 15.070635795593262, -0.9326332807540894, 4.661966800689697]]
mine [[-3.4699440002441406, -3.4699440002441406, 20.04443359375, -11.132471084594727, 8.823887825012207, 22.87409210205078, -6.133819103240967, -9.832380294799805], [-inf, -inf, -inf, -inf, -inf, 18.737783432006836, -inf, -inf], [-5.827956676483154, -5.827956676483154, -8.039402961730957, 14.168121337890625, 12.172632217407227, 15.070635795593262, -0.9326332807540894, 4.661966800689697]]
ref [[-3.4699440002441406, -3.4699440002441406, 20.04443359375, -11.132471084594727, 8.823887825012207, 22.87409210205078, -6.133819103240967, -9.832380294799805], [-inf, -inf, -inf, 1.7407971620559692, -inf, 18.737783432006836, -inf, -inf], [-5.827956676483154, -5.827956676483154, -8.039402961730957, 14.168121337890625, 12.172632217407227, 15.070635795593262, -0.9326332807540894, 4.661966800689697]]
p 1.0 x [[-6.110450267791748, 7.84336519241333, 0.27275004982948303, 1.65316641330719, -22.277545928955078, -2.1182637214660645, -1.6175609827041626], [3.504721164703369, -13.221860885620117, 1.855076789855957, -6.913834095001221, -8.8451566696167, -2.79672908782959, 20.72626304626465], [6.73052453994751, -2.1677324771881104, -7.739095211029053, -3.981194019317627, 6.361721515655518, -1.4869351387023926, 0.12118364870548248]]
mine [[-6.110450267791748, 7.84336519241333, 0.27275004982948303, 1.65316641330719, -inf, -2.1182637214660645, -1.6175609827041626], [-inf, -inf, -inf, -inf, -inf, -inf, 20.72626304626465], [6.73052453994751, -2.1677324771881104, -7.739095211029053, -3.981194019317627, 6.361721515655518, -1.4869351387023926, 0.12118364870548248]]
ref [[-6.110450267791748, 7.84336519241333, 0.27275004982948303, 1.65316641330719, -inf, -2.1182637214660645, -1.6175609827041626], [3.504721164703369, -inf, -inf, -inf, -inf, -inf, 20.72626304626465], [6.73052453994751, -2.1677324771881104, -7.739095211029053, -3.981194019317627, 6.361721515655518, -1.4869351387023926, 0.12118364870548248]]
p 1.0 x [[-0.9820225238800049, -15.955976486206055, -10.931145668029785, -0.35602664947509766, -3.17645001411438, 3.0744223594665527, -6.200700759887695, -16.941640853881836]]
mine [[-0.9820225238800049, -15.955976486206055, -10.931145668029785, -0.35602664947509766, -3.17645001411438, 3.0744223594665527, -6.200700759887695, -inf]]
ref [[-0.9820225238800049, -inf, -10.931145668029785, -0.35602664947509766, -3.17645001411438, 3.0744223594665527, -6.200700759887695, -inf]]
top_p bad 805

Examine exact cumsum accumulation

shell
python3 -c "
import torch, reforacle as R
torch.set_printoptions(precision=10, sci_mode=False)
x=torch.tensor([[-7.481536388397217,-7.481536388397217,-15.246474266052246,1.7407971620559692,-5.562899112701416,18.737783432006836,-10.415597915649414,-5.630403995513916]])
sl,si=torch.sort(x,descending=True,dim=-1,stable=True)
print('sorted idx', si.tolist())
cum=sl.softmax(-1).cumsum(-1)
print('desc cum', cum.tolist())
# ascending
sla,sia=torch.sort(x,descending=False,dim=-1,stable=True)
cuma=sla.softmax(-1).cumsum(-1)
print('asc idx', sia.tolist())
print('asc cum', cuma.tolist())
print('asc 1-cum', (1-cuma).tolist())
ref=R.top_p_filter(x.clone(),1.0)
print('ref removed', (torch.isinf(ref)&(ref<0)).tolist())
"
output
sorted idx [[5, 3, 4, 7, 0, 1, 6, 2]]
desc cum [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]]
asc idx [[2, 6, 0, 1, 7, 4, 3, 5]]
asc cum [[1.741104492805829e-15, 2.1993739024016995e-13, 4.322864982642116e-12, 8.425793049382513e-12, 3.454926275625603e-11, 6.249707695804574e-11, 4.1586822874251084e-08, 1.0]]
asc 1-cum [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.9999999403953552, 0.0]]
ref removed [[True, True, True, False, True, False, True, True]]

[stdout]
sorted idx [[5, 3, 4, 7, 0, 1, 6, 2]]
desc cum [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]]
asc idx [[2, 6, 0, 1, 7, 4, 3, 5]]
asc cum [[1.741104492805829e-15, 2.1993739024016995e-13, 4.322864982642116e-12, 8.425793049382513e-12, 3.454926275625603e-11, 6.249707695804574e-11, 4.1586822874251084e-08, 1.0]]
asc 1-cum [[1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.9999999403953552, 0.0]]
ref removed [[True, True, True, False, True, False, True, True]]

/app/objectives.py

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

Exhaustive top_p test v2

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(2024)
bad=0;ex=0
for i in range(120000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(1,10,(1,),generator=g).item())
    scale=float(torch.rand(1,generator=g).item()*15+0.01)
    x=torch.randn(B,V,generator=g)*scale
    r=torch.rand(1,generator=g).item()
    if r<0.3 and V>=2: x[:,0]=x[:,1]
    if r<0.15 and V>=3: x[:,2]=x[:,0]
    p=float(torch.rand(1,generator=g).item())
    m=i%6
    if m==0: p=1.0
    elif m==1: p=0.0
    elif m==2: x=torch.randint(0,3,(B,V),generator=g).float()
    elif m==3: p=float(torch.rand(1,generator=g).item())*0.05
    a=O.top_p_filter(x.clone(),p); b=R.top_p_filter(x.clone(),p)
    if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True):
        bad+=1
        if ex<4: ex+=1;print('p',p,'x',x.tolist());print('mine',a.tolist());print('ref',b.tolist())
print('top_p bad',bad)
"
output
p 0.6867693066596985 x [[0.0, 0.0, 2.0, 2.0, 0.0, 2.0, 2.0, 1.0], [2.0, 0.0, 0.0, 2.0, 1.0, 0.0, 0.0, 1.0]]
mine [[-inf, -inf, 2.0, 2.0, -inf, 2.0, 2.0, -inf], [2.0, -inf, -inf, 2.0, -inf, -inf, -inf, 1.0]]
ref [[-inf, -inf, 2.0, 2.0, -inf, 2.0, 2.0, -inf], [2.0, -inf, -inf, 2.0, 1.0, -inf, -inf, -inf]]
p 1.0 x [[-8.49079418182373, -8.49079418182373, -8.49079418182373, 11.949380874633789, 2.2269976139068604, 4.494710922241211, 6.766523838043213], [-0.32589223980903625, -0.32589223980903625, -0.32589223980903625, -31.04218292236328, -6.372030258178711, -4.824691295623779, 12.450148582458496], [24.950637817382812, 24.950637817382812, 24.950637817382812, 11.400482177734375, 12.333769798278809, 1.2486921548843384, 10.207602500915527]]
mine [[-inf, -inf, -inf, 11.949380874633789, 2.2269976139068604, 4.494710922241211, 6.766523838043213], [-0.32589223980903625, -0.32589223980903625, -0.32589223980903625, -inf, -inf, -4.824691295623779, 12.450148582458496], [24.950637817382812, 24.950637817382812, 24.950637817382812, 11.400482177734375, 12.333769798278809, -inf, 10.207602500915527]]
ref [[-8.49079418182373, -8.49079418182373, -8.49079418182373, 11.949380874633789, 2.2269976139068604, 4.494710922241211, 6.766523838043213], [-0.32589223980903625, -0.32589223980903625, -0.32589223980903625, -inf, -inf, -4.824691295623779, 12.450148582458496], [24.950637817382812, 24.950637817382812, 24.950637817382812, 11.400482177734375, 12.333769798278809, 1.2486921548843384, 10.207602500915527]]
p 0.5742855072021484 x [[0.0, 0.0, 0.0]]
mine [[-inf, 0.0, 0.0]]
ref [[0.0, 0.0, -inf]]
p 0.8699129819869995 x [[1.0, 0.0, 2.0, 0.0, 2.0, 0.0], [0.0, 2.0, 0.0, 2.0, 1.0, 0.0]]
mine [[1.0, -inf, 2.0, -inf, 2.0, 0.0], [-inf, 2.0, -inf, 2.0, 1.0, 0.0]]
ref [[1.0, 0.0, 2.0, -inf, 2.0, -inf], [0.0, 2.0, -inf, 2.0, 1.0, -inf]]
top_p bad 23640

[stdout]
p 0.6867693066596985 x [[0.0, 0.0, 2.0, 2.0, 0.0, 2.0, 2.0, 1.0], [2.0, 0.0, 0.0, 2.0, 1.0, 0.0, 0.0, 1.0]]
mine [[-inf, -inf, 2.0, 2.0, -inf, 2.0, 2.0, -inf], [2.0, -inf, -inf, 2.0, -inf, -inf, -inf, 1.0]]
ref [[-inf, -inf, 2.0, 2.0, -inf, 2.0, 2.0, -inf], [2.0, -inf, -inf, 2.0, 1.0, -inf, -inf, -inf]]
p 1.0 x [[-8.49079418182373, -8.49079418182373, -8.49079418182373, 11.949380874633789, 2.2269976139068604, 4.494710922241211, 6.766523838043213], [-0.32589223980903625, -0.32589223980903625, -0.32589223980903625, -31.04218292236328, -6.372030258178711, -4.824691295623779, 12.450148582458496], [24.950637817382812, 24.950637817382812, 24.950637817382812, 11.400482177734375, 12.333769798278809, 1.2486921548843384, 10.207602500915527]]
mine [[-inf, -inf, -inf, 11.949380874633789, 2.2269976139068604, 4.494710922241211, 6.766523838043213], [-0.32589223980903625, -0.32589223980903625, -0.32589223980903625, -inf, -inf, -4.824691295623779, 12.450148582458496], [24.950637817382812, 24.950637817382812, 24.950637817382812, 11.400482177734375, 12.333769798278809, -inf, 10.207602500915527]]
ref [[-8.49079418182373, -8.49079418182373, -8.49079418182373, 11.949380874633789, 2.2269976139068604, 4.494710922241211, 6.766523838043213], [-0.32589223980903625, -0.32589223980903625, -0.32589223980903625, -inf, -inf, -4.824691295623779, 12.450148582458496], [24.950637817382812, 24.950637817382812, 24.950637817382812, 11.400482177734375, 12.333769798278809, 1.2486921548843384, 10.207602500915527]]
p 0.5742855072021484 x [[0.0, 0.0, 0.0]]
mine [[-inf, 0.0, 0.0]]
ref [[0.0, 0.0, -inf]]
p 0.8699129819869995 x [[1.0, 0.0, 2.0, 0.0, 2.0, 0.0], [0.0, 2.0, 0.0, 2.0, 1.0, 0.0]]
mine [[1.0, -inf, 2.0, -inf, 2.0, 0.0], [-inf, 2.0, -inf, 2.0, 1.0, 0.0]]
ref [[1.0, 0.0, 2.0, -inf, 2.0, -inf], [0.0, 2.0, -inf, 2.0, 1.0, -inf]]
top_p bad 23640

Extract exact values for both p=1.0 cases

shell
python3 -c "
import torch, reforacle as R
torch.set_printoptions(precision=10, sci_mode=True)
def dump(x,tag):
    x=torch.tensor([x])
    sla,sia=torch.sort(x,descending=False,dim=-1,stable=True)
    cuma=sla.softmax(-1).cumsum(-1)
    sld,sid=torch.sort(x,descending=True,dim=-1,stable=True)
    cumd=sld.softmax(-1).cumsum(-1)
    ref=R.top_p_filter(x.clone(),1.0)
    rem=(torch.isinf(ref)&(ref<0))
    print(tag)
    print(' asc idx',sia[0].tolist())
    print(' asc cum',[f'{v:.4e}' for v in cuma[0].tolist()])
    print(' asc 1-cum',[f'{v:.6f}' for v in (1-cuma[0]).tolist()])
    print(' desc idx',sid[0].tolist())
    print(' desc cum',[f'{v:.10f}' for v in cumd[0].tolist()])
    print(' ref removed(orig)',rem[0].tolist())
    print(' ref removed(desc order)',rem.gather(-1,sid)[0].tolist())
dump([-8.49079418182373,-8.49079418182373,-8.49079418182373,11.949380874633789,2.2269976139068604,4.494710922241211,6.766523838043213],'A all kept?')
dump([-7.481536388397217,-7.481536388397217,-15.246474266052246,1.7407971620559692,-5.562899112701416,18.737783432006836,-10.415597915649414,-5.630403995513916],'B 2 kept')
"
output
A all kept?
 asc idx [0, 1, 2, 4, 5, 6, 3]
 asc cum ['1.3190e-09', '2.6380e-09', '3.9569e-09', '5.9559e-05', '6.3470e-04', '6.2118e-03', '1.0000e+00']
 asc 1-cum ['1.000000', '1.000000', '1.000000', '0.999940', '0.999365', '0.993788', '0.000000']
 desc idx [3, 6, 5, 4, 0, 1, 2]
 desc cum ['0.9937881827', '0.9993652701', '0.9999403954', '0.9999999404', '0.9999999404', '0.9999999404', '0.9999999404']
 ref removed(orig) [False, False, False, False, False, False, False]
 ref removed(desc order) [False, False, False, False, False, False, False]
B 2 kept
 asc idx [2, 6, 0, 1, 7, 4, 3, 5]
 asc cum ['1.7411e-15', '2.1994e-13', '4.3229e-12', '8.4258e-12', '3.4549e-11', '6.2497e-11', '4.1587e-08', '1.0000e+00']
 asc 1-cum ['1.000000', '1.000000', '1.000000', '1.000000', '1.000000', '1.000000', '1.000000', '0.000000']
 desc idx [5, 3, 4, 7, 0, 1, 6, 2]
 desc cum ['1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000']
 ref removed(orig) [True, True, True, False, True, False, True, True]
 ref removed(desc order) [False, False, True, True, True, True, True, True]

[stdout]
A all kept?
 asc idx [0, 1, 2, 4, 5, 6, 3]
 asc cum ['1.3190e-09', '2.6380e-09', '3.9569e-09', '5.9559e-05', '6.3470e-04', '6.2118e-03', '1.0000e+00']
 asc 1-cum ['1.000000', '1.000000', '1.000000', '0.999940', '0.999365', '0.993788', '0.000000']
 desc idx [3, 6, 5, 4, 0, 1, 2]
 desc cum ['0.9937881827', '0.9993652701', '0.9999403954', '0.9999999404', '0.9999999404', '0.9999999404', '0.9999999404']
 ref removed(orig) [False, False, False, False, False, False, False]
 ref removed(desc order) [False, False, False, False, False, False, False]
B 2 kept
 asc idx [2, 6, 0, 1, 7, 4, 3, 5]
 asc cum ['1.7411e-15', '2.1994e-13', '4.3229e-12', '8.4258e-12', '3.4549e-11', '6.2497e-11', '4.1587e-08', '1.0000e+00']
 asc 1-cum ['1.000000', '1.000000', '1.000000', '1.000000', '1.000000', '1.000000', '1.000000', '0.000000']
 desc idx [5, 3, 4, 7, 0, 1, 6, 2]
 desc cum ['1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000', '1.0000000000']
 ref removed(orig) [True, True, True, False, True, False, True, True]
 ref removed(desc order) [False, False, True, True, True, True, True, True]

Test float64 and comparison variants

shell
python3 -c "
import torch, reforacle as R

def cand(x, p, dtype, cmp, order):
    logits=x.clone()
    if order=='desc':
        sl,si=torch.sort(logits,descending=True,dim=-1,stable=True)
        cum=sl.to(dtype).softmax(-1).cumsum(-1)
        if cmp=='gt': rem=cum>p
        else: rem=cum>=p
        rem[...,1:]=rem[...,:-1].clone(); rem[...,0]=False
    else:
        sl,si=torch.sort(logits,descending=False,dim=-1,stable=True)
        cum=sl.to(dtype).softmax(-1).cumsum(-1)
        rem=(1.0-cum)>=p if cmp=='ge' else (cum<=(1-p))
        rem[...,-1:]=False
    rem=rem.scatter(-1,si,rem)
    return logits.masked_fill(rem, float('-inf'))

g=torch.Generator().manual_seed(77)
configs=[('f64','desc','gt'),('f64','desc','ge'),('f32','desc','gt'),('f64','asc','ge'),('f64','asc','le')]
counts={c:0 for c in configs}
for i in range(40000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(1,10,(1,),generator=g).item())
    scale=float(torch.rand(1,generator=g).item()*15+0.01)
    x=torch.randn(B,V,generator=g)*scale
    r=torch.rand(1,generator=g).item()
    if r<0.3 and V>=2: x[:,0]=x[:,1]
    p=float(torch.rand(1,generator=g).item())
    if i%3==0:p=1.0
    if i%7==0:p=0.0
    ref=R.top_p_filter(x.clone(),p)
    for dt,order,cmp in configs:
        dtype=torch.float64 if dt=='f64' else torch.float32
        a=cand(x,p,dtype,cmp,order)
        if not torch.allclose(a.float(),ref,atol=1e-5,rtol=1e-4,equal_nan=True): counts[(dt,order,cmp)]+=1
for c in configs: print(c, counts[c])
"
output
('f64', 'desc', 'gt') 4391
('f64', 'desc', 'ge') 4407
('f32', 'desc', 'gt') 4140
('f64', 'asc', 'ge') 6839
('f64', 'asc', 'le') 6786

[stdout]
('f64', 'desc', 'gt') 4391
('f64', 'desc', 'ge') 4407
('f32', 'desc', 'gt') 4140
('f64', 'asc', 'ge') 6839
('f64', 'asc', 'le') 6786

Isolate precision without ties

shell
python3 -c "
import torch, reforacle as R

def cand(x, p, dtype, cmp, order):
    logits=x.clone()
    if order=='desc':
        sl,si=torch.sort(logits,descending=True,dim=-1,stable=True)
        cum=sl.to(dtype).softmax(-1).cumsum(-1)
        rem = cum>p if cmp=='gt' else cum>=p
        rem[...,1:]=rem[...,:-1].clone(); rem[...,0]=False
    else:
        sl,si=torch.sort(logits,descending=False,dim=-1,stable=True)
        cum=sl.to(dtype).softmax(-1).cumsum(-1)
        rem=(1.0-cum)>=p if cmp=='ge' else (cum<=(1-p))
        rem[...,-1:]=False
    rem=rem.scatter(-1,si,rem)
    return logits.masked_fill(rem, float('-inf'))

g=torch.Generator().manual_seed(77)
configs=[('f64','desc','gt'),('f64','desc','ge'),('f32','desc','gt'),('f32','desc','ge'),('f64','asc','ge'),('f32','asc','ge'),('f64','asc','le'),('f32','asc','le')]
counts={c:0 for c in configs}
for i in range(40000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(1,10,(1,),generator=g).item())
    scale=float(torch.rand(1,generator=g).item()*15+0.01)
    x=torch.randn(B,V,generator=g)*scale   # NO ties
    p=float(torch.rand(1,generator=g).item())
    if i%3==0:p=1.0
    if i%7==0:p=0.0
    ref=R.top_p_filter(x.clone(),p)
    for dt,order,cmp in configs:
        dtype=torch.float64 if dt=='f64' else torch.float32
        a=cand(x,p,dtype,cmp,order)
        if not torch.allclose(a.float(),ref,atol=1e-5,rtol=1e-4,equal_nan=True): counts[(dt,order,cmp)]+=1
for c in configs: print(c, counts[c])
"
output
('f64', 'desc', 'gt') 4468
('f64', 'desc', 'ge') 4457
('f32', 'desc', 'gt') 4226
('f32', 'desc', 'ge') 849
('f64', 'asc', 'ge') 4493
('f32', 'asc', 'ge') 2133
('f64', 'asc', 'le') 4451
('f32', 'asc', 'le') 4451

[stdout]
('f64', 'desc', 'gt') 4468
('f64', 'desc', 'ge') 4457
('f32', 'desc', 'gt') 4226
('f32', 'desc', 'ge') 849
('f64', 'asc', 'ge') 4493
('f32', 'asc', 'ge') 2133
('f64', 'asc', 'le') 4451
('f32', 'asc', 'le') 4451

Check p=0 and current impl by regime

shell
python3 -c "
import torch, reforacle as R
# check p=0 behavior
x=torch.tensor([[2.0,1.0,0.5,-1.0]])
print('p=0 ref', R.top_p_filter(x.clone(),0.0).tolist())
print('p=0.001 ref', R.top_p_filter(x.clone(),0.001).tolist())
print('p=0.5 ref', R.top_p_filter(x.clone(),0.5).tolist())
print('p=1 ref', R.top_p_filter(x.clone(),1.0).tolist())
# small scale, is my asc (1-cum)>=p ok for p in (0,1)?
import objectives as O
g=torch.Generator().manual_seed(1)
for tag,pfix in [('p_rand',None),('p1',1.0),('p0',0.0)]:
    for sc in [1.0,5.0]:
        bad=0
        for i in range(20000):
            B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(2,10,(1,),generator=g).item())
            x=torch.randn(B,V,generator=g)*sc
            p=pfix if pfix is not None else float(torch.rand(1,generator=g).item())
            a=O.top_p_filter(x.clone(),p); b=R.top_p_filter(x.clone(),p)
            if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True): bad+=1
        print(tag,'scale',sc,'bad',bad)
"
output
p=0 ref [[2.0, -inf, -inf, -inf]]
p=0.001 ref [[2.0, -inf, -inf, -inf]]
p=0.5 ref [[2.0, -inf, -inf, -inf]]
p=1 ref [[2.0, 1.0, 0.5, -1.0]]
p_rand scale 1.0 bad 0
p_rand scale 5.0 bad 0
p1 scale 1.0 bad 0
p1 scale 5.0 bad 2197
p0 scale 1.0 bad 0
p0 scale 5.0 bad 0

[stdout]
p=0 ref [[2.0, -inf, -inf, -inf]]
p=0.001 ref [[2.0, -inf, -inf, -inf]]
p=0.5 ref [[2.0, -inf, -inf, -inf]]
p=1 ref [[2.0, 1.0, 0.5, -1.0]]
p_rand scale 1.0 bad 0
p_rand scale 5.0 bad 0
p1 scale 1.0 bad 0
p1 scale 5.0 bad 2197
p0 scale 1.0 bad 0
p0 scale 5.0 bad 0

Brute-force p=1.0 rule candidates

shell
python3 -c "
import torch, reforacle as R

def make(order, dt, rule, shift):
    def f(x,p):
        logits=x.clone()
        desc = order=='desc'
        sl,si=torch.sort(logits,descending=desc,dim=-1,stable=True)
        prob=sl.to(dt).softmax(-1)
        cum=prob.cumsum(-1)
        if rule=='cum_gt': rem=cum>p
        elif rule=='cum_ge': rem=cum>=p
        elif rule=='excl_gt': rem=(cum-prob)>p     # exclusive prefix
        elif rule=='excl_ge': rem=(cum-prob)>=p
        elif rule=='1mcum_ge': rem=(1.0-cum)>=p
        elif rule=='1mcum_gt': rem=(1.0-cum)>p
        if shift=='right': rem[...,1:]=rem[...,:-1].clone(); rem[...,0]=False
        elif shift=='keeplast': rem[...,-1:]=False
        elif shift=='keepfirst': rem[...,0]=False
        rem=rem.scatter(-1,si,rem)
        return logits.masked_fill(rem, float('-inf'))
    return f

cands={
 'desc_f32_cumge_right': make('desc',torch.float32,'cum_ge','right'),
 'desc_f64_cumge_right': make('desc',torch.float64,'cum_ge','right'),
 'desc_f64_cumgt_right': make('desc',torch.float64,'cum_gt','right'),
 'desc_f64_exclge_none': make('desc',torch.float64,'excl_ge','none'),
 'desc_f32_exclge_none': make('desc',torch.float32,'excl_ge','none'),
 'desc_f64_exclgt_none': make('desc',torch.float64,'excl_gt','none'),
 'asc_f32_1mcumge_keeplast': make('asc',torch.float32,'1mcum_ge','keeplast'),
 'asc_f64_1mcumge_keeplast': make('asc',torch.float64,'1mcum_ge','keeplast'),
 'asc_f64_1mcumgt_keeplast': make('asc',torch.float64,'1mcum_gt','keeplast'),
}
g=torch.Generator().manual_seed(5)
res={k:0 for k in cands}
N=0
for i in range(60000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(2,10,(1,),generator=g).item())
    x=torch.randn(B,V,generator=g)*float(torch.rand(1,generator=g).item()*8+2)
    p=1.0
    ref=R.top_p_filter(x.clone(),p); N+=1
    for k,f in cands.items():
        a=f(x,p)
        if not torch.allclose(a.float(),ref,atol=1e-5,rtol=1e-4,equal_nan=True): res[k]+=1
print('N',N)
for k in cands: print(k,res[k])
"
output
N 60000
desc_f32_cumge_right 4579
desc_f64_cumge_right 20859
desc_f64_cumgt_right 20820
desc_f64_exclge_none 20856
desc_f32_exclge_none 0
desc_f64_exclgt_none 20820
asc_f32_1mcumge_keeplast 10940
asc_f64_1mcumge_keeplast 20881
asc_f64_1mcumgt_keeplast 20803

[stdout]
N 60000
desc_f32_cumge_right 4579
desc_f64_cumge_right 20859
desc_f64_cumgt_right 20820
desc_f64_exclge_none 20856
desc_f32_exclge_none 0
desc_f64_exclgt_none 20820
asc_f32_1mcumge_keeplast 10940
asc_f64_1mcumge_keeplast 20881
asc_f64_1mcumgt_keeplast 20803

/app/objectives.py

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

Exhaustive top_p final test

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(2024)
bad=0;ex=0
for i in range(200000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(1,11,(1,),generator=g).item())
    scale=float(torch.rand(1,generator=g).item()*15+0.01)
    x=torch.randn(B,V,generator=g)*scale
    r=torch.rand(1,generator=g).item()
    if r<0.3 and V>=2: x[:,0]=x[:,1]
    if r<0.15 and V>=3: x[:,2]=x[:,0]
    if r<0.08 and V>=4: x[:,3]=x[:,1]
    p=float(torch.rand(1,generator=g).item())
    m=i%6
    if m==0: p=1.0
    elif m==1: p=0.0
    elif m==2: x=torch.randint(0,3,(B,V),generator=g).float()
    elif m==3: p=float(torch.rand(1,generator=g).item())*0.02
    a=O.top_p_filter(x.clone(),p); b=R.top_p_filter(x.clone(),p)
    if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True):
        bad+=1
        if ex<5: ex+=1;print('p',round(p,4),'x',x.tolist());print(' mine',a.tolist());print(' ref ',b.tolist())
print('top_p bad',bad,'/ 200000')
"
output
top_p bad 0 / 200000

[stdout]
top_p bad 0 / 200000

Run both verification suites

shell
python3 verify.py && echo "===" && python3 verify2.py
output
total fails: 0
ALL PASS
===
total fails: 0
ALL PASS

[stdout]
total fails: 0
ALL PASS
===
total fails: 0
ALL PASS

Final 1500-trial high-variance sweep

shell
python3 -c "
import torch, reforacle as R, objectives as O
from collections import Counter
fails=[]
def cmp(n,a,b):
    if torch.is_tensor(a):
        if a.dtype!=b.dtype: fails.append((n,f'dtype {a.dtype}/{b.dtype}'));return
        if a.shape!=b.shape: fails.append((n,f'shape {tuple(a.shape)}/{tuple(b.shape)}'));return
        if a.is_floating_point():
            if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True): fails.append((n,'vals'))
        else:
            if not torch.equal(a,b): fails.append((n,'int'))
    elif a!=b: fails.append((n,f'{a}/{b}'))
def run(n,*a):
    try: m=getattr(O,n)(*[x.clone() if torch.is_tensor(x) else x for x in a])
    except Exception as e: fails.append((n,f'MINE {e!r}'));return
    try: r=getattr(R,n)(*[x.clone() if torch.is_tensor(x) else x for x in a])
    except Exception as e: fails.append((n,f'REF {e!r}'));return
    cmp(n,m,r)
g=torch.Generator().manual_seed(31337)
for t in range(1500):
    B=int(torch.randint(1,7,(1,),generator=g).item());V=int(torch.randint(1,20,(1,),generator=g).item());T=int(torch.randint(1,14,(1,),generator=g).item())
    sc=float(torch.rand(1,generator=g).item()*25+0.01)
    logits=torch.randn(B,V,generator=g)*sc
    if t%4==0 and V>=2: logits[:,0]=logits[:,1]
    run('greedy_token',logits)
    for T_ in [0.0,float(torch.rand(1,generator=g).item()*5+0.01)]: run('temperature_scale',logits,T_)
    run('top_k_filter',logits,int(torch.randint(1,V+3,(1,),generator=g).item()))
    for pp in [0.0,1.0,float(torch.rand(1,generator=g).item())]: run('top_p_filter',logits,pp)
    run('min_p_filter',logits,float(torch.rand(1,generator=g).item()))
    run('epsilon_filter',logits,float(torch.rand(1,generator=g).item()))
    run('renormalize_logits',logits); run('entropy_from_logits',logits.unsqueeze(0))
    run('topk_indices',logits,int(torch.randint(1,V+1,(1,),generator=g).item()))
    ids=torch.randint(0,V,(B,T),generator=g)
    run('repetition_penalty',logits,ids,float(torch.rand(1,generator=g).item()*4+0.1))
    run('frequency_penalty',logits,ids,float(torch.rand(1,generator=g).item()*3-0.5))
    run('presence_penalty',logits,ids,float(torch.rand(1,generator=g).item()*3-0.5))
    run('suppress_tokens',logits,torch.randint(0,V,(int(torch.randint(1,4,(1,),generator=g).item()),),generator=g))
    run('min_length_filter',logits,int(torch.randint(0,10,(1,),generator=g).item()),int(torch.randint(0,10,(1,),generator=g).item()),int(torch.randint(0,V,(1,),generator=g).item()))
    keep=(torch.rand(B,T,generator=g)>0.5).long()
    if t%9==0: keep[0]=0
    ids2=torch.randint(1,V+1,(B,T),generator=g)*keep
    run('attention_mask',ids2,0); run('position_ids',ids2,0); run('causal_mask',T)
    run('shift_labels',ids2,-100); run('sequence_lengths',ids2,0)
    pl=torch.randint(0,T+2,(B,),generator=g)
    run('completion_loss_labels',ids2,pl,-100); run('prepare_labels',ids2,pl,-100)
    run('mask_after_eos',ids2,int(torch.randint(0,V+1,(1,),generator=g).item()))
    run('generation_stop_mask',ids2,int(torch.randint(0,V+1,(1,),generator=g).item()),int(torch.randint(0,T+3,(1,),generator=g).item()))
    run('last_token_indices',keep)
    lg3=torch.randn(B,T,V,generator=g)*sc
    run('gather_last_logits',lg3,keep); run('rollout_last_logits',lg3,ids2,0)
    run('action_logprobs',lg3,torch.randint(0,V,(B,T),generator=g))
    run('truncate_left',ids2,int(torch.randint(1,T+3,(1,),generator=g).item()))
    run('pad_to_multiple',ids2,0,int(torch.randint(1,7,(1,),generator=g).item()))
    run('right_pad_to',ids2,0,int(torch.randint(1,T+4,(1,),generator=g).item()))
    rew=torch.randn(B,T,generator=g)*sc; mk=(torch.rand(B,T,generator=g)>0.4).float()
    run('reward_to_go',rew,mk); run('discount_cumsum',rew,float(torch.rand(1,generator=g).item()))
    run('discount_cumsum',rew[0],float(torch.rand(1,generator=g).item()))
    sl=torch.randint(0,5,(int(torch.randint(1,5,(1,),generator=g).item()),),generator=g)
    run('build_position_ids_packed',sl,max(1,int(sl.sum().item())+int(torch.randint(-3,4,(1,),generator=g).item())))
    run('logits_warper',logits,float(torch.rand(1,generator=g).item()*3+0.1),int(torch.randint(1,V+2,(1,),generator=g).item()),float(torch.rand(1,generator=g).item()))
    run('sample_greedy_warped',logits,float(torch.rand(1,generator=g).item()*3+0.1),int(torch.randint(1,V+2,(1,),generator=g).item()),float(torch.rand(1,generator=g).item()))
    m2=(torch.rand(B,V,generator=g)>0.4).long(); m2[:,0]=1
    run('masked_argmax',logits,m2)
    run('first_stop_index',ids2,torch.randint(0,V+1,(int(torch.randint(1,4,(1,),generator=g).item()),),generator=g))
c=Counter(n for n,_ in fails)
print('total',len(fails))
for n in sorted(c): print(' ',n,c[n],next(m for nn,m in fails if nn==n))
print('OK' if not fails else 'FAIL')
"
output
total 40
  logits_warper 6 vals
  sample_greedy_warped 6 int
  top_p_filter 28 vals
FAIL

[stdout]
total 40
  logits_warper 6 vals
  sample_greedy_warped 6 int
  top_p_filter 28 vals
FAIL

Find remaining top_p failures

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(31337)
ex=0
for i in range(400000):
    B=int(torch.randint(1,7,(1,),generator=g).item());V=int(torch.randint(1,20,(1,),generator=g).item())
    sc=float(torch.rand(1,generator=g).item()*25+0.01)
    x=torch.randn(B,V,generator=g)*sc
    if i%4==0 and V>=2: x[:,0]=x[:,1]
    for p in [0.0,1.0,float(torch.rand(1,generator=g).item())]:
        a=O.top_p_filter(x.clone(),p); b=R.top_p_filter(x.clone(),p)
        if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True):
            ex+=1
            if ex<=6:
                d=(torch.isinf(a)&(a<0)).int()-(torch.isinf(b)&(b<0)).int()
                print('p',round(p,5),'sc',round(sc,2))
                print(' x',[round(v,3) for v in x[(d!=0).any(-1).nonzero()[0,0]].tolist()])
                print(' mine rm',(torch.isinf(a)&(a<0)).int().tolist())
                print(' ref  rm',(torch.isinf(b)&(b<0)).int().tolist())
    if ex>6: break
print('ex',ex)
"
output
p 1.0 sc 4.46
 x [-3.315, -3.315, -5.347, 2.07, 3.477, 4.952, -1.493, 2.133, 12.583, -2.49, -10.532, -2.76, 3.208, -3.957, 2.542, 1.193, -0.533]
 mine rm [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0], [0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0]]
 ref  rm [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0], [1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0]]
p 0.0 sc 0.49
 x [0.832, 0.832, 0.452, -0.378, 0.541, -0.721, 0.383, 0.115, -0.303, 0.09, -0.405, 0.188, -0.569, 0.28, -0.637, 0.428, -0.558]
 mine rm [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1], [1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
 ref  rm [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1], [1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
p 1.0 sc 8.07
 x [6.119, 6.119, 8.081, -5.293, -1.422, -11.518, 3.418, 2.087, -4.743, 23.557, 0.925, -0.145, 14.443, -2.873, 1.341, -12.777, 2.793, 1.18, -4.655]
 mine rm [[0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1]]
 ref  rm [[1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1]]
p 0.0 sc 0.61
 x [1.058, 1.058, -0.899, -0.777, -0.452, 0.598, 0.855, 0.593, -1.183, 0.812, 0.851, 0.006, 0.871, -0.071, -0.192, -0.16, 0.072, 0.327, -0.889]
 mine rm [[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1], [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
 ref  rm [[1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1], [1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
p 0.42669 sc 1.67
 x [2.924, 2.924, 1.192, -2.948, -0.809, -1.546, -3.528, -2.567, -2.207, -1.892, -1.498, -1.55, -0.409, 2.914, 0.598, 3.804, -3.337, 1.108]
 mine rm [[1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1], [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1]]
 ref  rm [[1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1], [1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1]]
p 0.89883 sc 2.54
 x [1.764, 1.764, -5.029, -3.229, -0.654, -1.711, -0.966, 2.617, 2.384, -2.613, 3.33, -2.227, -0.698, 1.438, -0.123, 4.444, -0.621]
 mine rm [[0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 1], [1, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0]]
 ref  rm [[1, 0, 1, 1, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 1], [1, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0]]
ex 7

[stdout]
p 1.0 sc 4.46
 x [-3.315, -3.315, -5.347, 2.07, 3.477, 4.952, -1.493, 2.133, 12.583, -2.49, -10.532, -2.76, 3.208, -3.957, 2.542, 1.193, -0.533]
 mine rm [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0], [0, 1, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0]]
 ref  rm [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0], [1, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0]]
p 0.0 sc 0.49
 x [0.832, 0.832, 0.452, -0.378, 0.541, -0.721, 0.383, 0.115, -0.303, 0.09, -0.405, 0.188, -0.569, 0.28, -0.637, 0.428, -0.558]
 mine rm [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1], [1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
 ref  rm [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1], [1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
p 1.0 sc 8.07
 x [6.119, 6.119, 8.081, -5.293, -1.422, -11.518, 3.418, 2.087, -4.743, 23.557, 0.925, -0.145, 14.443, -2.873, 1.341, -12.777, 2.793, 1.18, -4.655]
 mine rm [[0, 1, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1]]
 ref  rm [[1, 0, 0, 1, 1, 1, 1, 1, 1, 0, 1, 1, 0, 1, 1, 1, 1, 1, 1]]
p 0.0 sc 0.61
 x [1.058, 1.058, -0.899, -0.777, -0.452, 0.598, 0.855, 0.593, -1.183, 0.812, 0.851, 0.006, 0.871, -0.071, -0.192, -0.16, 0.072, 0.327, -0.889]
 mine rm [[0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1], [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
 ref  rm [[1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1], [1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]]
p 0.42669 sc 1.67
 x [2.924, 2.924, 1.192, -2.948, -0.809, -1.546, -3.528, -2.567, -2.207, -1.892, -1.498, -1.55, -0.409, 2.914, 0.598, 3.804, -3.337, 1.108]
 mine rm [[1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1], [0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1]]
 ref  rm [[1, 1, 1, 1, 1, 0, 1, 1, 1, 1, 0, 1, 1, 1, 1, 1, 0, 1], [1, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 1, 1]]
p 0.89883 sc 2.54
 x [1.764, 1.764, -5.029, -3.229, -0.654, -1.711, -0.966, 2.617, 2.384, -2.613, 3.33, -2.227, -0.698, 1.438, -0.123, 4.444, -0.621]
 mine rm [[0, 1, 1, 1, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 1], [1, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0]]
 ref  rm [[1, 0, 1, 1, 1, 1, 1, 0, 0, 1, 0, 1, 1, 1, 1, 0, 1], [1, 1, 0, 0, 1, 1, 1, 1, 0, 1, 1, 1, 0, 1, 1, 1, 0]]
ex 7

Nail down reference tie ordering

shell
python3 -c "
import torch, reforacle as R
def rm(x,p): 
    o=R.top_p_filter(torch.tensor([x]),p); return (torch.isinf(o)&(o<0)).int()[0].tolist()
print('[0,0,0] p0.574 removed', rm([0.,0.,0.],0.574))
print('[0,0,0] p0.4 removed', rm([0.,0.,0.],0.4))
print('[0,0,0,0] p0.6 removed', rm([0.,0.,0.,0.],0.6))
print('[5,5,1] p0.5 removed', rm([5.,5.,1.],0.5))   # tie at top
print('[5,5,1] p0.9 removed', rm([5.,5.,1.],0.9))
print('[5,5,5,1] p0.5 removed', rm([5.,5.,5.,1.],0.5))
print('[1,5,5] p0.5 removed', rm([1.,5.,5.],0.5))
print('[5,1,5] p0.5 removed', rm([5.,1.,5.],0.5))
# torch orders
for desc in [True,False]:
  for st in [True,False]:
    v,ix=torch.sort(torch.tensor([[5.,5.,1.]]),descending=desc,stable=st)
    print(f'sort desc={desc} stable={st}',ix.tolist())
"
output
[0,0,0] p0.574 removed [0, 0, 1]
[0,0,0] p0.4 removed [0, 0, 1]
[0,0,0,0] p0.6 removed [0, 0, 0, 1]
[5,5,1] p0.5 removed [0, 0, 1]
[5,5,1] p0.9 removed [0, 0, 1]
[5,5,5,1] p0.5 removed [0, 0, 1, 1]
[1,5,5] p0.5 removed [1, 0, 0]
[5,1,5] p0.5 removed [0, 1, 0]
sort desc=True stable=True [[0, 1, 2]]
sort desc=True stable=False [[0, 1, 2]]
sort desc=False stable=True [[2, 0, 1]]
sort desc=False stable=False [[2, 0, 1]]

[stdout]
[0,0,0] p0.574 removed [0, 0, 1]
[0,0,0] p0.4 removed [0, 0, 1]
[0,0,0,0] p0.6 removed [0, 0, 0, 1]
[5,5,1] p0.5 removed [0, 0, 1]
[5,5,1] p0.9 removed [0, 0, 1]
[5,5,5,1] p0.5 removed [0, 0, 1, 1]
[1,5,5] p0.5 removed [1, 0, 0]
[5,1,5] p0.5 removed [0, 1, 0]
sort desc=True stable=True [[0, 1, 2]]
sort desc=True stable=False [[0, 1, 2]]
sort desc=False stable=True [[2, 0, 1]]
sort desc=False stable=False [[2, 0, 1]]

/app/objectives.py

− removed
    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1, stable=True)
    probs = sorted_logits.softmax(dim=-1)
+ added
    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
    probs = sorted_logits.softmax(dim=-1)

Re-test top_p with non-stable sort

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(31337)
bad=0
for i in range(400000):
    B=int(torch.randint(1,7,(1,),generator=g).item());V=int(torch.randint(1,20,(1,),generator=g).item())
    sc=float(torch.rand(1,generator=g).item()*25+0.01)
    x=torch.randn(B,V,generator=g)*sc
    if i%4==0 and V>=2: x[:,0]=x[:,1]
    if i%6==0 and V>=3: x[:,2]=x[:,1]
    for p in [0.0,1.0,float(torch.rand(1,generator=g).item())]:
        a=O.top_p_filter(x.clone(),p); b=R.top_p_filter(x.clone(),p)
        if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True): bad+=1
print('top_p bad',bad)
"
output
top_p bad 0

[stdout]
top_p bad 0

Heavy-tie topk test

shell
python3 -c "
import torch, reforacle as R, objectives as O
# dedicated topk heavy-tie test, larger V
g=torch.Generator().manual_seed(8)
bad_stable=0; bad_topk=0
for i in range(30000):
    B=int(torch.randint(1,4,(1,),generator=g).item()); V=int(torch.randint(2,15,(1,),generator=g).item())
    x=torch.randint(0,4,(B,V),generator=g).float()
    k=int(torch.randint(1,V+1,(1,),generator=g).item())
    ref=R.topk_indices(x.clone(),k)
    if not torch.equal(O.topk_indices(x.clone(),k),ref): bad_stable+=1
print('topk_indices bad',bad_stable)
"
output
topk_indices bad 0

[stdout]
topk_indices bad 0

Final full sweep rerun

shell
python3 -c "
import torch, reforacle as R, objectives as O
from collections import Counter
fails=[]
def cmp(n,a,b):
    if torch.is_tensor(a):
        if a.dtype!=b.dtype: fails.append((n,f'dtype {a.dtype}/{b.dtype}'));return
        if a.shape!=b.shape: fails.append((n,f'shape'));return
        if a.is_floating_point():
            if not torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True): fails.append((n,'vals'))
        else:
            if not torch.equal(a,b): fails.append((n,'int'))
    elif a!=b: fails.append((n,f'{a}/{b}'))
def run(n,*a):
    try: m=getattr(O,n)(*[x.clone() if torch.is_tensor(x) else x for x in a])
    except Exception as e: fails.append((n,f'MINE {e!r}'));return
    try: r=getattr(R,n)(*[x.clone() if torch.is_tensor(x) else x for x in a])
    except Exception as e: fails.append((n,f'REF {e!r}'));return
    cmp(n,m,r)
g=torch.Generator().manual_seed(31337)
for t in range(1500):
    B=int(torch.randint(1,7,(1,),generator=g).item());V=int(torch.randint(1,20,(1,),generator=g).item());T=int(torch.randint(1,14,(1,),generator=g).item())
    sc=float(torch.rand(1,generator=g).item()*25+0.01)
    logits=torch.randn(B,V,generator=g)*sc
    if t%4==0 and V>=2: logits[:,0]=logits[:,1]
    if t%6==0 and V>=3: logits[:,2]=logits[:,1]
    run('greedy_token',logits)
    for T_ in [0.0,float(torch.rand(1,generator=g).item()*5+0.01)]: run('temperature_scale',logits,T_)
    run('top_k_filter',logits,int(torch.randint(1,V+3,(1,),generator=g).item()))
    for pp in [0.0,1.0,float(torch.rand(1,generator=g).item())]: run('top_p_filter',logits,pp)
    run('min_p_filter',logits,float(torch.rand(1,generator=g).item()))
    run('epsilon_filter',logits,float(torch.rand(1,generator=g).item()))
    run('renormalize_logits',logits); run('entropy_from_logits',logits.unsqueeze(0))
    run('topk_indices',logits,int(torch.randint(1,V+1,(1,),generator=g).item()))
    ids=torch.randint(0,V,(B,T),generator=g)
    run('repetition_penalty',logits,ids,float(torch.rand(1,generator=g).item()*4+0.1))
    run('frequency_penalty',logits,ids,float(torch.rand(1,generator=g).item()*3-0.5))
    run('presence_penalty',logits,ids,float(torch.rand(1,generator=g).item()*3-0.5))
    run('suppress_tokens',logits,torch.randint(0,V,(int(torch.randint(1,4,(1,),generator=g).item()),),generator=g))
    run('min_length_filter',logits,int(torch.randint(0,10,(1,),generator=g).item()),int(torch.randint(0,10,(1,),generator=g).item()),int(torch.randint(0,V,(1,),generator=g).item()))
    keep=(torch.rand(B,T,generator=g)>0.5).long()
    if t%9==0: keep[0]=0
    ids2=torch.randint(1,V+1,(B,T),generator=g)*keep
    run('attention_mask',ids2,0); run('position_ids',ids2,0); run('causal_mask',T)
    run('shift_labels',ids2,-100); run('sequence_lengths',ids2,0)
    pl=torch.randint(0,T+2,(B,),generator=g)
    run('completion_loss_labels',ids2,pl,-100); run('prepare_labels',ids2,pl,-100)
    run('mask_after_eos',ids2,int(torch.randint(0,V+1,(1,),generator=g).item()))
    run('generation_stop_mask',ids2,int(torch.randint(0,V+1,(1,),generator=g).item()),int(torch.randint(0,T+3,(1,),generator=g).item()))
    run('last_token_indices',keep)
    lg3=torch.randn(B,T,V,generator=g)*sc
    run('gather_last_logits',lg3,keep); run('rollout_last_logits',lg3,ids2,0)
    run('action_logprobs',lg3,torch.randint(0,V,(B,T),generator=g))
    run('truncate_left',ids2,int(torch.randint(1,T+3,(1,),generator=g).item()))
    run('pad_to_multiple',ids2,0,int(torch.randint(1,7,(1,),generator=g).item()))
    run('right_pad_to',ids2,0,int(torch.randint(1,T+4,(1,),generator=g).item()))
    rew=torch.randn(B,T,generator=g)*sc; mk=(torch.rand(B,T,generator=g)>0.4).float()
    run('reward_to_go',rew,mk); run('discount_cumsum',rew,float(torch.rand(1,generator=g).item()))
    run('discount_cumsum',rew[0],float(torch.rand(1,generator=g).item()))
    sl=torch.randint(0,5,(int(torch.randint(1,5,(1,),generator=g).item()),),generator=g)
    run('build_position_ids_packed',sl,max(1,int(sl.sum().item())+int(torch.randint(-3,4,(1,),generator=g).item())))
    run('logits_warper',logits,float(torch.rand(1,generator=g).item()*3+0.1),int(torch.randint(1,V+2,(1,),generator=g).item()),float(torch.rand(1,generator=g).item()))
    run('sample_greedy_warped',logits,float(torch.rand(1,generator=g).item()*3+0.1),int(torch.randint(1,V+2,(1,),generator=g).item()),float(torch.rand(1,generator=g).item()))
    m2=(torch.rand(B,V,generator=g)>0.4).long(); m2[:,0]=1
    run('masked_argmax',logits,m2)
    run('first_stop_index',ids2,torch.randint(0,V+1,(int(torch.randint(1,4,(1,),generator=g).item()),),generator=g))
c=Counter(n for n,_ in fails)
print('total',len(fails))
for n in sorted(c): print(' ',n,c[n],next(m for nn,m in fails if nn==n))
print('OK' if not fails else 'FAIL')
"
output
total 0
OK

[stdout]
total 0
OK

Warper T=0 edge and standalone import verification

shell
python3 -c "
import torch, reforacle as R, objectives as O
g=torch.Generator().manual_seed(9)
# warper with T=0 edge
bad=0
for i in range(3000):
    B=int(torch.randint(1,4,(1,),generator=g).item());V=int(torch.randint(2,10,(1,),generator=g).item())
    x=torch.randn(B,V,generator=g)*float(torch.rand(1,generator=g).item()*10)
    T=0.0 if i%3==0 else float(torch.rand(1,generator=g).item()*3+0.1)
    k=int(torch.randint(1,V+2,(1,),generator=g).item()); p=float(torch.rand(1,generator=g).item())
    for fn in ['logits_warper','sample_greedy_warped']:
        a=getattr(O,fn)(x.clone(),T,k,p); b=getattr(R,fn)(x.clone(),T,k,p)
        ok=torch.equal(a,b) if not a.is_floating_point() else torch.allclose(a,b,atol=1e-5,rtol=1e-4,equal_nan=True)
        if not ok: bad+=1
print('warper T=0 edge bad',bad)
" 
echo "--- standalone import check ---"
python3 -c "
import sys, types
# block reforacle to simulate clean sandbox
sys.modules['reforacle']=None
import importlib
sys.modules.pop('objectives',None)
import objectives
print('imports OK, deps:', [m for m in ['reforacle'] if getattr(objectives,m,'x')!='x'])
import torch
print('greedy works:', objectives.greedy_token(torch.tensor([[1.,3.,2.]])).tolist())
"
grep -n "import" /app/objectives.py
output
warper T=0 edge bad 0
--- standalone import check ---
imports OK, deps: []
greedy works: [1]
1:from __future__ import annotations
3:import torch

[stdout]
warper T=0 edge bad 0
--- standalone import check ---
imports OK, deps: []
greedy works: [1]
1:from __future__ import annotations
3:import torch

Clean up scratch files and check budget

shell
rm -f /app/p.py /app/verify.py /app/verify2.py; bash /app/timer.sh; echo "---"; ls /app
output
Remaining time (hours:minutes):
1:32
---
__pycache__
objectives.py
probe.py
reforacle.cpython-312-aarch64-linux-gnu.so
timer.sh

[stdout]
Remaining time (hours:minutes):
1:32
---
__pycache__
objectives.py
probe.py
reforacle.cpython-312-aarch64-linux-gnu.so
timer.sh
All 37 functions are implemented and verified. Here's a summary of what I did. ## Approach Since `reforacle` was the only ground truth, I recovered each function's exact behavior purely by probing it and comparing, then wrote a standalone `torch`-only implementation. I built randomized comparison harnesses (200k+ trials total) covering ties, extreme magnitudes, all-pad rows, underflow, and `p=0`/`p=1` boundaries to catch cases that random continuous inputs miss. ## Key findings that differed from "standard" conventions - **`repetition_penalty`** is applied *per occurrence* (`penalty**count`), not once per unique token. - **`argmax` family** (`greedy_token`, `masked_argmax`, `temperature_scale(T=0)`, `sample_greedy_warped`) breaks ties toward the **highest** index , implemented via flipped argmax. - **`position_ids`** uses `clamp(cumsum(mask)-1, min=0)` (not the HF `masked_fill(1)` variant). - **`last_token_indices`** returns `-1` for an all-pad row, while **`gather_last_logits`/`rollout_last_logits`** clamp that to index `0`. - **`prepare_labels`** shifts labels left *and* masks the first `prompt_length-1`; **`completion_loss_labels`** masks the prompt with no shift. - **`generation_stop_mask`** keeps up to the first EOS but forces at least `min_len` tokens. - **`epsilon_filter`** keeps exactly one token (the last-argmax) when the threshold removes everything. - **`build_position_ids_packed`** fills leftover positions with zeros, not a continuing range. - **`top_p_filter`** was the trickiest: it uses **descending non-stable `torch.sort`**, computes the exclusive prefix sum `cumsum(probs) - probs`, and removes where that is `>= p` (keeping index 0). This exactly reproduces the reference's float32 rounding behavior at `p=1.0` with dominating tokens *and* its content-dependent tie ordering , both of which the algebraically-equivalent HF formulations get wrong. The final `/app/objectives.py` depends only on `torch`, imports cleanly with `reforacle` absent, and passed every function across all stress sweeps.
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: 36/37  (partial_score=0.973)
REWARD = 0   (1 iff ALL 37 functions match the sealed reference exactly)

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

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

Trial trial_53291fb542564c98 · verifier authoritative; classifier explanatory.