File size: 2,486 Bytes
2ef407e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | import torch
from trident.config import TridentConfig
from trident.patching import entropy_budget_patch
def test_patch_constraints():
cfg = TridentConfig.tiny_test()
B, L = 4, 60
torch.manual_seed(0)
entropy = torch.rand(B, L) * 4.0
valid = torch.ones(B, L, dtype=torch.bool)
# size the patch budget so there is no overflow (P >= L worst case)
out = entropy_budget_patch(entropy, valid, cfg.b_max, cfg.tau_b, max_patches=L + 1)
seg = out["seg_ids"]
# segment ids are non-decreasing along time and start at 0
assert (seg[:, 0] == 0).all()
assert (seg[:, 1:] - seg[:, :-1] >= 0).all()
assert (seg[:, 1:] - seg[:, :-1] <= 1).all() # increments by at most 1
# no patch exceeds b_max bytes
assert (out["patch_len"] <= cfg.b_max).all()
# prev_patch = seg - 1
assert torch.equal(out["prev_patch"], seg - 1)
# determinism
out2 = entropy_budget_patch(entropy, valid, cfg.b_max, cfg.tau_b, max_patches=L + 1)
assert torch.equal(seg, out2["seg_ids"])
def test_bmax_forces_boundary():
cfg = TridentConfig.tiny_test(tau_b=1e9) # never trigger entropy boundary
B, L = 1, 40
entropy = torch.zeros(B, L)
valid = torch.ones(B, L, dtype=torch.bool)
out = entropy_budget_patch(entropy, valid, cfg.b_max, cfg.tau_b, cfg.max_patches)
# with no entropy trigger, boundaries occur exactly every b_max bytes
assert (out["patch_len"][out["patch_valid"]] <= cfg.b_max).all()
assert out["patch_len"][0, 0] == cfg.b_max
def test_low_threshold_many_patches():
cfg = TridentConfig.tiny_test(tau_b=0.1)
B, L = 1, 30
entropy = torch.ones(B, L)
valid = torch.ones(B, L, dtype=torch.bool)
out = entropy_budget_patch(entropy, valid, cfg.b_max, cfg.tau_b, max_patches=L + 1)
# every byte should be its own patch (budget exceeded immediately)
assert out["num_patches"][0].item() == L
def test_overflow_merges_into_last_patch():
"""When the patch budget is exceeded, the remainder merges into the last
slot rather than corrupting indices."""
cfg = TridentConfig.tiny_test(tau_b=0.1)
B, L, P = 1, 30, 8
entropy = torch.ones(B, L)
valid = torch.ones(B, L, dtype=torch.bool)
out = entropy_budget_patch(entropy, valid, cfg.b_max, cfg.tau_b, max_patches=P)
assert out["num_patches"][0].item() == P
assert out["seg_ids"].max().item() == P - 1
# all bytes still assigned; total length preserved
assert out["patch_len"].sum().item() == L
|