[content] tests/test_tokenizer_reproducibility.py (python)
    def test_realistic_training_setting_reproducibility(self):
        """Test tokenizer with realistic settings matching actual training code."""
        # Use realistic parameters matching train_tokenizer.py defaults
        realistic_params = {
            'in_dim': 3,  # image_channels
            'model_dim': 512,  # model_dim
            'latent_dim': 32,  # latent_dim
            'num_latents': 1024,  # num_latents
            'patch_size': 4,  # patch_size
            'num_blocks': 8,  # num_blocks
            'num_heads': 8,  # num_heads
            'dropout': 0.0,  # dropout
            'codebook_dropout': 0.01,  # codebook_dropout
        }
        
        # Use realistic input dimensions matching training
        realistic_batch_size = 2  # Smaller for testing
        realistic_seq_len = 16  # seq_len
        realistic_height = 90  # image_height
        realistic_width = 160  # image_width
        realistic_channels = 3  # image_channels
        
        # Create realistic test videos
        rng = jax.random.PRNGKey(self.seed)
        rng, input_rng = jax.random.split(rng)
        realistic_videos = jax.random.uniform(
            input_rng,
            (realistic_batch_size, realistic_seq_len, realistic_height, realistic_width, realistic_channels)
        )
        
        # Create two identical tokenizers
        tokenizer1 = TokenizerVQVAE(**realistic_params)
        tokenizer2 = TokenizerVQVAE(**realistic_params)
        
        # Initialize with proper input format (dictionary with "videos" key)
        inputs = {"videos": realistic_videos}
        params1 = tokenizer1.init(self.rng, inputs)
        params2 = tokenizer2.init(self.rng, inputs)
        
        # Test full forward pass with proper input format and RNG handling
        rng_dropout = jax.random.PRNGKey(123)
        output1 = tokenizer1.apply(
            params1, 
            inputs, 
            training=True, 
            rngs={"dropout": rng_dropout}
        )
        output2 = tokenizer2.apply(
            params2, 
            inputs, 
            training=True, 
            rngs={"dropout": rng_dropout}
        )
        
        # Verify all outputs are identical
        expected_keys = ["recon", "z_q", "z", "emb", "indices"]
        for key in expected_keys:
            self.assertIn(key, output1)
            self.assertIn(key, output2)
            np.testing.assert_array_equal(output1[key], output2[key])
        
        # Test that outputs have correct shapes
        self.assertEqual(output1["recon"].shape, realistic_videos.shape)
        self.assertEqual(output1["indices"].shape, (realistic_batch_size, realistic_seq_len, 
                                                   (realistic_height // realistic_params['patch_size']) * 
                                                   (realistic_width // realistic_params['patch_size'])))
        
        # Test multiple forward passes with same RNG are reproducible
        outputs = []
        for _ in range(3):
            output = tokenizer1.apply(
                params1, 
                inputs, 
                training=True, 
                rngs={"dropout": rng_dropout}
            )
            outputs.append(output)
        
        for i in range(1, len(outputs)):
            for key in expected_keys:
                np.testing.assert_array_equal(outputs[0][key], outputs[i][key])
        
        # Test that different dropout RNGs produce different outputs
        rng_dropout1 = jax.random.PRNGKey(123)
        rng_dropout2 = jax.random.PRNGKey(456)
        
        output_diff1 = tokenizer1.apply(
            params1, 
            inputs, 
            training=True, 
            rngs={"dropout": rng_dropout1}
        )
        output_diff2 = tokenizer1.apply(
            params1, 
            inputs, 
            training=True, 
            rngs={"dropout": rng_dropout2}
        )
        
        # With dropout=0.0, outputs should be identical even with different RNGs
        for key in expected_keys:
            np.testing.assert_array_equal(output_diff1[key], output_diff2[key])
        
        # Test with non-zero dropout that outputs differ
        tokenizer_with_dropout = TokenizerVQVAE(
            **{**realistic_params, 'dropout': 0.1, 'codebook_dropout': 0.1}
        )
        params_with_dropout = tokenizer_with_dropout.init(self.rng, inputs)
        
        output_dropout1 = tokenizer_with_dropout.apply(
            params_with_dropout, 
            inputs, 
            training=True, 
            rngs={"dropout": rng_dropout1}
        )
        output_dropout2 = tokenizer_with_dropout.apply(
            params_with_dropout, 
            inputs, 
            training=True, 
            rngs={"dropout": rng_dropout2}
        )
        
        # With dropout > 0, some outputs should differ
        different_outputs = 0
        for key in ["recon", "z_q", "z", "emb"]:
            try:
                np.testing.assert_array_equal(output_dropout1[key], output_dropout2[key])
            except AssertionError:
                different_outputs += 1
        
        # At least some outputs should be different with dropout
        self.assertGreater(different_outputs, 0, 
                          "With dropout > 0, some outputs should differ between different RNGs")