seeker
64 lines
1.9 KiB
Python
64 lines
1.9 KiB
Python
import torch
|
|
|
|
from pi07mini.memory import MEMConfig, MEMVideoEncoder
|
|
|
|
|
|
def tiny_encoder() -> MEMVideoEncoder:
|
|
return MEMVideoEncoder(
|
|
MEMConfig(
|
|
image_size=16,
|
|
patch_size=8,
|
|
d_model=16,
|
|
n_heads=4,
|
|
n_layers=2,
|
|
max_history_frames=6,
|
|
)
|
|
).eval()
|
|
|
|
|
|
def test_mem_output_token_count_is_independent_of_history_length() -> None:
|
|
encoder = tiny_encoder()
|
|
current = torch.randn(2, 3, 16, 16)
|
|
with torch.no_grad():
|
|
no_history = encoder(current)
|
|
long_history = encoder(current, torch.randn(2, 6, 3, 16, 16))
|
|
assert no_history.shape == (2, 4, 16)
|
|
assert long_history.shape == no_history.shape
|
|
|
|
|
|
def test_mem_is_sensitive_to_temporal_order() -> None:
|
|
torch.manual_seed(7)
|
|
encoder = tiny_encoder()
|
|
current = torch.randn(1, 3, 16, 16)
|
|
history = torch.stack(
|
|
(
|
|
torch.zeros(3, 16, 16),
|
|
torch.ones(3, 16, 16),
|
|
torch.full((3, 16, 16), 2.0),
|
|
)
|
|
)[None]
|
|
with torch.no_grad():
|
|
chronological = encoder(current, history)
|
|
reversed_history = encoder(current, history.flip(1))
|
|
assert not torch.allclose(chronological, reversed_history)
|
|
|
|
|
|
def test_masked_history_frame_cannot_change_memory() -> None:
|
|
torch.manual_seed(9)
|
|
encoder = tiny_encoder()
|
|
current = torch.randn(1, 3, 16, 16)
|
|
history = torch.randn(1, 2, 3, 16, 16)
|
|
changed = history.clone()
|
|
changed[:, 0] = 1000.0
|
|
mask = torch.tensor([[False, True]])
|
|
with torch.no_grad():
|
|
first = encoder(current, history, mask)
|
|
second = encoder(current, changed, mask)
|
|
torch.testing.assert_close(first, second)
|
|
|
|
|
|
def test_current_frame_temporal_position_is_exactly_zero() -> None:
|
|
encoding = MEMVideoEncoder.temporal_position(torch.tensor([3, 2, 1, 0]), 15)
|
|
assert torch.count_nonzero(encoding[-1]) == 0
|
|
|