| import unittest |
|
|
| import torch |
| from torch import nn |
|
|
| from predictor_training.single_block import SingleBlockPredictor |
|
|
|
|
| class EchoBlock(nn.Module): |
| def forward(self, value: torch.Tensor, **kwargs) -> torch.Tensor: |
| return value |
|
|
|
|
| def build_model(input_variant: str) -> SingleBlockPredictor: |
| model = SingleBlockPredictor( |
| EchoBlock(), |
| dim=8, |
| gradient_checkpointing=False, |
| input_variant=input_variant, |
| ).eval() |
| with torch.no_grad(): |
| model.residual_out.weight.copy_(torch.eye(8)) |
| return model |
|
|
|
|
| def inputs() -> dict[str, object]: |
| return { |
| "current_tokens": torch.randn(2, 3, 8), |
| "anchor_hidden": torch.randn(2, 3, 8), |
| "timestep_modulation": torch.randn(2, 6, 8), |
| "grid_sizes": torch.tensor([[1, 1, 3], [1, 1, 3]]), |
| "freqs": torch.empty(0), |
| "history_k": torch.randn(2, 2, 1, 8), |
| "history_v": torch.randn(2, 2, 1, 8), |
| "cross_k": torch.randn(2, 1, 1, 8), |
| "cross_v": torch.randn(2, 1, 1, 8), |
| "current_start": 2, |
| } |
|
|
|
|
| class DiscaInputVariantTest(unittest.TestCase): |
| def test_disca_output_ignores_previous_chunk_hidden(self) -> None: |
| torch.manual_seed(0) |
| model = build_model(input_variant="disca") |
| common = inputs() |
| first = model(previous_hidden=torch.randn(2, 3, 8), **common) |
| second = model(previous_hidden=torch.randn(2, 3, 8), **common) |
| self.assertTrue(torch.equal(first, second)) |
|
|
| def test_disca_physically_removes_previous_chunk_channel(self) -> None: |
| baseline = build_model(input_variant="self_forcing") |
| disca = build_model(input_variant="disca") |
| self.assertEqual(baseline.fusion.proj_in.weight.shape, (16, 24)) |
| self.assertEqual(disca.fusion.proj_in.weight.shape, (16, 16)) |
| self.assertTrue(hasattr(baseline.fusion, "previous_norm")) |
| self.assertFalse(hasattr(disca.fusion, "previous_norm")) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
