[content] tests/test_tokenizer_edge_cases.py (python)
import unittest
import jax
import numpy as np

from models.tokenizer import TokenizerVQVAE
from utils.nn import VectorQuantizer, STBlock


class TokenizerEdgeCaseTest(unittest.TestCase):
    """Test edge cases and potential sources of non-reproducibility."""

    def setUp(self):
        super().setUp()
        self.seed = 42
        self.rng = jax.random.PRNGKey(self.seed)
        
        self.batch_size = 2
        self.seq_len = 8
        self.image_height = 64
        self.image_width = 64
        self.image_channels = 3
        self.patch_size = 4
        self.model_dim = 128
        self.latent_dim = 32
        self.num_latents = 256
        self.num_blocks = 2
        self.num_heads = 4
        self.dropout = 0.1
        self.codebook_dropout = 0.01

    def test_vector_quantizer_codebook_initialization_reproducibility(self):
        """Test that VectorQuantizer codebook initialization is reproducible."""
        latent_dim = 32
        num_latents = 256
        dropout = 0.01
        
        codebooks = []
        for i in range(3):
            rng = jax.random.PRNGKey(self.seed + i)
            vq = VectorQuantizer(latent_dim, num_latents, dropout)
            x = jax.random.uniform(rng, (100, latent_dim))
            params = vq.init(rng, x, training=True)
            codebook = params["params"]["codebook"]
            codebooks.append(codebook)
        
        for i in range(1, len(codebooks)):
            np.testing.assert_array_equal(codebooks[0], codebooks[i])

    def test_st_block_forward_pass_reproducibility(self):
        """Test that STBlock forward pass is reproducible across multiple calls."""
        dim = 64
        num_heads = 4
        dropout = 0.1
        
        block = STBlock(dim, num_heads, dropout)
        
        rng = jax.random.PRNGKey(self.seed)
        x = jax.random.uniform(rng, (2, 3, 16, dim))
        params = block.init(rng, x)
        
        rng_dropout = jax.random.PRNGKey(123)
        outputs = []
        for _ in range(3):
            output = block.apply(params, x, rngs={"dropout": rng_dropout})
            outputs.append(output)
        
        for i in range(1, len(outputs)):
            np.testing.assert_array_equal(outputs[0], outputs[i])

    def test_tokenizer_gradient_reproducibility(self):
        """Test that gradients are reproducible."""
        class MockArgs:
            def __init__(self):
                self.vq_beta = 0.25
                self.num_latents = 64
                self.log_gradients = False
        
        args = MockArgs()
        
        from train_tokenizer import create_training_functions
        tokenizer_loss_fn, _ = create_training_functions(args)
        
        tokenizer = TokenizerVQVAE(
            in_dim=self.image_channels,
            model_dim=self.model_dim,
            latent_dim=self.latent_dim,
            num_latents=self.num_latents,
            patch_size=self.patch_size,
            num_blocks=self.num_blocks,
            num_heads=self.num_heads,
            dropout=self.dropout,
            codebook_dropout=self.codebook_dropout,
        )
        
        rng = jax.random.PRNGKey(self.seed)
        test_videos = jax.random.uniform(rng, (self.batch_size, self.seq_len, self.image_height, self.image_width, self.image_channels))
        params = tokenizer.init(rng, {"videos": test_videos})
        
        rng_input = jax.random.PRNGKey(456)
        inputs = {"videos": test_videos, "rng": rng_input}
        
        class MockState:
            def __init__(self, apply_fn):
                self.apply_fn = apply_fn
        
        state = MockState(tokenizer.apply)
        
        grad_fn = jax.value_and_grad(tokenizer_loss_fn, has_aux=True, allow_int=True)
        (loss1, (recon1, metrics1)), grads1 = grad_fn(params, state, inputs)
        (loss2, (recon2, metrics2)), grads2 = grad_fn(params, state, inputs)
        
        np.testing.assert_array_equal(loss1, loss2)
        np.testing.assert_array_equal(recon1, recon2)
        for key in metrics1.keys():
            np.testing.assert_array_equal(metrics1[key], metrics2[key])
        
        jax.tree.map(lambda x, y: np.testing.assert_array_equal(x, y), grads1, grads2)

    def test_tokenizer_with_different_batch_sizes(self):
        """Test that tokenizer works consistently with different batch sizes."""
        tokenizer = TokenizerVQVAE(
            in_dim=self.image_channels,
            model_dim=self.model_dim,
            latent_dim=self.latent_dim,
            num_latents=self.num_latents,
            patch_size=self.patch_size,
            num_blocks=self.num_blocks,
            num_heads=self.num_heads,
            dropout=0.0,
            codebook_dropout=0.0,
        )
        
        rng = jax.random.PRNGKey(self.seed)
        test_videos_1 = jax.random.uniform(rng, (1, self.seq_len, self.image_height, self.image_width, self.image_channels))
        params = tokenizer.init(rng, {"videos": test_videos_1})
        
        output_1 = tokenizer.apply(params, {"videos": test_videos_1}, training=True)
        
        test_videos_2 = jax.random.uniform(rng, (2, self.seq_len, self.image_height, self.image_width, self.image_channels))
        output_2 = tokenizer.apply(params, {"videos": test_videos_2}, training=True)
        
        for key in output_1.keys():
            np.testing.assert_array_equal(output_1[key], jax.tree.map(lambda x: x[0:1], output_2[key]))

    def test_tokenizer_indices_consistency(self):
        """Test that tokenizer produces consistent indices for same input."""
        tokenizer = TokenizerVQVAE(
            in_dim=self.image_channels,
            model_dim=self.model_dim,
            latent_dim=self.latent_dim,
            num_latents=self.num_latents,
            patch_size=self.patch_size,
            num_blocks=self.num_blocks,
            num_heads=self.num_heads,
            dropout=0.0,  # Zero dropout for deterministic test
            codebook_dropout=0.0,
        )
        
        # Initialize
        rng = jax.random.PRNGKey(self.seed)
        test_videos = jax.random.uniform(rng, (self.batch_size, self.seq_len, self.image_height, self.image_width, self.image_channels))
        params = tokenizer.init(rng, {"videos": test_videos})
        
        # Run vq_encode multiple times
        indices_list = []
        for _ in range(3):
            output = tokenizer.apply(params, test_videos, method=tokenizer.vq_encode, training=True)
            indices_list.append(output["indices"])
        
        # All indices should be identical
        for i in range(1, len(indices_list)):
            np.testing.assert_array_equal(indices_list[0], indices_list[i])


if __name__ == "__main__":
    unittest.main()