tests/test_patching.py

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