revert

Author: hengtaoguo

tests/unit/qwen3_omni_layers_test.py

@@ -775,99 +775,6 @@ def test_padded_video_valid_outputs_match_unpadded(self):
         atol=1e-4,
     )
 
-  def test_padded_video_valid_outputs_match_unpadded_batched(self):
-    """Padded tokens do not change valid outputs anywhere in the ViT for batched inputs."""
-    raw_grid_1 = (2, 2, 2)
-    raw_grid_2 = (1, 4, 4)
-    max_grid = (3, 4, 4)
-    config = pyconfig.initialize(
-        ["", base_config_path],
-        model_name="qwen3-omni-30b-a3b",
-        attention="dot_product",
-        attention_type="full",
-        dtype="float32",
-        dtype_mm="float32",
-        weight_dtype="float32",
-        override_model_config=True,
-        attention_for_vit="dot_product",
-        hidden_size_for_vit=16,
-        num_attention_heads_for_vit=2,
-        intermediate_size_for_vit=32,
-        num_hidden_layers_for_vit=1,
-        deepstack_visual_indexes_for_vit=[],
-        video_max_grid_t=max_grid[0],
-        video_max_grid_h=max_grid[1],
-        video_max_grid_w=max_grid[2],
-    )
-    patch_size = config.patch_size_for_vit
-    temporal_patch_size = config.temporal_patch_size_for_vit
-
-    # Generate raw videos
-    raw_shape_1 = (
-        1,
-        config.num_channels_for_vit,
-        raw_grid_1[0] * temporal_patch_size,
-        raw_grid_1[1] * patch_size,
-        raw_grid_1[2] * patch_size,
-    )
-    raw_shape_2 = (
-        1,
-        config.num_channels_for_vit,
-        raw_grid_2[0] * temporal_patch_size,
-        raw_grid_2[1] * patch_size,
-        raw_grid_2[2] * patch_size,
-    )
-    raw_video_1, _ = create_random_jax_torch(*raw_shape_1)
-    raw_video_2, _ = create_random_jax_torch(*raw_shape_2)
-
-    # Pad them individually
-    padded_video_1, _, video_mask_1 = maybe_pad_video_values_to_max_grid(
-        np.asarray(raw_video_1), np.asarray([raw_grid_1]), config
-    )
-    padded_video_2, _, video_mask_2 = maybe_pad_video_values_to_max_grid(
-        np.asarray(raw_video_2), np.asarray([raw_grid_2]), config
-    )
-
-    # Batch them
-    batched_padded_video = jnp.concatenate([padded_video_1, padded_video_2], axis=0)
-    batched_video_mask = jnp.concatenate([video_mask_1, video_mask_2], axis=0)
-    batched_grid_thw = jnp.array([raw_grid_1, raw_grid_2], dtype=jnp.int32)
-
-    encoder = JaxQwen3OmniMoeVisionEncoder(config=config, mesh=self.mesh, rngs=nnx.Rngs(42))
-
-    # Run individual unpadded
-    raw_output_1, _ = encoder(raw_video_1)
-    raw_output_2, _ = encoder(raw_video_2)
-
-    # Run batched padded
-    batched_padded_output, _ = encoder(
-        batched_padded_video,
-        video_mask=batched_video_mask,
-        video_grid_thw=batched_grid_thw,
-    )
-
-    # Verify equivalence for batch item 0
-    patch_mask_1 = (
-        np.asarray(video_mask_1).reshape(1, -1, temporal_patch_size * patch_size * patch_size).max(-1).astype(bool)
-    )
-    np.testing.assert_allclose(
-        np.asarray(raw_output_1).reshape(-1, config.hidden_size_for_vit),
-        np.asarray(batched_padded_output[0])[np.asarray(patch_mask_1)[0]],
-        rtol=1e-4,
-        atol=1e-4,
-    )
-
-    # Verify equivalence for batch item 1
-    patch_mask_2 = (
-        np.asarray(video_mask_2).reshape(1, -1, temporal_patch_size * patch_size * patch_size).max(-1).astype(bool)
-    )
-    np.testing.assert_allclose(
-        np.asarray(raw_output_2).reshape(-1, config.hidden_size_for_vit),
-        np.asarray(batched_padded_output[1])[np.asarray(patch_mask_2)[0]],
-        rtol=1e-4,
-        atol=1e-4,
-    )
-
   def test_padded_video_projected_outputs_match_unpadded_batched(self):
     """Padded tokens do not change valid projected outputs for batched inputs."""
     raw_grid_1 = (2, 2, 2)