File size: 3,767 Bytes
c34ff1f | 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 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 | """Unit tests for P-EAGLE COD sampling with max_anchors."""
import torch
from speculators.models.peagle.data import generate_cod_sample_indices
def _loss_mask(seq_length: int) -> torch.Tensor:
return torch.ones(1, seq_length, dtype=torch.float32)
class TestMaxAnchors:
def test_depth0_is_full_sequence(self):
"""Depth 0 should always be the full sequence regardless of max_anchors."""
seq_len = 32
anchor_pos, depth = generate_cod_sample_indices(
seq_length=seq_len,
loss_mask=_loss_mask(seq_len),
num_depths=4,
max_anchors=4,
)
depth0_positions = anchor_pos[depth == 0]
assert depth0_positions.shape[0] == seq_len
assert torch.equal(depth0_positions, torch.arange(seq_len))
def test_max_anchors_caps_chains(self):
"""With max_anchors, depth-1+ chains should not exceed max_anchors."""
seq_len = 64
max_anchors = 8
anchor_pos, depth = generate_cod_sample_indices(
seq_length=seq_len,
loss_mask=_loss_mask(seq_len),
num_depths=4,
max_anchors=max_anchors,
)
for d in range(1, 4):
assert (depth == d).sum().item() <= max_anchors
def test_max_anchors_preserves_full_depth0(self):
"""Depth 0 count should equal seq_length even with small max_anchors."""
seq_len = 128
anchor_pos, depth = generate_cod_sample_indices(
seq_length=seq_len,
loss_mask=_loss_mask(seq_len),
num_depths=8,
max_anchors=4,
)
assert (depth == 0).sum().item() == seq_len
def test_max_anchors_none_uses_all(self):
"""max_anchors=None should use all valid positions (default behavior)."""
seq_len = 32
torch.manual_seed(42)
anchor_pos_none, depth_none = generate_cod_sample_indices(
seq_length=seq_len,
loss_mask=_loss_mask(seq_len),
num_depths=4,
max_anchors=None,
)
torch.manual_seed(42)
anchor_pos_default, depth_default = generate_cod_sample_indices(
seq_length=seq_len,
loss_mask=_loss_mask(seq_len),
num_depths=4,
)
assert torch.equal(anchor_pos_none, anchor_pos_default)
assert torch.equal(depth_none, depth_default)
def test_max_anchors_fewer_valid_than_cap(self):
"""When valid positions < max_anchors, all valid positions are used."""
seq_len = 16
loss_mask = torch.zeros(1, seq_len)
loss_mask[0, :5] = 1
torch.manual_seed(42)
anchor_pos_capped, depth_capped = generate_cod_sample_indices(
seq_length=seq_len,
loss_mask=loss_mask,
num_depths=4,
max_anchors=100,
)
torch.manual_seed(42)
anchor_pos_none, depth_none = generate_cod_sample_indices(
seq_length=seq_len,
loss_mask=loss_mask,
num_depths=4,
max_anchors=None,
)
assert (depth_capped == 0).sum().item() == seq_len
assert torch.equal(anchor_pos_capped, anchor_pos_none)
assert torch.equal(depth_capped, depth_none)
def test_max_anchors_sorted_order(self):
"""Subsampled anchors should be in sorted order for causal masking."""
seq_len = 64
anchor_pos, depth = generate_cod_sample_indices(
seq_length=seq_len,
loss_mask=_loss_mask(seq_len),
num_depths=4,
max_anchors=8,
)
for d in range(1, 4):
d_anchors = anchor_pos[depth == d]
if d_anchors.shape[0] > 1:
assert torch.all(d_anchors[1:] >= d_anchors[:-1])
|