p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] tests/test_tokenizer_reproducibility.py (python)
        # Create mock args object
        class MockArgs:
            def __init__(self):
                self.vq_beta = 0.25
                self.num_latents = 64
                self.log_gradients = False
        
        args = MockArgs()
[content] tests/test_tokenizer_reproducibility.py (python)
        # Handle potential tuple returns
        if isinstance(output1, tuple):
            output1 = output1[0]
        if isinstance(output2, tuple):
            output2 = output2[0]
        
        # Check if outputs are dictionaries
        if isinstance(output1, dict) and isinstance(output2, dict):
            for key in output1.keys():
                np.testing.assert_array_equal(output1[key], output2[key])
        else:
            np.testing.assert_array_equal(output1, output2)
[content] tests/test_tokenizer_reproducibility.py (python)
        # Handle tuple returns for vector quantizer
        if isinstance(output1, tuple):
            for o1, o2 in zip(output1, output2):
                np.testing.assert_array_equal(o1, o2)
        else:
            np.testing.assert_array_equal(output1, output2)
[content] tests/test_tokenizer_reproducibility.py (python)
        if isinstance(output1, dict) and isinstance(output2, dict):
            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])
        else:
            with self.assertRaises(AssertionError):
                np.testing.assert_array_equal(output1, output2)
[content] tests/test_tokenizer_reproducibility.py (python)
            if isinstance(outputs[0], dict):
                for key in outputs[0].keys():
                    np.testing.assert_array_equal(outputs[0][key], outputs[i][key])
            else:
                np.testing.assert_array_equal(outputs[0], outputs[i])
[content] tests/test_tokenizer_reproducibility.py (python)
        loss1, (recon1, metrics1) = tokenizer_loss_fn(params1, state1, inputs)
        loss2, (recon2, metrics2) = tokenizer_loss_fn(params2, state2, 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])

    def test_training_step_reproducibility(self):
        """Test that training step is reproducible."""
        # Create mock args object
[content] tests/test_tokenizer_reproducibility.py (python)
        # Outputs should be different (due to dropout) - both are dictionaries
        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])
[content] tests/test_tokenizer_reproducibility.py (python)
        # Now run the same sequence again with the same seed
        outputs2 = []
        for i in range(3):
            rng_dropout = jax.random.fold_in(base_rng, i)
            output = tokenizer.apply(params, {"videos": self.test_videos}, training=True, rngs={"dropout": rng_dropout})
            outputs2.append(output)
        
        # Verify that corresponding outputs from both runs are identical
        for i in range(3):
            for key in outputs[i].keys():
                np.testing.assert_array_equal(outputs[i][key], outputs2[i][key])
[content] tests/test_tokenizer_edge_cases.py (python)
        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)
