[tab] tests/test_tokenizer_reproducibility.py (python)
import unittest
import jax
import numpy as np

from models.tokenizer import TokenizerVQVAE
from utils.nn import STTransformer, VectorQuantizer, PositionalEncoding, STBlock
from utils.preprocess import patchify, unpatchify


class TokenizerReproducibilityTest(unittest.TestCase):
    """Test reproducibility of tokenizer components and full model."""

    def setUp(self):
        super().setUp()
        self.seed = 42
        self.rng = jax.random.PRNGKey(self.seed)
        
        self.batch_size = 1
        self.seq_len = 4
        self.image_height = 32
        self.image_width = 32
        self.image_channels = 3
        self.patch_size = 4
        self.model_dim = 64
        self.latent_dim = 16
        self.num_latents = 64
        self.num_blocks = 1
        self.num_heads = 2
        self.dropout = 0.1
        self.codebook_dropout = 0.01
        
        self.rng, input_rng = jax.random.split(self.rng)
        self.test_videos = jax.random.uniform(
            input_rng, 
            (self.batch_size, self.seq_len, self.image_height, self.image_width, self.image_channels)
        )

    def test_positional_encoding_reproducibility(self):
        """Test that positional encoding is deterministic."""
        dim = 64
        max_len = 100
        
        pe1 = PositionalEncoding(dim, max_len)
        pe2 = PositionalEncoding(dim, max_len)
        
        rng1, rng2 = jax.random.split(self.rng)
        x = jax.random.uniform(rng1, (2, 3, 50, dim))
        
        params1 = pe1.init(rng1, x)
        params2 = pe2.init(rng2, x)
        
        output1 = pe1.apply(params1, x)
        output2 = pe2.apply(params2, x)
        
        np.testing.assert_array_equal(output1, output2)

    def test_st_block_reproducibility(self):
        """Test that STBlock is reproducible."""
        dim = 64
        num_heads = 4
        dropout = 0.1
        
        block1 = STBlock(dim, num_heads, dropout)
        block2 = STBlock(dim, num_heads, dropout)
        
        x = jax.random.uniform(self.rng, (2, 3, 16, dim))
        
        params1 = block1.init(self.rng, x)
        params2 = block2.init(self.rng, x)
        
        rng_dropout = jax.random.PRNGKey(123)
        output1 = block1.apply(params1, x, rngs={"dropout": rng_dropout})
        output2 = block2.apply(params2, x, rngs={"dropout": rng_dropout})
        
        np.testing.assert_array_equal(output1, output2)

    def test_st_transformer_reproducibility(self):
        """Test that STTransformer is reproducible."""
        model_dim = 64
        out_dim = 32
        num_blocks = 2
        num_heads = 4
        dropout = 0.1
        
        transformer1 = STTransformer(model_dim, out_dim, num_blocks, num_heads, dropout)
        transformer2 = STTransformer(model_dim, out_dim, num_blocks, num_heads, dropout)
        
        x = jax.random.uniform(self.rng, (2, 3, 16, 48))
        
        params1 = transformer1.init(self.rng, x)
        params2 = transformer2.init(self.rng, x)
        
        rng_dropout = jax.random.PRNGKey(123)
        output1 = transformer1.apply(params1, x, rngs={"dropout": rng_dropout})
        output2 = transformer2.apply(params2, x, rngs={"dropout": rng_dropout})
        
        np.testing.assert_array_equal(output1, output2)

    def test_vector_quantizer_reproducibility(self):
        """Test that VectorQuantizer is reproducible."""
        latent_dim = 32
        num_latents = 256
        dropout = 0.01
        
        vq1 = VectorQuantizer(latent_dim, num_latents, dropout)
        vq2 = VectorQuantizer(latent_dim, num_latents, dropout)
        
        x = jax.random.uniform(self.rng, (100, latent_dim))
        
        params1 = vq1.init(self.rng, x, training=True)
        params2 = vq2.init(self.rng, x, training=True)
        
        rng_dropout = jax.random.PRNGKey(123)
        output1 = vq1.apply(params1, x, training=True, rngs={"dropout": rng_dropout})
        output2 = vq2.apply(params2, x, training=True, rngs={"dropout": rng_dropout})
        
        for o1, o2 in zip(output1, output2):
            np.testing.assert_array_equal(o1, o2)

    def test_patchify_unpatchify_reproducibility(self):
        """Test that patchify/unpatchify operations are deterministic."""
        videos = self.test_videos
        patch_size = self.patch_size
        
        # Apply patchify twice
        patches1 = patchify(videos, patch_size)
        patches2 = patchify(videos, patch_size)
        
        # Should be identical
        np.testing.assert_array_equal(patches1, patches2)
        
        # Apply unpatchify twice
        H, W = videos.shape[2:4]
        recon1 = unpatchify(patches1, patch_size, H, W)
        recon2 = unpatchify(patches2, patch_size, H, W)
        
        # Should be identical
        np.testing.assert_array_equal(recon1, recon2)

    def test_tokenizer_vq_encode_reproducibility(self):
        """Test that tokenizer vq_encode is reproducible."""
        # Create two instances
        tokenizer1 = 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,
        )
        tokenizer2 = 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 both with the SAME rng to ensure identical parameters
        params1 = tokenizer1.init(self.rng, {"videos": self.test_videos})
        params2 = tokenizer2.init(self.rng, {"videos": self.test_videos})  # Use same rng
        
        # Apply vq_encode with same dropout rng
        rng_dropout = jax.random.PRNGKey(123)
        output1 = tokenizer1.apply(params1, self.test_videos, method=tokenizer1.vq_encode, training=True, rngs={"dropout": rng_dropout})
        output2 = tokenizer2.apply(params2, self.test_videos, method=tokenizer2.vq_encode, training=True, rngs={"dropout": rng_dropout})
        
        # vq_encode returns a dict
        # Should be identical
        for key in output1.keys():
            np.testing.assert_array_equal(output1[key], output2[key])

    def test_full_tokenizer_reproducibility(self):
        """Test that full tokenizer forward pass is reproducible."""
        # Create two instances
        tokenizer1 = 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,
        )
        tokenizer2 = 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 both with the SAME rng to ensure identical parameters
        params1 = tokenizer1.init(self.rng, {"videos": self.test_videos})
        params2 = tokenizer2.init(self.rng, {"videos": self.test_videos})  # Use same rng
        
        # Apply full forward pass with same dropout rng
        rng_dropout = jax.random.PRNGKey(123)
        output1 = tokenizer1.apply(params1, {"videos": self.test_videos}, training=True, rngs={"dropout": rng_dropout})
        output2 = tokenizer2.apply(params2, {"videos": self.test_videos}, training=True, rngs={"dropout": rng_dropout})
        
        # Full tokenizer returns a dict
        # Should be identical
        for key in output1.keys():
            np.testing.assert_array_equal(output1[key], output2[key])

    def test_tokenizer_loss_reproducibility(self):
        """Test that tokenizer loss computation is reproducible."""
        from train_tokenizer import tokenizer_loss_fn
        
        # Create tokenizer
        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 with the SAME rng to ensure identical parameters
        params1 = tokenizer.init(self.rng, {"videos": self.test_videos})
        params2 = tokenizer.init(self.rng, {"videos": self.test_videos})  # Use same rng
        
        # Create inputs with same rng
        rng_input = jax.random.PRNGKey(456)
        inputs = {"videos": self.test_videos, "rng": rng_input}
        
        # Create a mock state object
        class MockState:
            def __init__(self, apply_fn):
                self.apply_fn = apply_fn
        
        state1 = MockState(tokenizer.apply)
        state2 = MockState(tokenizer.apply)
        
        # Create mock args object
        class MockArgs:
            def __init__(self):
                self.vq_beta = 0.25
                self.num_latents = 64  # Use a fixed value
                self.log_gradients = False
        
        args = MockArgs()
        
        # Compute loss twice
        loss1, (recon1, metrics1) = tokenizer_loss_fn(params1, state1, inputs, args)
        loss2, (recon2, metrics2) = tokenizer_loss_fn(params2, state2, inputs, args)
        
        # 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])

    def test_training_step_reproducibility(self):
        """Test that training step is reproducible."""
        from train_tokenizer import create_training_functions
        
        # Create tokenizer
        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 with the SAME rng to ensure identical parameters
        params1 = tokenizer.init(self.rng, {"videos": self.test_videos})
        params2 = tokenizer.init(self.rng, {"videos": self.test_videos})  # Use same rng
        
        # Create train state
        import optax
        tx = optax.adamw(learning_rate=1e-4)
        from flax.training.train_state import TrainState
        
        state1 = TrainState.create(apply_fn=tokenizer.apply, params=params1, tx=tx)
        state2 = TrainState.create(apply_fn=tokenizer.apply, params=params2, tx=tx)
        
        # Create inputs with same rng
        rng_input = jax.random.PRNGKey(456)
        inputs = {"videos": self.test_videos, "rng": rng_input}
        
        # Create mock args object
        class MockArgs:
            def __init__(self):
                self.vq_beta = 0.25
                self.num_latents = 64  # Use a fixed value
                self.log_gradients = False
        
        args = MockArgs()
        
        # Create training functions with args captured in closure
        tokenizer_loss_fn, train_step = create_training_functions(args)
        
        # Run training step twice
        state1_out, loss1, recon1, metrics1 = train_step(state1, inputs)
        state2_out, loss2, recon2, metrics2 = train_step(state2, 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])

    def test_multiple_forward_passes_reproducibility(self):
        """Test that multiple forward passes with same seed are reproducible."""
        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
        params = tokenizer.init(self.rng, {"videos": self.test_videos})
        
        # Run multiple forward passes with same dropout rng
        rng_dropout = jax.random.PRNGKey(123)
        outputs = []
        for _ in range(3):
            output = tokenizer.apply(params, {"videos": self.test_videos}, training=True, rngs={"dropout": rng_dropout})
            outputs.append(output)
        
        # All outputs should be identical
        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_different_dropout_rngs_produce_different_outputs(self):
        """Test that different dropout RNGs produce different outputs (sanity check)."""
        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
        params = tokenizer.init(self.rng, {"videos": self.test_videos})
        
        # Run forward passes with different dropout rngs
        rng1 = jax.random.PRNGKey(123)
        rng2 = jax.random.PRNGKey(456)
        
        output1 = tokenizer.apply(params, {"videos": self.test_videos}, training=True, rngs={"dropout": rng1})
        output2 = tokenizer.apply(params, {"videos": self.test_videos}, training=True, rngs={"dropout": rng2})
        
        # Outputs should be different (due to dropout)
        for key in output1.keys():
            if key in ["recon", "z_q", "z", "emb"]:  # These should be affected by dropout
                with self.assertRaises(AssertionError):
                    np.testing.assert_array_equal(output1[key], output2[key])


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