[tab] tests/test_tokenizer_edge_cases.py (python)
import unittest
import jax
import jax.numpy as jnp
import numpy as np
from typing import Dict, Any

from models.tokenizer import TokenizerVQVAE
from utils.nn import STTransformer, 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)
        
        # Model parameters
        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
        
        # Create multiple instances and check codebook initialization
        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)
        
        # All codebooks should be identical when using same seed
        for i in range(1, len(codebooks)):
            np.testing.assert_array_equal(codebooks[0], codebooks[i])

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

    def test_tokenizer_with_zero_dropout_reproducibility(self):
        """Test that tokenizer with zero dropout is deterministic."""
        # Create tokenizer with zero dropout
        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
            codebook_dropout=0.0,  # Zero codebook dropout
        )
        
        # 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 multiple forward passes without dropout rng
        outputs = []
        for _ in range(3):
            output = tokenizer.apply(params, {"videos": test_videos}, training=True)
            outputs.append(output)
        
        # All outputs should be identical (no dropout)
        for i in range(1, len(outputs)):
            for key in outputs[0].keys():
                np.testing.assert_array_equal(outputs[0][key], outputs[i][key])

    def test_tokenizer_training_vs_eval_mode_reproducibility(self):
        """Test that tokenizer produces same output in training and eval mode with zero dropout."""
        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
            codebook_dropout=0.0,  # Zero codebook dropout
        )
        
        # 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 in training mode
        train_output = tokenizer.apply(params, {"videos": test_videos}, training=True)
        
        # Run in eval mode
        eval_output = tokenizer.apply(params, {"videos": test_videos}, training=False)
        
        # Should be identical with zero dropout
        for key in train_output.keys():
            np.testing.assert_array_equal(train_output[key], eval_output[key])

    def test_vector_quantizer_training_vs_eval_mode(self):
        """Test VectorQuantizer behavior in training vs eval mode."""
        latent_dim = 32
        num_latents = 256
        dropout = 0.1
        
        vq = VectorQuantizer(latent_dim, num_latents, dropout)
        
        # Initialize
        rng = jax.random.PRNGKey(self.seed)
        x = jax.random.uniform(rng, (100, latent_dim))
        params = vq.init(rng, x, training=True)
        
        # Run in training mode with dropout
        rng_dropout = jax.random.PRNGKey(123)
        train_output = vq.apply(params, x, training=True, rngs={"dropout": rng_dropout})
        
        # Run in eval mode (should ignore dropout)
        eval_output = vq.apply(params, x, training=False)
        
        # Run in training mode without dropout
        train_no_dropout_output = vq.apply(params, x, training=True, rngs={"dropout": jax.random.PRNGKey(456)})
        
        # Eval mode should be same as training without dropout
        for o1, o2 in zip(eval_output, train_no_dropout_output):
            np.testing.assert_array_equal(o1, o2)

    def test_tokenizer_gradient_reproducibility(self):
        """Test that gradients are reproducible."""
        from train_tokenizer import tokenizer_loss_fn
        
        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,
        )
        
        # 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})
        
        # Create inputs with same rng
        rng_input = jax.random.PRNGKey(456)
        inputs = {"videos": test_videos, "rng": rng_input}
        
        # Compute gradients twice
        grad_fn = jax.value_and_grad(tokenizer_loss_fn, has_aux=True, allow_int=True)
        (loss1, (recon1, metrics1)), grads1 = grad_fn(params, tokenizer, inputs)
        (loss2, (recon2, metrics2)), grads2 = grad_fn(params, tokenizer, inputs)
        
        # Should be identical
        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])
        
        # Check gradients are identical
        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,  # Zero dropout for deterministic test
            codebook_dropout=0.0,
        )
        
        # Initialize with batch size 1
        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})
        
        # Test with batch size 1
        output_1 = tokenizer.apply(params, {"videos": test_videos_1}, training=True)
        
        # Test with batch size 2 (first element should match)
        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)
        
        # First element of batch size 2 should match batch size 1
        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()