[content] tests/test_tokenizer_edge_cases.py (python)
        # With zero dropout, outputs should be identical
        np.testing.assert_array_equal(train_output["recon"], eval_output["recon"])
        np.testing.assert_array_equal(train_output["z_q"], eval_output["z_q"])
        np.testing.assert_array_equal(train_output["z"], eval_output["z"])
        np.testing.assert_array_equal(train_output["emb"], eval_output["emb"])
        np.testing.assert_array_equal(train_output["indices"], eval_output["indices"])

    def test_tokenizer_training_vs_eval_mode_with_dropout(self):
        """Test that tokenizer produces different outputs in training vs eval mode with 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.1,
            codebook_dropout=0.1,
        )
        
        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})
        
        # Use same RNG for both to ensure fair comparison
        rng_dropout = jax.random.PRNGKey(123)
        
        train_output = tokenizer.apply(params, {"videos": test_videos}, training=True, rngs={"dropout": rng_dropout})
        
        eval_output = tokenizer.apply(params, {"videos": test_videos}, training=False)
        
        # With dropout, outputs should be different
        with self.assertRaises(AssertionError):
            np.testing.assert_array_equal(train_output["recon"], eval_output["recon"])
        
        # But indices should be the same (deterministic)
        np.testing.assert_array_equal(train_output["indices"], eval_output["indices"])
[content] tests/test_tokenizer_edge_cases.py (python)
        # Both elements in batch should be identical since input is identical
        np.testing.assert_array_equal(output_batch["recon"][0], output_batch["recon"][1])
        np.testing.assert_array_equal(output_batch["z_q"][0], output_batch["z_q"][1])
        np.testing.assert_array_equal(output_batch["z"][0], output_batch["z"][1])
        np.testing.assert_array_equal(output_batch["emb"][0], output_batch["emb"][1])
        np.testing.assert_array_equal(output_batch["indices"][0], output_batch["indices"][1])
[content] tests/test_tokenizer_edge_cases.py (python)
        # First element of batch should match single video output
        np.testing.assert_array_equal(output_single["recon"], output_batch["recon"][0:1])
        np.testing.assert_array_equal(output_single["z_q"], output_batch["z_q"][0:1])
        np.testing.assert_array_equal(output_single["z"], output_batch["z"][0:1])
        np.testing.assert_array_equal(output_single["emb"], output_batch["emb"][0:1])
        np.testing.assert_array_equal(output_single["indices"], output_batch["indices"][0:1])
[content] tests/test_tokenizer_edge_cases.py (python)
        # Create a batch with the same video repeated
        batch_videos = jnp.concatenate([single_video, single_video], axis=0)
        
        # Initialize with single video
        params = tokenizer.init(rng, {"videos": single_video})
        
        # Process single video
        output_single = tokenizer.apply(params, {"videos": single_video}, training=True)
        
        # Process batch with same video repeated
        output_batch = tokenizer.apply(params, {"videos": batch_videos}, training=True)
[content] tests/test_tokenizer_edge_cases.py (python)
        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)
[content] tests/test_tokenizer_edge_cases.py (python)
        # Create a batch with the same video repeated
        batch_videos = jnp.concatenate([single_video, single_video], axis=0)
        
        # Initialize with single video
        params = tokenizer.init(rng, {"videos": single_video})
        
        # Process single video
        output_single = tokenizer.apply(params, {"videos": single_video}, training=True)
        
        # Process batch with same video repeated
        output_batch = tokenizer.apply(params, {"videos": batch_videos}, training=True)
        
        # Cast to dict to help linter understand the type
        output_single = dict(output_single)
        output_batch = dict(output_batch)
[content] tests/test_tokenizer_edge_cases.py (python)
        np.testing.assert_array_equal(batch_recon[0], batch_recon[1])
        np.testing.assert_array_equal(batch_z_q[0], batch_z_q[1])
        np.testing.assert_array_equal(batch_z[0], batch_z[1])
        np.testing.assert_array_equal(batch_emb[0], batch_emb[1])
        np.testing.assert_array_equal(batch_indices[0], batch_indices[1])
[content] tests/test_tokenizer_edge_cases.py (python)
        single_z_q = getattr(output_single, "z_q", output_single["z_q"])
        batch_z_q = getattr(output_batch, "z_q", output_batch["z_q"])
        np.testing.assert_array_equal(single_z_q, batch_z_q[0:1])
        
        single_z = getattr(output_single, "z", output_single["z"])
        batch_z = getattr(output_batch, "z", output_batch["z"])
        np.testing.assert_array_equal(single_z, batch_z[0:1])
        
        single_emb = getattr(output_single, "emb", output_single["emb"])
        batch_emb = getattr(output_batch, "emb", output_batch["emb"])
        np.testing.assert_array_equal(single_emb, batch_emb[0:1])
        
        single_indices = getattr(output_single, "indices", output_single["indices"])
        batch_indices = getattr(output_batch, "indices", output_batch["indices"])
        np.testing.assert_array_equal(single_indices, batch_indices[0:1])
[content] tests/test_tokenizer_edge_cases.py (python)
        # First element of batch should match single video output
        # Use getattr to avoid linter issues with dictionary access
        single_recon = getattr(output_single, "recon", output_single["recon"])
        batch_recon = getattr(output_batch, "recon", output_batch["recon"])
        np.testing.assert_array_equal(single_recon, batch_recon[0:1])
[content] tests/test_tokenizer_edge_cases.py (python)
        np.testing.assert_array_equal(output_batch["recon"][0], output_batch["recon"][1])
        np.testing.assert_array_equal(output_batch["z_q"][0], output_batch["z_q"][1])
        np.testing.assert_array_equal(output_batch["z"][0], output_batch["z"][1])
        np.testing.assert_array_equal(output_batch["emb"][0], output_batch["emb"][1])
        np.testing.assert_array_equal(output_batch["indices"][0], output_batch["indices"][1])
[content] tests/test_tokenizer_edge_cases.py (python)
            output = tokenizer.apply(params, {"videos": test_videos}, method=tokenizer.vq_encode, training=True)
            indices = getattr(output, "indices", output["indices"])
            indices_list.append(indices)
[content] tests/test_tokenizer_edge_cases.py (python)
        # Both elements in batch should be identical since input is identical
        np.testing.assert_array_equal(batch_recon[0], batch_recon[1])
        np.testing.assert_array_equal(batch_z_q[0], batch_z_q[1])
        np.testing.assert_array_equal(batch_z[0], batch_z[1])
        np.testing.assert_array_equal(batch_emb[0], batch_emb[1])
        np.testing.assert_array_equal(batch_indices[0], batch_indices[1])
[content] tests/test_tokenizer_edge_cases.py (python)
        single_indices = getattr(output_single, "indices", output_single["indices"])
        batch_indices = getattr(output_batch, "indices", output_batch["indices"])
        np.testing.assert_array_equal(single_indices, batch_indices[0:1])
[content] tests/test_tokenizer_edge_cases.py (python)
        single_emb = getattr(output_single, "emb", output_single["emb"])
        batch_emb = getattr(output_batch, "emb", output_batch["emb"])
        np.testing.assert_array_equal(single_emb, batch_emb[0:1])
[content] tests/test_tokenizer_edge_cases.py (python)
        # Create a single video
        single_video = jax.random.uniform(rng, (1, self.seq_len, self.image_height, self.image_width, self.image_channels))
        
        # Create a batch with the same video repeated
        batch_videos = jnp.concatenate([single_video, single_video], axis=0)
        
        # Initialize with single video
        params = tokenizer.init(rng, {"videos": single_video})
        
        # Process single video
        output_single = tokenizer.apply(params, {"videos": single_video}, training=True)
        
        # Process batch with same video repeated
        output_batch = tokenizer.apply(params, {"videos": batch_videos}, training=True)
        
        # First element of batch should match single video output
        # Use getattr to avoid linter issues with dictionary access
        single_recon = getattr(output_single, "recon", output_single["recon"])
        batch_recon = getattr(output_batch, "recon", output_batch["recon"])
        np.testing.assert_array_equal(single_recon, batch_recon[0:1])
        
        single_z_q = getattr(output_single, "z_q", output_single["z_q"])
        batch_z_q = getattr(output_batch, "z_q", output_batch["z_q"])
        np.testing.assert_array_equal(single_z_q, batch_z_q[0:1])
        
        single_z = getattr(output_single, "z", output_single["z"])
        batch_z = getattr(output_batch, "z", output_batch["z"])
        np.testing.assert_array_equal(single_z, batch_z[0:1])
[content] tests/test_tokenizer_edge_cases.py (python)
        np.testing.assert_array_equal(output_single["recon"], output_batch["recon"][0:1])
        np.testing.assert_array_equal(output_single["z_q"], output_batch["z_q"][0:1])
        np.testing.assert_array_equal(output_single["z"], output_batch["z"][0:1])
        np.testing.assert_array_equal(output_single["emb"], output_batch["emb"][0:1])
        np.testing.assert_array_equal(output_single["indices"], output_batch["indices"][0:1])
[content] tests/test_tokenizer_edge_cases.py (python)
        # Use getattr to avoid linter issues with dictionary access
        single_recon = getattr(output_single, "recon", output_single["recon"])
        batch_recon = getattr(output_batch, "recon", output_batch["recon"])
        np.testing.assert_array_equal(single_recon, batch_recon[0:1])
        
        single_z_q = getattr(output_single, "z_q", output_single["z_q"])
        batch_z_q = getattr(output_batch, "z_q", output_batch["z_q"])
        np.testing.assert_array_equal(single_z_q, batch_z_q[0:1])
        
        single_z = getattr(output_single, "z", output_single["z"])
        batch_z = getattr(output_batch, "z", output_batch["z"])
        np.testing.assert_array_equal(single_z, batch_z[0:1])
        
        single_emb = getattr(output_single, "emb", output_single["emb"])
        batch_emb = getattr(output_batch, "emb", output_batch["emb"])
        np.testing.assert_array_equal(single_emb, batch_emb[0:1])
        
        single_indices = getattr(output_single, "indices", output_single["indices"])
        batch_indices = getattr(output_batch, "indices", output_batch["indices"])
        np.testing.assert_array_equal(single_indices, batch_indices[0:1])
[content] tests/test_tokenizer_edge_cases.py (python)
        # Initialize with single video
        params = tokenizer.init(rng, {"videos": single_video})
        
        # Process single video
        output_single = tokenizer.apply(params, {"videos": single_video}, training=True)
        
        # Process batch with same video repeated
        output_batch = tokenizer.apply(params, {"videos": batch_videos}, training=True)
        
        # First element of batch should match single video output
        np.testing.assert_array_equal(output_single["recon"], output_batch["recon"][0:1])
        np.testing.assert_array_equal(output_single["z_q"], output_batch["z_q"][0:1])
        np.testing.assert_array_equal(output_single["z"], output_batch["z"][0:1])
        np.testing.assert_array_equal(output_single["emb"], output_batch["emb"][0:1])
        np.testing.assert_array_equal(output_single["indices"], output_batch["indices"][0:1])
[content] tests/test_tokenizer_reproducibility.py (python)
        # Handle both tuple and single value returns
        if isinstance(output1, tuple):
            output1 = output1[0]
        if isinstance(output2, tuple):
            output2 = output2[0]
[content] tests/test_tokenizer_reproducibility.py (python)
        if isinstance(output1, dict):
            for key in output1.keys():
                np.testing.assert_array_equal(output1[key], output2[key])
        else:
            # Handle tuple return
            for o1, o2 in zip(output1, output2):
                np.testing.assert_array_equal(o1, o2)
[content] tests/test_tokenizer_reproducibility.py (python)
        if isinstance(output1, dict):
            for key in output1.keys():
                if key in ["recon", "z_q", "z", "emb"]:
                    with self.assertRaises(AssertionError):
                        np.testing.assert_array_equal(output1[key], output2[key])
        else:
            # Handle tuple return - check first 4 elements (z_q, z, emb, indices)
            for i in range(4):
                with self.assertRaises(AssertionError):
                    np.testing.assert_array_equal(output1[i], output2[i])
[content] tests/test_tokenizer_reproducibility.py (python)
        if isinstance(output_single, dict):
            np.testing.assert_array_equal(output_single["recon"], output_batch["recon"][0:1])
            np.testing.assert_array_equal(output_single["z_q"], output_batch["z_q"][0:1])
            np.testing.assert_array_equal(output_single["z"], output_batch["z"][0:1])
            np.testing.assert_array_equal(output_single["emb"], output_batch["emb"][0:1])
            np.testing.assert_array_equal(output_single["indices"], output_batch["indices"][0:1])
        else:
            # Handle tuple return - (z_q, z, emb, indices) + recon from full forward pass
            np.testing.assert_array_equal(output_single[0], output_batch[0][0:1])  # z_q
            np.testing.assert_array_equal(output_single[1], output_batch[1][0:1])  # z
            np.testing.assert_array_equal(output_single[2], output_batch[2][0:1])  # emb
            np.testing.assert_array_equal(output_single[3], output_batch[3][0:1])  # indices
[content] tests/test_tokenizer_reproducibility.py (python)
            if isinstance(output, dict):
                indices_list.append(output["indices"])
            else:
                # Handle tuple return - indices is the 4th element
                indices_list.append(output[3])
[content] tests/test_tokenizer_reproducibility.py (python)
        for key in output1.keys():
            if key in ["recon", "z_q", "z", "emb"]:
                with self.assertRaises(AssertionError):
                    np.testing.assert_array_equal(output1[key], output2[key])
[content] tests/test_tokenizer_reproducibility.py (python)
        np.testing.assert_array_equal(output_single["recon"], output_batch["recon"][0:1])
        np.testing.assert_array_equal(output_single["z_q"], output_batch["z_q"][0:1])
        np.testing.assert_array_equal(output_single["z"], output_batch["z"][0:1])
        np.testing.assert_array_equal(output_single["emb"], output_batch["emb"][0:1])
        np.testing.assert_array_equal(output_single["indices"], output_batch["indices"][0:1])
[content] tests/test_tokenizer_reproducibility.py (python)
        # Full tokenizer returns a dictionary
        if isinstance(output1, dict):
            for key in output1.keys():
                if key in ["recon", "z_q", "z", "emb"]:
                    with self.assertRaises(AssertionError):
                        np.testing.assert_array_equal(output1[key], output2[key])
        else:
            # Handle tuple return - check first 4 elements (z_q, z, emb, indices)
            for i in range(4):
                with self.assertRaises(AssertionError):
                    np.testing.assert_array_equal(output1[i], output2[i])
[content] tests/test_tokenizer_reproducibility.py (python)
        # Compare the first video from batch with single video
        if isinstance(output_single, dict):
            np.testing.assert_array_equal(output_single["recon"], output_batch["recon"][0:1])
            np.testing.assert_array_equal(output_single["z_q"], output_batch["z_q"][0:1])
            np.testing.assert_array_equal(output_single["z"], output_batch["z"][0:1])
            np.testing.assert_array_equal(output_single["emb"], output_batch["emb"][0:1])
            np.testing.assert_array_equal(output_single["indices"], output_batch["indices"][0:1])
        else:
            # Handle tuple return - (z_q, z, emb, indices) + recon from full forward pass
            np.testing.assert_array_equal(output_single[0], output_batch[0][0:1])  # z_q
            np.testing.assert_array_equal(output_single[1], output_batch[1][0:1])  # z
            np.testing.assert_array_equal(output_single[2], output_batch[2][0:1])  # emb
            np.testing.assert_array_equal(output_single[3], output_batch[3][0:1])  # indices
[content] tests/test_tokenizer_reproducibility.py (python)
        # Full tokenizer returns a dictionary
        if isinstance(output1, dict):
            for key in output1.keys():
                np.testing.assert_array_equal(output1[key], output2[key])
        else:
            # Handle tuple return
            for o1, o2 in zip(output1, output2):
                np.testing.assert_array_equal(o1, o2)
[content] tests/test_tokenizer_reproducibility.py (python)
        # vq_encode returns a dictionary
        if isinstance(output1, dict):
            for key in output1.keys():
                np.testing.assert_array_equal(output1[key], output2[key])
        else:
            # Handle tuple return
            for o1, o2 in zip(output1, output2):
                np.testing.assert_array_equal(o1, o2)
[content] tests/test_tokenizer_reproducibility.py (python)
        # Calculate number of patches per video
        H, W = video1.shape[2:4]
        H_pad = H + (-H % self.patch_size)
        W_pad = W + (-W % self.patch_size)
        num_patches = (H_pad // self.patch_size) * (W_pad // self.patch_size)
        patches_per_video = self.seq_len * num_patches
[content] tests/test_tokenizer_reproducibility.py (python)
        # Calculate number of patches per video
        H, W = video1.shape[2:4]
        H_pad = H + (-H % self.patch_size)
        W_pad = W + (-W % self.patch_size)
        num_patches = (H_pad // self.patch_size) * (W_pad // self.patch_size)
        patches_per_video = self.seq_len * num_patches
        
        # Compare the first video in the batch with the single video output
[content] tests/test_tokenizer_reproducibility.py (python)
    def test_training_rng_reproducibility_issue(self):
        """Test that demonstrates the RNG reproducibility issue in training."""
        class MockArgs:
            def __init__(self):
                self.vq_beta = 0.25
                self.num_latents = 64
                self.log_gradients = False
        
        args = MockArgs()
        
        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,
        )
        
        params = tokenizer.init(self.rng, {"videos": self.test_videos})
        
        class MockState:
            def __init__(self, apply_fn):
                self.apply_fn = apply_fn
        
        state = MockState(tokenizer.apply)
        
        # Simulate two different training runs with different RNG evolution
        # Run 1: RNG splits 5 times before reaching our step
        rng1 = jax.random.PRNGKey(self.seed)
        for i in range(5):
            rng1, _ = jax.random.split(rng1)
        
        # Run 2: RNG splits 3 times before reaching our step (different data ordering)
        rng2 = jax.random.PRNGKey(self.seed)
        for i in range(3):
            rng2, _ = jax.random.split(rng2)
        
        # Now both runs reach the same training step but with different RNG states
        inputs1 = {"videos": self.test_videos, "rng": rng1}
        inputs2 = {"videos": self.test_videos, "rng": rng2}
        
        loss1, (recon1, metrics1) = tokenizer_loss_fn(params, state, inputs1)
        loss2, (recon2, metrics2) = tokenizer_loss_fn(params, state, inputs2)
        
        # These should be different due to different RNG states!
        with self.assertRaises(AssertionError):
            np.testing.assert_array_equal(loss1, loss2)
        
        # But if we use the same RNG state, they should be identical
        rng_same = jax.random.PRNGKey(123)
        inputs_same1 = {"videos": self.test_videos, "rng": rng_same}
        inputs_same2 = {"videos": self.test_videos, "rng": rng_same}
        
        loss_same1, (recon_same1, metrics_same1) = tokenizer_loss_fn(params, state, inputs_same1)
        loss_same2, (recon_same2, metrics_same2) = tokenizer_loss_fn(params, state, inputs_same2)
        
        np.testing.assert_array_equal(loss_same1, loss_same2)