tests/test.sh exit code (0 → resolved); the classification below is post-hoc and cannot change it.Check budget and probe pattern
bash /app/timer.sh; echo "---PROBE---"; cat /app/probe.py
Remaining time (hours:minutes):
1:59
---PROBE---
from __future__ import annotations
import torch
import reforacle
g = torch.Generator().manual_seed(0)
logits = torch.randn(2, 6, generator=g)
print("logits:", logits.tolist())
print("reference greedy_token ->", reforacle.greedy_token(logits).tolist())
print("reference temperature_scale(T=0) ->", reforacle.temperature_scale(logits, 0.0).tolist())
print()
print("Probe any reference function this way to discover its exact behavior,")
print("then implement it in /app/objectives.py to match.")
[stdout]
Remaining time (hours:minutes):
1:59
---PROBE---
from __future__ import annotations
import torch
import reforacle
g = torch.Generator().manual_seed(0)
logits = torch.randn(2, 6, generator=g)
print("logits:", logits.tolist())
print("reference greedy_token ->", reforacle.greedy_token(logits).tolist())
print("reference temperature_scale(T=0) ->", reforacle.temperature_scale(logits, 0.0).tolist())
print()
print("Probe any reference function this way to discover its exact behavior,")
print("then implement it in /app/objectives.py to match.")
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
/app/objectives.py
Remaining time (hours:minutes):
1:59
---PROBE---
from __future__ import annotations
import torch
import reforacle
g = torch.Generator().manual_seed(0)
logits = torch.randn(2, 6, generator=g)
print("logits:", logits.tolist())
print("reference greedy_token ->", reforacle.greedy_token(logits).tolist())
print("reference temperature_scale(T=0) ->", reforacle.temperature_scale(logits, 0.0).tolist())
print()
print("Probe any reference function this way to discover its exact behavior,")
print("then implement it in /app/objectives.py to match.")
[stdout]
Remaining time (hours:minutes):
1:59
---PROBE---
from __future__ import annotations
import torch
import reforacle
g = torch.Generator().manual_seed(0)
logits = torch.randn(2, 6, generator=g)
print("logits:", logits.tolist())
print("reference greedy_token ->", reforacle.greedy_token(logits).tolist())
print("reference temperature_scale(T=0) ->", reforacle.temperature_scale(logits, 0.0).tolist())
print()
print("Probe any reference function this way to discover its exact behavior,")
print("then implement it in /app/objectives.py to match.")
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
Probe basic functions
python3 -c "
import torch, reforacle
g = torch.Generator().manual_seed(0)
logits = torch.randn(2,6,generator=g)
print('logits', logits)
print('greedy', reforacle.greedy_token(logits), reforacle.greedy_token(logits).dtype, reforacle.greedy_token(logits).shape)
print('temp0', reforacle.temperature_scale(logits,0.0))
print('temp2', reforacle.temperature_scale(logits,2.0))
print('renorm', reforacle.renormalize_logits(logits))
print('logsoftmax', torch.log_softmax(logits,dim=-1))
"
logits tensor([[ 1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986],
[ 0.4033, 0.8380, -0.7193, -0.4033, -0.5966, 0.1820]])
greedy tensor([0, 1]) torch.int64 torch.Size([2])
temp0 tensor([[0., -inf, -inf, -inf, -inf, -inf],
[-inf, 0., -inf, -inf, -inf, -inf]])
temp2 tensor([[ 0.7705, -0.1467, -1.0894, 0.2842, -0.5423, -0.6993],
[ 0.2017, 0.4190, -0.3596, -0.2017, -0.2983, 0.0910]])
renorm tensor([[-0.5232, -2.3576, -4.2429, -1.4957, -3.1487, -3.4627],
[-1.5009, -1.0662, -2.6235, -2.3076, -2.5009, -1.7222]])
logsoftmax tensor([[-0.5232, -2.3576, -4.2429, -1.4957, -3.1487, -3.4627],
[-1.5009, -1.0662, -2.6235, -2.3076, -2.5009, -1.7222]])
[stdout]
logits tensor([[ 1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986],
[ 0.4033, 0.8380, -0.7193, -0.4033, -0.5966, 0.1820]])
greedy tensor([0, 1]) torch.int64 torch.Size([2])
temp0 tensor([[0., -inf, -inf, -inf, -inf, -inf],
[-inf, 0., -inf, -inf, -inf, -inf]])
temp2 tensor([[ 0.7705, -0.1467, -1.0894, 0.2842, -0.5423, -0.6993],
[ 0.2017, 0.4190, -0.3596, -0.2017, -0.2983, 0.0910]])
renorm tensor([[-0.5232, -2.3576, -4.2429, -1.4957, -3.1487, -3.4627],
[-1.5009, -1.0662, -2.6235, -2.3076, -2.5009, -1.7222]])
logsoftmax tensor([[-0.5232, -2.3576, -4.2429, -1.4957, -3.1487, -3.4627],
[-1.5009, -1.0662, -2.6235, -2.3076, -2.5009, -1.7222]])Probe filters
python3 -c "
import torch, reforacle
g = torch.Generator().manual_seed(0)
logits = torch.randn(2,6,generator=g)
print('logits', logits)
print('top_k2', reforacle.top_k_filter(logits,2))
print('top_p0.8', reforacle.top_p_filter(logits,0.8))
print('min_p0.3', reforacle.min_p_filter(logits,0.3))
print('epsilon0.1', reforacle.epsilon_filter(logits,0.1))
print('topk_indices2', reforacle.topk_indices(logits,2))
"
logits tensor([[ 1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986],
[ 0.4033, 0.8380, -0.7193, -0.4033, -0.5966, 0.1820]])
top_k2 tensor([[1.5410, -inf, -inf, 0.5684, -inf, -inf],
[0.4033, 0.8380, -inf, -inf, -inf, -inf]])
top_p0.8 tensor([[ 1.5410, -inf, -inf, 0.5684, -inf, -inf],
[ 0.4033, 0.8380, -inf, -0.4033, -inf, 0.1820]])
min_p0.3 tensor([[1.5410, -inf, -inf, 0.5684, -inf, -inf],
[0.4033, 0.8380, -inf, -inf, -inf, 0.1820]])
epsilon0.1 tensor([[1.5410, -inf, -inf, 0.5684, -inf, -inf],
[0.4033, 0.8380, -inf, -inf, -inf, 0.1820]])
topk_indices2 tensor([[0, 3],
[1, 0]])
[stdout]
logits tensor([[ 1.5410, -0.2934, -2.1788, 0.5684, -1.0845, -1.3986],
[ 0.4033, 0.8380, -0.7193, -0.4033, -0.5966, 0.1820]])
top_k2 tensor([[1.5410, -inf, -inf, 0.5684, -inf, -inf],
[0.4033, 0.8380, -inf, -inf, -inf, -inf]])
top_p0.8 tensor([[ 1.5410, -inf, -inf, 0.5684, -inf, -inf],
[ 0.4033, 0.8380, -inf, -0.4033, -inf, 0.1820]])
min_p0.3 tensor([[1.5410, -inf, -inf, 0.5684, -inf, -inf],
[0.4033, 0.8380, -inf, -inf, -inf, 0.1820]])
epsilon0.1 tensor([[1.5410, -inf, -inf, 0.5684, -inf, -inf],
[0.4033, 0.8380, -inf, -inf, -inf, 0.1820]])
topk_indices2 tensor([[0, 3],
[1, 0]])Probe penalties
python3 -c "
import torch, reforacle
g = torch.Generator().manual_seed(0)
logits = torch.randn(2,6,generator=g)
ids = torch.tensor([[0,0,3,1],[1,1,1,5]])
print('rep1.5', reforacle.repetition_penalty(logits,ids,1.5))
print('freq0.5', reforacle.frequency_penalty(logits,ids,0.5))
print('pres0.5', reforacle.presence_penalty(logits,ids,0.5))
print('suppress', reforacle.suppress_tokens(logits,[1,2]))
print('minlen', reforacle.min_length_filter(logits,3,5,4))
"
rep1.5 tensor([[ 0.6849, -0.4401, -2.1788, 0.3790, -1.0845, -1.3986],
[ 0.4033, 0.2483, -0.7193, -0.4033, -0.5966, 0.1214]])
freq0.5 tensor([[ 0.5410, -0.7934, -2.1788, 0.0684, -1.0845, -1.3986],
[ 0.4033, -0.6620, -0.7193, -0.4033, -0.5966, -0.3180]])
pres0.5 tensor([[ 1.0410, -0.7934, -2.1788, 0.0684, -1.0845, -1.3986],
[ 0.4033, 0.3380, -0.7193, -0.4033, -0.5966, -0.3180]])
suppress tensor([[ 1.5410, -inf, -inf, 0.5684, -1.0845, -1.3986],
[ 0.4033, -inf, -inf, -0.4033, -0.5966, 0.1820]])
minlen tensor([[ 1.5410, -0.2934, -2.1788, 0.5684, -inf, -1.3986],
[ 0.4033, 0.8380, -0.7193, -0.4033, -inf, 0.1820]])
[stdout]
rep1.5 tensor([[ 0.6849, -0.4401, -2.1788, 0.3790, -1.0845, -1.3986],
[ 0.4033, 0.2483, -0.7193, -0.4033, -0.5966, 0.1214]])
freq0.5 tensor([[ 0.5410, -0.7934, -2.1788, 0.0684, -1.0845, -1.3986],
[ 0.4033, -0.6620, -0.7193, -0.4033, -0.5966, -0.3180]])
pres0.5 tensor([[ 1.0410, -0.7934, -2.1788, 0.0684, -1.0845, -1.3986],
[ 0.4033, 0.3380, -0.7193, -0.4033, -0.5966, -0.3180]])
suppress tensor([[ 1.5410, -inf, -inf, 0.5684, -1.0845, -1.3986],
[ 0.4033, -inf, -inf, -0.4033, -0.5966, 0.1820]])
minlen tensor([[ 1.5410, -0.2934, -2.1788, 0.5684, -inf, -1.3986],
[ 0.4033, 0.8380, -0.7193, -0.4033, -inf, 0.1820]])Test rep penalty power and minlen boundary
python3 -c "
import torch, reforacle
logits = torch.tensor([[2.0, -1.0, 0.5, -2.0, 1.0, 0.0]])
ids = torch.tensor([[1,1,3]]) # token1 (neg) twice, token3 (neg) once
print('rep2', reforacle.repetition_penalty(logits,ids,2.0))
# single-apply HF would give: token1 -1*2=-2, token3 -2*2=-4
# power: token1 -1*4=-4, token3 -4
print('minlen_ge', reforacle.min_length_filter(logits,5,5,4))
print('minlen_gt', reforacle.min_length_filter(logits,6,5,4))
"
rep2 tensor([[ 2.0000, -4.0000, 0.5000, -4.0000, 1.0000, 0.0000]]) minlen_ge tensor([[ 2.0000, -1.0000, 0.5000, -2.0000, 1.0000, 0.0000]]) minlen_gt tensor([[ 2.0000, -1.0000, 0.5000, -2.0000, 1.0000, 0.0000]]) [stdout] rep2 tensor([[ 2.0000, -4.0000, 0.5000, -4.0000, 1.0000, 0.0000]]) minlen_ge tensor([[ 2.0000, -1.0000, 0.5000, -2.0000, 1.0000, 0.0000]]) minlen_gt tensor([[ 2.0000, -1.0000, 0.5000, -2.0000, 1.0000, 0.0000]])
Probe collation utils
python3 -c "
import torch, reforacle
ids = torch.tensor([[5,6,7,0,0],[1,2,0,0,0]])
print('attn', reforacle.attention_mask(ids,0), reforacle.attention_mask(ids,0).dtype)
print('posids', reforacle.position_ids(ids,0), reforacle.position_ids(ids,0).dtype)
print('causal', reforacle.causal_mask(3), reforacle.causal_mask(3).dtype)
print('seqlen', reforacle.sequence_lengths(ids,0))
print('lasttok', reforacle.last_token_indices(reforacle.attention_mask(ids,0)))
"
attn tensor([[1., 1., 1., 0., 0.],
[1., 1., 0., 0., 0.]]) torch.float32
posids tensor([[0, 1, 2, 2, 2],
[0, 1, 1, 1, 1]]) torch.int64
causal tensor([[0., -inf, -inf],
[0., 0., -inf],
[0., 0., 0.]]) torch.float32
seqlen tensor([3, 2])
lasttok tensor([2, 1])
[stdout]
attn tensor([[1., 1., 1., 0., 0.],
[1., 1., 0., 0., 0.]]) torch.float32
posids tensor([[0, 1, 2, 2, 2],
[0, 1, 1, 1, 1]]) torch.int64
causal tensor([[0., -inf, -inf],
[0., 0., -inf],
[0., 0., 0.]]) torch.float32
seqlen tensor([3, 2])
lasttok tensor([2, 1])Probe label/mask utils
python3 -c "
import torch, reforacle
ids = torch.tensor([[5,6,7,8],[1,2,3,4]])
print('shift', reforacle.shift_labels(ids,-100))
print('mask_after_eos', reforacle.mask_after_eos(torch.tensor([[5,9,6,9,7],[9,1,2,3,4]]),9))
print('completion', reforacle.completion_loss_labels(ids, torch.tensor([2,1]), -100))
print('prepare_labels', reforacle.prepare_labels(ids, torch.tensor([2,1]), -100))
g=torch.Generator().manual_seed(1)
logits=torch.randn(2,4,3,generator=g)
am=torch.tensor([[1.,1.,0.,0.],[1.,1.,1.,0.]])
print('gather_last', reforacle.gather_last_logits(logits,am))
print('leftpos', reforacle.position_ids(torch.tensor([[0,0,5,6]]),0))
"
shift tensor([[ 6, 7, 8, -100],
[ 2, 3, 4, -100]])
mask_after_eos tensor([[1., 1., 0., 0., 0.],
[1., 0., 0., 0., 0.]])
completion tensor([[-100, -100, 7, 8],
[-100, 2, 3, 4]])
prepare_labels tensor([[-100, 7, 8, -100],
[ 2, 3, 4, -100]])
gather_last tensor([[-1.6095, -0.1002, -0.6092],
[ 0.2753, 1.7163, -0.0561]])
leftpos tensor([[0, 0, 0, 1]])
[stdout]
shift tensor([[ 6, 7, 8, -100],
[ 2, 3, 4, -100]])
mask_after_eos tensor([[1., 1., 0., 0., 0.],
[1., 0., 0., 0., 0.]])
completion tensor([[-100, -100, 7, 8],
[-100, 2, 3, 4]])
prepare_labels tensor([[-100, 7, 8, -100],
[ 2, 3, 4, -100]])
gather_last tensor([[-1.6095, -0.1002, -0.6092],
[ 0.2753, 1.7163, -0.0561]])
leftpos tensor([[0, 0, 0, 1]])Probe truncate, pad, action_logprobs, entropy
python3 -c "
import torch, reforacle
ids = torch.tensor([[1,2,3,4,5],[6,7,8,9,10]])
print('trunc_left3', reforacle.truncate_left(ids,3))
print('pad_mult', reforacle.pad_to_multiple(torch.tensor([[1,2,3],[4,5,6]]),0,4))
print('rightpad', reforacle.right_pad_to(torch.tensor([[1,2],[3,4]]),0,5))
g=torch.Generator().manual_seed(2)
logits=torch.randn(2,5,4,generator=g)
actions=torch.tensor([[0,1,2,3,0],[3,2,1,0,1]])
print('action_lp', reforacle.action_logprobs(logits,actions))
lp=torch.log_softmax(logits,-1)
print('manual', lp.gather(-1,actions.unsqueeze(-1)).squeeze(-1))
print('entropy', reforacle.entropy_from_logits(logits))
"
trunc_left3 tensor([[ 3, 4, 5],
[ 8, 9, 10]])
pad_mult tensor([[1, 2, 3, 0],
[4, 5, 6, 0]])
rightpad tensor([[1, 2, 0, 0, 0],
[3, 4, 0, 0, 0]])
action_lp tensor([[-2.2806, -0.3895, -1.7254, -1.9711, -2.5974],
[-1.1030, -1.3539, -1.0903, -1.0662, -3.1551]])
manual tensor([[-2.2806, -0.3895, -1.7254, -1.9711, -2.5974],
[-1.1030, -1.3539, -1.0903, -1.0662, -3.1551]])
entropy tensor([[0.8911, 0.9418, 1.2761, 1.1208, 1.1223],
[1.2744, 1.2201, 1.0745, 1.2135, 0.9645]])
[stdout]
trunc_left3 tensor([[ 3, 4, 5],
[ 8, 9, 10]])
pad_mult tensor([[1, 2, 3, 0],
[4, 5, 6, 0]])
rightpad tensor([[1, 2, 0, 0, 0],
[3, 4, 0, 0, 0]])
action_lp tensor([[-2.2806, -0.3895, -1.7254, -1.9711, -2.5974],
[-1.1030, -1.3539, -1.0903, -1.0662, -3.1551]])
manual tensor([[-2.2806, -0.3895, -1.7254, -1.9711, -2.5974],
[-1.1030, -1.3539, -1.0903, -1.0662, -3.1551]])
entropy tensor([[0.8911, 0.9418, 1.2761, 1.1208, 1.1223],
[1.2744, 1.2201, 1.0745, 1.2135, 0.9645]])Probe entropy, reward_to_go, discount, masked_argmax, first_stop
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(2)
logits=torch.randn(2,5,4,generator=g)
p=torch.softmax(logits,-1)
ent=-(p*torch.log_softmax(logits,-1)).sum(-1)
print('ent match', torch.allclose(ent, reforacle.entropy_from_logits(logits)))
# reward_to_go and discount_cumsum
r=torch.tensor([[1.,2.,3.,4.],[5.,6.,7.,8.]])
mask=torch.tensor([[1.,1.,1.,0.],[1.,1.,0.,0.]])
print('r2g', reforacle.reward_to_go(r,mask))
print('disc', reforacle.discount_cumsum(r,0.9))
print('masked_argmax', reforacle.masked_argmax(torch.tensor([[1.,5.,2.,3.],[4.,1.,9.,2.]]), torch.tensor([[1.,0.,1.,1.],[1.,1.,0.,1.]])))
print('first_stop', reforacle.first_stop_index(torch.tensor([[1,2,3,4,2],[5,2,3,2,1]]),[2,3]))
"
Exit code 1
ent match True
r2g tensor([[ 6., 5., 3., 0.],
[11., 6., 0., 0.]])
disc tensor([[ 8.1460, 7.9400, 6.6000, 4.0000],
[21.9020, 18.7800, 14.2000, 8.0000]])
masked_argmax tensor([3, 0])
Traceback (most recent call last):
File "<string>", line 14, in <module>
File "reforacle.py", line 243, in reforacle.first_stop_index
AttributeError: 'list' object has no attribute 'tolist'
[error] tool reported failureProbe first_stop, gen_stop, build_pos_packed
python3 -c "
import torch, reforacle
print('first_stop', reforacle.first_stop_index(torch.tensor([[1,2,3,4,2],[5,2,3,2,1]]),torch.tensor([2,3])))
print('first_stop_none', reforacle.first_stop_index(torch.tensor([[1,4,7,8,9]]),torch.tensor([2,3])))
print('gen_stop', reforacle.generation_stop_mask(torch.tensor([[5,9,6,9],[9,1,2,3]]),9,2))
print('build_pos_packed', reforacle.build_position_ids_packed(torch.tensor([2,3,1]),6))
"
first_stop tensor([1, 1])
first_stop_none tensor([-1])
gen_stop tensor([[1., 1., 0., 0.],
[1., 1., 0., 0.]])
build_pos_packed tensor([0, 1, 0, 1, 2, 0])
[stdout]
first_stop tensor([1, 1])
first_stop_none tensor([-1])
gen_stop tensor([[1., 1., 0., 0.],
[1., 1., 0., 0.]])
build_pos_packed tensor([0, 1, 0, 1, 2, 0])Nail generation_stop_mask
python3 -c "
import torch, reforacle
# gen_stop cases
print('A eos after minlen', reforacle.generation_stop_mask(torch.tensor([[1,2,9,4,5]]),9,2))
print('B no eos', reforacle.generation_stop_mask(torch.tensor([[1,2,3,4,5]]),9,2))
print('C eos at end', reforacle.generation_stop_mask(torch.tensor([[1,2,3,9]]),9,2))
print('D minlen0', reforacle.generation_stop_mask(torch.tensor([[9,1,2]]),9,0))
print('E minlen1', reforacle.generation_stop_mask(torch.tensor([[9,1,2]]),9,1))
print('F eos exactly minlen', reforacle.generation_stop_mask(torch.tensor([[1,9,3,9]]),9,3))
"
A eos after minlen tensor([[1., 1., 1., 0., 0.]]) B no eos tensor([[1., 1., 1., 1., 1.]]) C eos at end tensor([[1., 1., 1., 1.]]) D minlen0 tensor([[1., 0., 0.]]) E minlen1 tensor([[1., 0., 0.]]) F eos exactly minlen tensor([[1., 1., 1., 0.]]) [stdout] A eos after minlen tensor([[1., 1., 1., 0., 0.]]) B no eos tensor([[1., 1., 1., 1., 1.]]) C eos at end tensor([[1., 1., 1., 1.]]) D minlen0 tensor([[1., 0., 0.]]) E minlen1 tensor([[1., 0., 0.]]) F eos exactly minlen tensor([[1., 1., 1., 0.]])
Probe warper, sample_greedy, rollout_last
python3 -c "
import torch, reforacle
g=torch.Generator().manual_seed(3)
logits=torch.randn(2,6,generator=g)
print('logits',logits)
print('warper T1.5 k3 p0.9', reforacle.logits_warper(logits,1.5,3,0.9))
print('sample_greedy_warped', reforacle.sample_greedy_warped(logits,1.5,3,0.9))
# rollout_last_logits
g2=torch.Generator().manual_seed(4)
lg=torch.randn(2,5,4,generator=g2)
ids=torch.tensor([[1,2,3,0,0],[4,5,6,7,0]])
print('rollout_last', reforacle.rollout_last_logits(lg,ids,0))
"
logits tensor([[ 0.8033, 0.1748, 0.0890, -0.6137, 0.0462, -1.3683],
[ 0.3375, 1.0111, -1.4352, 0.9774, 0.5220, 1.2379]])
warper T1.5 k3 p0.9 tensor([[0.5355, 0.1166, 0.0593, -inf, -inf, -inf],
[ -inf, 0.6741, -inf, 0.6516, -inf, 0.8253]])
sample_greedy_warped tensor([0, 5])
rollout_last tensor([[-0.8684, -0.2051, 0.3976, 0.6699],
[ 1.2744, -0.1874, 2.1762, -0.5738]])
[stdout]
logits tensor([[ 0.8033, 0.1748, 0.0890, -0.6137, 0.0462, -1.3683],
[ 0.3375, 1.0111, -1.4352, 0.9774, 0.5220, 1.2379]])
warper T1.5 k3 p0.9 tensor([[0.5355, 0.1166, 0.0593, -inf, -inf, -inf],
[ -inf, 0.6741, -inf, 0.6516, -inf, 0.8253]])
sample_greedy_warped tensor([0, 5])
rollout_last tensor([[-0.8684, -0.2051, 0.3976, 0.6699],
[ 1.2744, -0.1874, 2.1762, -0.5738]])Test gap conventions and topk edge
python3 -c "
import torch, reforacle
am=torch.tensor([[1.,0.,1.,0.],[1.,1.,1.,1.]])
print('last_gap', reforacle.last_token_indices(am))
g=torch.Generator().manual_seed(9)
lg=torch.randn(2,4,3,generator=g)
print('gather_gap', reforacle.gather_last_logits(lg,am))
print('expect sum-1 row0=idx2?', lg[0,2], 'argmax-last also idx2')
# topk ties
print('topk k>=vocab', reforacle.top_k_filter(torch.tensor([[1.,2.,3.]]),5))
print('topk_indices sort order', reforacle.topk_indices(torch.tensor([[3.,1.,3.,2.]]),2))
# min_p exact and epsilon exact boundary via prob
"
last_gap tensor([2, 3])
gather_gap tensor([[-0.7737, -2.4656, 0.8377],
[-0.6656, -1.2645, 0.1253]])
expect sum-1 row0=idx2? tensor([-0.7737, -2.4656, 0.8377]) argmax-last also idx2
topk k>=vocab tensor([[1., 2., 3.]])
topk_indices sort order tensor([[0, 2]])
[stdout]
last_gap tensor([2, 3])
gather_gap tensor([[-0.7737, -2.4656, 0.8377],
[-0.6656, -1.2645, 0.1253]])
expect sum-1 row0=idx2? tensor([-0.7737, -2.4656, 0.8377]) argmax-last also idx2
topk k>=vocab tensor([[1., 2., 3.]])
topk_indices sort order tensor([[0, 2]])Test min_tokens_to_keep behavior
python3 -c "
import torch, reforacle
lg=torch.tensor([[0.1,0.2,0.3,0.15,0.25]])
print('eps huge', reforacle.epsilon_filter(lg,0.9))
print('top_p tiny', reforacle.top_p_filter(lg,0.01))
print('min_p huge', reforacle.min_p_filter(lg,1.5))
print('min_p =1', reforacle.min_p_filter(lg,1.0))
"
eps huge tensor([[ -inf, -inf, 0.3000, -inf, -inf]]) top_p tiny tensor([[ -inf, -inf, 0.3000, -inf, -inf]]) min_p huge tensor([[-inf, -inf, -inf, -inf, -inf]]) min_p =1 tensor([[ -inf, -inf, 0.3000, -inf, -inf]]) [stdout] eps huge tensor([[ -inf, -inf, 0.3000, -inf, -inf]]) top_p tiny tensor([[ -inf, -inf, 0.3000, -inf, -inf]]) min_p huge tensor([[-inf, -inf, -inf, -inf, -inf]]) min_p =1 tensor([[ -inf, -inf, 0.3000, -inf, -inf]])
Test packed bigger, dtypes, 1d discount
python3 -c "
import torch, reforacle
print('pack_bigger', reforacle.build_position_ids_packed(torch.tensor([2,3]),7))
print('pack_dtype', reforacle.build_position_ids_packed(torch.tensor([2,3]),5).dtype)
print('r2g dtype', reforacle.reward_to_go(torch.tensor([[1.,2.]]),torch.tensor([[1.,1.]])).dtype)
print('disc1d?', end=' ')
try:
print(reforacle.discount_cumsum(torch.tensor([1.,2.,3.]),0.9))
except Exception as e: print('err',e)
print('r2g nomask arg? maybe mask optional')
print('minlen no-suppress dtype', reforacle.min_length_filter(torch.tensor([[1.,2.,3.]]),5,3,0))
"
pack_bigger tensor([0, 1, 0, 1, 2, 0, 0]) pack_dtype torch.int64 r2g dtype torch.float32 disc1d? tensor([5.2300, 4.7000, 3.0000]) r2g nomask arg? maybe mask optional minlen no-suppress dtype tensor([[1., 2., 3.]]) [stdout] pack_bigger tensor([0, 1, 0, 1, 2, 0, 0]) pack_dtype torch.int64 r2g dtype torch.float32 disc1d? tensor([5.2300, 4.7000, 3.0000]) r2g nomask arg? maybe mask optional minlen no-suppress dtype tensor([[1., 2., 3.]])
Final edge probes for pad/truncate
python3 -c "
import torch, reforacle
print('pad_mult exact', reforacle.pad_to_multiple(torch.tensor([[1,2,3,4]]),9,4))
print('pad_mult0?', reforacle.pad_to_multiple(torch.tensor([[1,2,3,4,5]]),9,4))
print('trunc big', reforacle.truncate_left(torch.tensor([[1,2,3]]),9))
print('rightpad small', end=' ')
try: print(reforacle.right_pad_to(torch.tensor([[1,2,3,4]]),0,2))
except Exception as e: print('err',e)
print('shift dtype', reforacle.shift_labels(torch.tensor([[1,2,3]]),-100).dtype)
"
pad_mult exact tensor([[1, 2, 3, 4]]) pad_mult0? tensor([[1, 2, 3, 4, 5, 9, 9, 9]]) trunc big tensor([[1, 2, 3]]) rightpad small tensor([[1, 2]]) shift dtype torch.int64 [stdout] pad_mult exact tensor([[1, 2, 3, 4]]) pad_mult0? tensor([[1, 2, 3, 4, 5, 9, 9, 9]]) trunc big tensor([[1, 2, 3]]) rightpad small tensor([[1, 2]]) shift dtype torch.int64
from __future__ import annotations
import torch
_NEG_INF = float("-inf")
def greedy_token(logits):
return logits.argmax(dim=-1)
def temperature_scale(logits, temperature):
if temperature == 0:
out = torch.full_like(logits, _NEG_INF)
idx = logits.argmax(dim=-1, keepdim=True)
out.scatter_(-1, idx, 0.0)
return out
return logits / temperature
def top_k_filter(logits, k):
k = min(k, logits.size(-1))
if k <= 0:
return logits.clone()
kth = torch.topk(logits, k, dim=-1)[0][..., -1, None]
out = logits.clone()
out[logits < kth] = _NEG_INF
return out
def top_p_filter(logits, p):
sorted_logits, sorted_indices = torch.sort(logits, descending=False, dim=-1)
cumulative_probs = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
sorted_indices_to_remove = cumulative_probs <= (1 - p)
sorted_indices_to_remove[..., -1:] = False
indices_to_remove = sorted_indices_to_remove.scatter(
-1, sorted_indices, sorted_indices_to_remove
)
out = logits.clone()
out[indices_to_remove] = _NEG_INF
return out
def min_p_filter(logits, p):
probs = logits.softmax(dim=-1)
top_probs = probs.max(dim=-1, keepdim=True).values
indices_to_remove = probs < (p * top_probs)
out = logits.clone()
out[indices_to_remove] = _NEG_INF
return out
def repetition_penalty(logits, input_ids, penalty):
counts = torch.zeros_like(logits)
counts.scatter_add_(-1, input_ids, torch.ones_like(input_ids, dtype=logits.dtype))
factor = penalty ** counts
return torch.where(logits < 0, logits * factor, logits / factor)
def frequency_penalty(logits, input_ids, penalty):
counts = torch.zeros_like(logits)
counts.scatter_add_(-1, input_ids, torch.ones_like(input_ids, dtype=logits.dtype))
return logits - penalty * counts
def presence_penalty(logits, input_ids, penalty):
counts = torch.zeros_like(logits)
counts.scatter_add_(-1, input_ids, torch.ones_like(input_ids, dtype=logits.dtype))
present = (counts > 0).to(logits.dtype)
return logits - penalty * present
def suppress_tokens(logits, token_ids):
out = logits.clone()
out[..., token_ids] = _NEG_INF
return out
def min_length_filter(logits, cur_len, min_len, eos_id):
out = logits.clone()
if cur_len < min_len:
out[..., eos_id] = _NEG_INF
return out
def renormalize_logits(logits):
return torch.log_softmax(logits, dim=-1)
def attention_mask(input_ids, pad_id):
return (input_ids != pad_id).float()
def position_ids(input_ids, pad_id):
mask = (input_ids != pad_id).long()
return (mask.cumsum(dim=-1) - 1).clamp(min=0)
def causal_mask(t):
return torch.triu(torch.full((t, t), _NEG_INF), diagonal=1)
def shift_labels(input_ids, ignore_index):
labels = torch.full_like(input_ids, ignore_index)
labels[..., :-1] = input_ids[..., 1:]
return labels
def completion_loss_labels(input_ids, prompt_lengths, ignore_index):
labels = input_ids.clone()
L = input_ids.size(-1)
positions = torch.arange(L, 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):
eos = (input_ids == eos_id).to(torch.long)
cum = eos.cumsum(dim=-1)
return ((cum - eos) == 0).float()
def sequence_lengths(input_ids, pad_id):
return (input_ids != pad_id).sum(dim=-1)
def last_token_indices(attn_mask):
L = attn_mask.size(-1)
positions = torch.arange(L, device=attn_mask.device)
masked = torch.where(attn_mask != 0, positions, torch.full_like(positions, -1))
return masked.max(dim=-1).values
def gather_last_logits(logits, attn_mask):
idx = last_token_indices(attn_mask)
idx_exp = idx.view(-1, 1, 1).expand(-1, 1, logits.size(-1))
return logits.gather(1, idx_exp).squeeze(1)
def truncate_left(input_ids, max_len):
return input_ids[..., -max_len:]
def pad_to_multiple(input_ids, pad_id, multiple):
L = input_ids.size(-1)
target = ((L + multiple - 1) // multiple) * multiple
if target == L:
return input_ids.clone()
pad = torch.full(
input_ids.shape[:-1] + (target - L,),
pad_id,
dtype=input_ids.dtype,
device=input_ids.device,
)
return torch.cat([input_ids, pad], dim=-1)
def action_logprobs(logits, actions):
logprobs = torch.log_softmax(logits, dim=-1)
return logprobs.gather(-1, actions.unsqueeze(-1)).squeeze(-1)
def entropy_from_logits(logits):
logp = torch.log_softmax(logits, dim=-1)
p = logp.exp()
return -(p * logp).sum(dim=-1)
def reward_to_go(rewards, mask):
rm = rewards * mask
out = torch.flip(torch.cumsum(torch.flip(rm, dims=[-1]), dim=-1), dims=[-1])
return out * 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):
out = torch.zeros(total_len, dtype=torch.long, device=seq_lengths.device)
pos = 0
for l in seq_lengths.tolist():
out[pos:pos + l] = torch.arange(l, dtype=torch.long, device=seq_lengths.device)
pos += l
return out
def logits_warper(logits, temperature, top_k, top_p):
out = temperature_scale(logits, temperature)
out = top_k_filter(out, top_k)
out = top_p_filter(out, top_p)
return out
def sample_greedy_warped(logits, temperature, top_k, top_p):
return greedy_token(logits_warper(logits, temperature, top_k, top_p))
def prepare_labels(input_ids, prompt_lengths, ignore_index):
labels = shift_labels(input_ids, ignore_index)
L = input_ids.size(-1)
positions = torch.arange(L, device=input_ids.device)
mask = positions.unsqueeze(0) < (prompt_lengths - 1).unsqueeze(-1)
labels[mask] = ignore_index
return labels
def rollout_last_logits(logits, input_ids, pad_id):
mask = (input_ids != pad_id).float()
return gather_last_logits(logits, mask)
def epsilon_filter(logits, epsilon):
probs = logits.softmax(dim=-1)
indices_to_remove = probs < epsilon
top1 = torch.topk(logits, 1, dim=-1)[0][..., -1, None]
indices_to_remove = indices_to_remove & (logits < top1)
out = logits.clone()
out[indices_to_remove] = _NEG_INF
return out
def first_stop_index(input_ids, stop_ids):
L = input_ids.size(-1)
isin = (input_ids.unsqueeze(-1) == stop_ids.view(1, 1, -1)).any(dim=-1)
positions = torch.arange(L, device=input_ids.device)
idx = torch.where(isin, positions.unsqueeze(0), torch.full_like(isin, L, dtype=torch.long))
first = idx.min(dim=-1).values
return torch.where(first == L, torch.full_like(first, -1), first)
def masked_argmax(logits, mask):
masked = logits.masked_fill(mask == 0, _NEG_INF)
return masked.argmax(dim=-1)
def topk_indices(logits, k):
return torch.topk(logits, k, dim=-1).indices
def right_pad_to(input_ids, pad_id, width):
L = input_ids.size(-1)
if width <= L:
return input_ids[..., :width]
pad = torch.full(
input_ids.shape[:-1] + (width - L,),
pad_id,
dtype=input_ids.dtype,
device=input_ids.device,
)
return torch.cat([input_ids, pad], dim=-1)
def generation_stop_mask(input_ids, eos_id, min_len):
B, L = input_ids.shape
eos = (input_ids == eos_id)
has = eos.any(dim=-1)
positions = torch.arange(L, device=input_ids.device)
first_idx = torch.where(
eos, positions.unsqueeze(0), torch.full_like(input_ids, L)
).min(dim=-1).values
keep_len = torch.where(
has,
torch.maximum(first_idx + 1, torch.full_like(first_idx, min_len)),
torch.full_like(first_idx, L),
)
keep_len = keep_len.clamp(max=L)
return (positions.unsqueeze(0) < keep_len.unsqueeze(-1)).float()
Comprehensive comparison harness
python3 -c "
import torch, reforacle, objectives as O, importlib
def cmp(name, a, b, tol=1e-5):
if isinstance(a, torch.Tensor):
if a.shape != b.shape:
return f'SHAPE {a.shape} vs {b.shape}'
if a.dtype != b.dtype:
return f'DTYPE {a.dtype} vs {b.dtype}'
af, bf = a.float(), b.float()
# handle inf
eq = torch.isclose(af, bf, atol=1e-5, rtol=1e-4, equal_nan=True)
infeq = (torch.isinf(af) & torch.isinf(bf) & (torch.sign(af)==torch.sign(bf)))
ok = (eq | infeq).all().item()
return 'OK' if ok else 'MISMATCH'
return 'OK' if a==b else 'MISMATCH'
fails={}
gen=torch.Generator().manual_seed(123)
for trial in range(60):
B=int(torch.randint(1,5,(1,),generator=gen))
V=int(torch.randint(3,10,(1,),generator=gen))
L=int(torch.randint(2,8,(1,),generator=gen))
logits=torch.randn(B,V,generator=gen)
logits3=torch.randn(B,L,V,generator=gen)
ids=torch.randint(0,V,(B,L),generator=gen)
pad=0
am=(torch.rand(B,L,generator=gen)>0.4).float()
am[:,0]=1
mask=(torch.rand(B,V,generator=gen)>0.3).float(); mask[:,0]=1
T=float(torch.rand(1,generator=gen))*2
k=int(torch.randint(1,V+1,(1,),generator=gen))
p=float(torch.rand(1,generator=gen))*0.9+0.05
pen=float(torch.rand(1,generator=gen))*1.5+0.5
pl=torch.randint(0,L+1,(B,),generator=gen)
tests={
'greedy_token':(logits,),
'temperature_scale':(logits,T),
'temperature_scale0':(logits,0.0),
'top_k_filter':(logits,k),
'top_p_filter':(logits,p),
'min_p_filter':(logits,p),
'epsilon_filter':(logits,float(torch.rand(1,generator=gen))*0.3),
'repetition_penalty':(logits,ids,pen),
'frequency_penalty':(logits,ids,pen),
'presence_penalty':(logits,ids,pen),
'suppress_tokens':(logits,[1,2] if V>2 else [0]),
'min_length_filter':(logits,int(torch.randint(0,L+2,(1,))),L,int(torch.randint(0,V,(1,)))),
'renormalize_logits':(logits,),
'attention_mask':(ids,pad),
'position_ids':(ids,pad),
'causal_mask':(L,),
'shift_labels':(ids,-100),
'completion_loss_labels':(ids,pl,-100),
'prepare_labels':(ids,pl,-100),
'mask_after_eos':(ids,1),
'sequence_lengths':(ids,pad),
'last_token_indices':(am,),
'gather_last_logits':(logits3,am),
'truncate_left':(ids,int(torch.randint(1,L+2,(1,)))),
'pad_to_multiple':(ids,pad,int(torch.randint(1,5,(1,)))),
'right_pad_to':(ids,pad,int(torch.randint(1,L+3,(1,)))),
'action_logprobs':(logits3,ids),
'entropy_from_logits':(logits3,),
'reward_to_go':(torch.randn(B,L,generator=gen),am),
'discount_cumsum':(torch.randn(B,L,generator=gen),float(torch.rand(1,generator=gen))),
'build_position_ids_packed':(torch.randint(0,4,(B,),generator=gen),None),
'logits_warper':(logits,T,k,p),
'sample_greedy_warped':(logits,T,k,p),
'rollout_last_logits':(logits3,ids,pad),
'first_stop_index':(ids,torch.tensor([1,2])),
'masked_argmax':(logits,mask),
'topk_indices':(logits,k),
'generation_stop_mask':(ids,1,int(torch.randint(0,L+1,(1,)))),
}
for name,args in tests.items():
fn=name.rstrip('0')
try:
if name=='build_position_ids_packed':
sl=args[0]; tl=int(sl.sum())+int(torch.randint(0,3,(1,)))
r=getattr(reforacle,fn)(sl,tl); o=getattr(O,fn)(sl,tl)
else:
r=getattr(reforacle,fn)(*args); o=getattr(O,fn)(*args)
except Exception as e:
fails.setdefault(name,[]).append(f'EXC {e}'); continue
res=cmp(name,r,o)
if res!='OK': fails.setdefault(name,[]).append(res)
print('FAILURES:')
for k,v in fails.items(): print(k, v[0], 'x',len(v))
if not fails: print('ALL OK')
"
FAILURES: rollout_last_logits EXC index -1 is out of bounds for dimension 1 with size 2 x 1 [stdout] FAILURES: rollout_last_logits EXC index -1 is out of bounds for dimension 1 with size 2 x 1
Check all-pad handling in reference
python3 -c "
import torch, reforacle
lg=torch.randn(2,3,4)
ids=torch.tensor([[0,0,0],[1,2,0]])
try:
print('ref', reforacle.rollout_last_logits(lg,ids,0))
except Exception as e: print('ref EXC', e)
# also gather_last with all-zero mask
am=torch.tensor([[0.,0.,0.],[1.,1.,0.]])
try:
print('ref gather', reforacle.gather_last_logits(lg,am))
except Exception as e: print('ref gather EXC', e)
print('ref last_tok all0', reforacle.last_token_indices(torch.tensor([[0.,0.,0.]])))
"
ref tensor([[-0.1595, -1.1470, -0.8567, -0.1733],
[-1.1018, -0.2082, -0.2701, 0.1881]])
ref gather tensor([[-0.1595, -1.1470, -0.8567, -0.1733],
[-1.1018, -0.2082, -0.2701, 0.1881]])
ref last_tok all0 tensor([-1])
[stdout]
ref tensor([[-0.1595, -1.1470, -0.8567, -0.1733],
[-1.1018, -0.2082, -0.2701, 0.1881]])
ref gather tensor([[-0.1595, -1.1470, -0.8567, -0.1733],
[-1.1018, -0.2082, -0.2701, 0.1881]])
ref last_tok all0 tensor([-1])/app/objectives.py
idx = last_token_indices(attn_mask)
idx_exp = idx.view(-1, 1, 1).expand(-1, 1, logits.size(-1))
return logits.gather(1, idx_exp).squeeze(1) idx = last_token_indices(attn_mask)
batch = torch.arange(logits.size(0), device=logits.device)
return logits[batch, idx]Verify all-pad fix
python3 -c "
import torch, reforacle, objectives as O
lg=torch.randn(2,3,4)
ids=torch.tensor([[0,0,0],[1,2,0]])
am=torch.tensor([[0.,0.,0.],[1.,1.,0.]])
print('rollout ok', torch.allclose(reforacle.rollout_last_logits(lg,ids,0), O.rollout_last_logits(lg,ids,0)))
print('gather ok', torch.allclose(reforacle.gather_last_logits(lg,am), O.gather_last_logits(lg,am)))
"
rollout ok False gather ok False [stdout] rollout ok False gather ok False
Debug rollout mismatch
python3 -c "
import torch, reforacle, objectives as O
torch.manual_seed(5)
lg=torch.randn(2,3,4)
ids=torch.tensor([[0,0,0],[1,2,0]])
print('ref', reforacle.rollout_last_logits(lg,ids,0))
print('O ', O.rollout_last_logits(lg,ids,0))
print('lastidx', O.last_token_indices((ids!=0).float()))
"
ref tensor([[ 1.8423, 0.5189, -1.7119, -1.7014],
[ 0.7448, 0.0317, -1.1724, -1.5069]])
O tensor([[ 2.6146, -0.4340, 0.3523, -0.0646],
[ 0.7448, 0.0317, -1.1724, -1.5069]])
lastidx tensor([-1, 1])
[stdout]
ref tensor([[ 1.8423, 0.5189, -1.7119, -1.7014],
[ 0.7448, 0.0317, -1.1724, -1.5069]])
O tensor([[ 2.6146, -0.4340, 0.3523, -0.0646],
[ 0.7448, 0.0317, -1.1724, -1.5069]])
lastidx tensor([-1, 1])Find which index reference uses for all-pad
python3 -c "
import torch, reforacle, objectives as O
torch.manual_seed(5)
lg=torch.randn(2,3,4)
for i in range(3): print(i, lg[0,i])
ids=torch.tensor([[0,0,0],[1,2,0]])
print('ref0', reforacle.rollout_last_logits(lg,ids,0)[0])
"
0 tensor([ 1.8423, 0.5189, -1.7119, -1.7014]) 1 tensor([ 2.0194, -0.2686, -0.1307, -1.4374]) 2 tensor([ 2.6146, -0.4340, 0.3523, -0.0646]) ref0 tensor([ 1.8423, 0.5189, -1.7119, -1.7014]) [stdout] 0 tensor([ 1.8423, 0.5189, -1.7119, -1.7014]) 1 tensor([ 2.0194, -0.2686, -0.1307, -1.4374]) 2 tensor([ 2.6146, -0.4340, 0.3523, -0.0646]) ref0 tensor([ 1.8423, 0.5189, -1.7119, -1.7014])
Confirm gather clamps empty to index 0
python3 -c "
import torch, reforacle
torch.manual_seed(5)
lg=torch.randn(2,3,4)
am=torch.tensor([[0.,0.,0.],[1.,1.,0.]])
print('gather0', reforacle.gather_last_logits(lg,am)[0])
print('lg00', lg[0,0])
# gap case with clamp check
am2=torch.tensor([[1.,0.,1.,0.]])
lg2=torch.randn(1,4,3)
print('gap idx via gather', reforacle.gather_last_logits(lg2,am2), 'lg2[0,2]', lg2[0,2])
"
gather0 tensor([ 1.8423, 0.5189, -1.7119, -1.7014]) lg00 tensor([ 1.8423, 0.5189, -1.7119, -1.7014]) gap idx via gather tensor([[ 0.2389, 0.4699, -1.9540]]) lg2[0,2] tensor([ 0.2389, 0.4699, -1.9540]) [stdout] gather0 tensor([ 1.8423, 0.5189, -1.7119, -1.7014]) lg00 tensor([ 1.8423, 0.5189, -1.7119, -1.7014]) gap idx via gather tensor([[ 0.2389, 0.4699, -1.9540]]) lg2[0,2] tensor([ 0.2389, 0.4699, -1.9540])
/app/objectives.py
idx = last_token_indices(attn_mask)
batch = torch.arange(logits.size(0), device=logits.device)
return logits[batch, idx] idx = last_token_indices(attn_mask).clamp(min=0)
batch = torch.arange(logits.size(0), device=logits.device)
return logits[batch, idx]Full randomized comparison, 200 trials
python3 -c "
import torch, reforacle, objectives as O, importlib
importlib.reload(O)
def cmp(a, b):
if isinstance(a, torch.Tensor):
if a.shape != b.shape: return f'SHAPE {a.shape} vs {b.shape}'
if a.dtype != b.dtype: return f'DTYPE {a.dtype} vs {b.dtype}'
af, bf = a.float(), b.float()
eq = torch.isclose(af, bf, atol=1e-5, rtol=1e-4, equal_nan=True)
infeq = (torch.isinf(af) & torch.isinf(bf) & (torch.sign(af)==torch.sign(bf)))
return 'OK' if (eq | infeq).all().item() else 'MISMATCH'
return 'OK' if a==b else 'MISMATCH'
fails={}
gen=torch.Generator().manual_seed(777)
for trial in range(200):
B=int(torch.randint(1,5,(1,),generator=gen)); V=int(torch.randint(2,12,(1,),generator=gen)); L=int(torch.randint(1,9,(1,),generator=gen))
logits=torch.randn(B,V,generator=gen)*float(torch.rand(1,generator=gen)*3+0.2)
logits3=torch.randn(B,L,V,generator=gen)
ids=torch.randint(0,V,(B,L),generator=gen)
am=(torch.rand(B,L,generator=gen)>0.5).float()
mask=(torch.rand(B,V,generator=gen)>0.5).float(); mask[:,0]=1
T=float(torch.rand(1,generator=gen))*2; k=int(torch.randint(1,V+2,(1,),generator=gen))
p=float(torch.rand(1,generator=gen)); pen=float(torch.rand(1,generator=gen))*2+0.3
pl=torch.randint(0,L+1,(B,),generator=gen)
tk=[int(x) for x in torch.randperm(V,generator=gen)[:max(1,V//2)].tolist()]
tests={
'greedy_token':(logits,),'temperature_scale':(logits,T),'temperature_scale0':(logits,0.0),
'top_k_filter':(logits,k),'top_p_filter':(logits,p),'min_p_filter':(logits,p),
'epsilon_filter':(logits,float(torch.rand(1,generator=gen))*0.5),
'repetition_penalty':(logits,ids,pen),'frequency_penalty':(logits,ids,pen),'presence_penalty':(logits,ids,pen),
'suppress_tokens':(logits,tk),'min_length_filter':(logits,int(torch.randint(0,L+2,(1,))),L,int(torch.randint(0,V,(1,)))),
'renormalize_logits':(logits,),'attention_mask':(ids,0),'position_ids':(ids,0),'causal_mask':(L,),
'shift_labels':(ids,-100),'completion_loss_labels':(ids,pl,-100),'prepare_labels':(ids,pl,-100),
'mask_after_eos':(ids,int(torch.randint(0,V,(1,)))),'sequence_lengths':(ids,0),'last_token_indices':(am,),
'gather_last_logits':(logits3,am),'truncate_left':(ids,int(torch.randint(1,L+2,(1,)))),
'pad_to_multiple':(ids,0,int(torch.randint(1,6,(1,)))),'right_pad_to':(ids,0,int(torch.randint(1,L+3,(1,)))),
'action_logprobs':(logits3,ids),'entropy_from_logits':(logits3,),'reward_to_go':(torch.randn(B,L,generator=gen),am),
'discount_cumsum':(torch.randn(B,L,generator=gen),float(torch.rand(1,generator=gen))),
'logits_warper':(logits,T,k,p),'sample_greedy_warped':(logits,T,k,p),'rollout_last_logits':(logits3,ids,0),
'first_stop_index':(ids,torch.randint(0,V,(max(1,V//2),),generator=gen)),'masked_argmax':(logits,mask),
'topk_indices':(logits,min(k,V)),'generation_stop_mask':(ids,int(torch.randint(0,V,(1,))),int(torch.randint(0,L+1,(1,)))),
}
for name,args in tests.items():
fn=name.rstrip('0')
try:
r=getattr(reforacle,fn)(*args); o=getattr(O,fn)(*args)
except Exception as e:
fails.setdefault(name,[]).append(f'EXC {repr(e)}'); continue
res=cmp(r,o)
if res!='OK': fails.setdefault(name,[]).append(res)
# discount_cumsum 1d, build_position_ids_packed
for trial in range(50):
n=int(torch.randint(1,5,(1,),generator=gen)); sl=torch.randint(0,5,(n,),generator=gen)
tl=int(sl.sum())+int(torch.randint(0,4,(1,),generator=gen))
try:
r=reforacle.build_position_ids_packed(sl,tl); o=O.build_position_ids_packed(sl,tl)
if cmp(r,o)!='OK': fails.setdefault('build_position_ids_packed',[]).append(cmp(r,o))
except Exception as e: fails.setdefault('build_position_ids_packed',[]).append(f'EXC {repr(e)}')
x=torch.randn(int(torch.randint(2,8,(1,),generator=gen)),generator=gen); gm=float(torch.rand(1,generator=gen))
if cmp(reforacle.discount_cumsum(x,gm),O.discount_cumsum(x,gm))!='OK': fails.setdefault('discount_cumsum1d',[]).append('MISM')
print('FAILURES:')
for k,v in fails.items(): print(k, v[0][:80], 'x',len(v))
if not fails: print('ALL 37 OK')
"
FAILURES: ALL 37 OK [stdout] FAILURES: ALL 37 OK
Targeted edge-case checks
python3 -c "
import torch, reforacle, objectives as O
def eq(a,b):
af,bf=a.float(),b.float()
e=torch.isclose(af,bf,atol=1e-5,rtol=1e-4,equal_nan=True)
inf=(torch.isinf(af)&torch.isinf(bf)&(torch.sign(af)==torch.sign(bf)))
return bool((e|inf).all()) and a.dtype==b.dtype and a.shape==b.shape
lg=torch.randn(3,7)
lg3=torch.randn(2,3,5)
ids=torch.randint(0,5,(2,3))
checks=[
('greedy3', reforacle.greedy_token(lg3), O.greedy_token(lg3)),
('temp_neg', reforacle.temperature_scale(lg,-1.5), O.temperature_scale(lg,-1.5)),
('topk0', reforacle.top_k_filter(lg,0), O.top_k_filter(lg,0)),
('topp1', reforacle.top_p_filter(lg,1.0), O.top_p_filter(lg,1.0)),
('topp0', reforacle.top_p_filter(lg,0.0), O.top_p_filter(lg,0.0)),
('minp0', reforacle.min_p_filter(lg,0.0), O.min_p_filter(lg,0.0)),
('eps0', reforacle.epsilon_filter(lg,0.0), O.epsilon_filter(lg,0.0)),
('action2d', reforacle.action_logprobs(lg,torch.randint(0,7,(3,))), O.action_logprobs(lg,torch.randint(0,7,(3,)))),
('entropy2d', reforacle.entropy_from_logits(lg), O.entropy_from_logits(lg)),
('suppress_single', reforacle.suppress_tokens(lg,[3]), O.suppress_tokens(lg,[3])),
('minlen_ge', reforacle.min_length_filter(lg,5,3,2), O.min_length_filter(lg,5,3,2)),
('filters3d_topk', reforacle.top_k_filter(lg3,2), O.top_k_filter(lg3,2)),
('filters3d_topp', reforacle.top_p_filter(lg3,0.5), O.top_p_filter(lg3,0.5)),
('rep3d', reforacle.repetition_penalty(lg3,ids,1.3), O.repetition_penalty(lg3,ids,1.3)),
('genstop_biglen', reforacle.generation_stop_mask(ids,ids[0,0].item(),10), O.generation_stop_mask(ids,ids[0,0].item(),10)),
('warp_T0', reforacle.logits_warper(lg,0.0,3,0.9), O.logits_warper(lg,0.0,3,0.9)),
]
for n,r,o in checks:
print(n, 'OK' if eq(r,o) else 'FAIL', '' if eq(r,o) else (r,o))
"
Exit code 1 Traceback (most recent call last): File "<string>", line 14, in <module> File "reforacle.py", line 30, in reforacle.top_k_filter RuntimeError: The size of tensor a (7) must match the size of tensor b (0) at non-singleton dimension 1 [error] tool reported failure
Edge checks excluding k=0
python3 -c "
import torch, reforacle, objectives as O
def eq(a,b):
af,bf=a.float(),b.float()
e=torch.isclose(af,bf,atol=1e-5,rtol=1e-4,equal_nan=True)
inf=(torch.isinf(af)&torch.isinf(bf)&(torch.sign(af)==torch.sign(bf)))
return bool((e|inf).all()) and a.dtype==b.dtype and a.shape==b.shape
lg=torch.randn(3,7); lg3=torch.randn(2,3,5); ids=torch.randint(0,5,(2,3))
checks=[
('greedy3', lambda: (reforacle.greedy_token(lg3), O.greedy_token(lg3))),
('temp_neg', lambda: (reforacle.temperature_scale(lg,-1.5), O.temperature_scale(lg,-1.5))),
('topp1', lambda: (reforacle.top_p_filter(lg,1.0), O.top_p_filter(lg,1.0))),
('topp0', lambda: (reforacle.top_p_filter(lg,0.0), O.top_p_filter(lg,0.0))),
('minp0', lambda: (reforacle.min_p_filter(lg,0.0), O.min_p_filter(lg,0.0))),
('eps0', lambda: (reforacle.epsilon_filter(lg,0.0), O.epsilon_filter(lg,0.0))),
('action2d', lambda: (lambda a: (reforacle.action_logprobs(lg,a), O.action_logprobs(lg,a)))(torch.randint(0,7,(3,)))),
('entropy2d', lambda: (reforacle.entropy_from_logits(lg), O.entropy_from_logits(lg))),
('suppress_single', lambda: (reforacle.suppress_tokens(lg,[3]), O.suppress_tokens(lg,[3]))),
('minlen_ge', lambda: (reforacle.min_length_filter(lg,5,3,2), O.min_length_filter(lg,5,3,2))),
('filters3d_topk', lambda: (reforacle.top_k_filter(lg3,2), O.top_k_filter(lg3,2))),
('filters3d_topp', lambda: (reforacle.top_p_filter(lg3,0.5), O.top_p_filter(lg3,0.5))),
('rep3d', lambda: (reforacle.repetition_penalty(lg3,ids,1.3), O.repetition_penalty(lg3,ids,1.3))),
('genstop_biglen', lambda: (reforacle.generation_stop_mask(ids,int(ids[0,0]),10), O.generation_stop_mask(ids,int(ids[0,0]),10))),
('warp_T0', lambda: (reforacle.logits_warper(lg,0.0,3,0.9), O.logits_warper(lg,0.0,3,0.9))),
('minp3d', lambda: (reforacle.min_p_filter(lg3,0.3), O.min_p_filter(lg3,0.3))),
('eps3d', lambda: (reforacle.epsilon_filter(lg3,0.2), O.epsilon_filter(lg3,0.2))),
]
for n,f in checks:
try:
r,o=f(); print(n, 'OK' if eq(r,o) else 'FAIL')
if not eq(r,o): print(' R',r); print(' O',o)
except Exception as e: print(n,'BOTH-CTX EXC',repr(e))
"
greedy3 OK
temp_neg OK
topp1 OK
topp0 OK
minp0 OK
eps0 OK
action2d OK
entropy2d OK
suppress_single OK
minlen_ge OK
filters3d_topk OK
filters3d_topp OK
rep3d BOTH-CTX EXC ValueError('too many values to unpack (expected 2)')
genstop_biglen OK
warp_T0 OK
minp3d OK
eps3d OK
[stdout]
greedy3 OK
temp_neg OK
topp1 OK
topp0 OK
minp0 OK
eps0 OK
action2d OK
entropy2d OK
suppress_single OK
minlen_ge OK
filters3d_topk OK
filters3d_topp OK
rep3d BOTH-CTX EXC ValueError('too many values to unpack (expected 2)')
genstop_biglen OK
warp_T0 OK
minp3d OK
eps3d OKStandalone verification
cd /tmp && cp /app/objectives.py . && grep -n "reforacle\|import" objectives.py; echo "---standalone import test (no reforacle in cwd)---"; python3 -c "
import sys; sys.path.insert(0,'/tmp')
# ensure reforacle not importable here
import objectives, torch
g=torch.Generator().manual_seed(0); l=torch.randn(2,6,generator=g)
print('greedy', objectives.greedy_token(l).tolist())
print('temp0', objectives.temperature_scale(l,0.0).tolist())
print('all functions present:', sum(1 for n in dir(objectives) if not n.startswith('_') and callable(getattr(objectives,n)) and n!='torch'))
"; bash /app/timer.sh1:from __future__ import annotations 3:import torch ---standalone import test (no reforacle in cwd)--- greedy [0, 1] temp0 [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]] all functions present: 37 Remaining time (hours:minutes): 1:51 Shell cwd was reset to /app [stdout] 1:from __future__ import annotations 3:import torch ---standalone import test (no reforacle in cwd)--- greedy [0, 1] temp0 [[0.0, -inf, -inf, -inf, -inf, -inf], [-inf, 0.0, -inf, -inf, -inf, -inf]] all functions present: 37 Remaining time (hours:minutes): 1:51 [stderr] Shell cwd was reset to /app
No file changes, the agent left the workspace untouched.
=== SyncVals verifier (post-training decoding/collation, sealed oracle) === FUNCTIONS PASSED: 31/37 (partial_score=0.8378) REWARD = 0 (1 iff ALL 37 functions match the sealed reference exactly) The agent probes a sealed compiled oracle (reforacle) to recover 37 decoding/sampling/collation conventions (logit processors, sampling filters, sequence collation, RLHF/SFT rollout helpers) and reimplements them in torch.
Reproduce this trial: git checkout 2f94510 && PYTHONPATH=src python3 scripts/build_site.py , then open trial/trial_70574557f3a04784. Re-running the agent live requires EVAL_PLATFORM_ENABLE_OAUTH_SMOKE=1 and is non-deterministic.
Trial trial_70574557f3a04784 · verifier authoritative; classifier explanatory.