batched padded encoder+projector

Author: hengtaoguo

src/maxtext/layers/embeddings.py

@@ -1424,7 +1424,12 @@ def __call__(
     batch_size = inputs.shape[0]
     max_patches = num_frames * height * width
 
-    if valid_grid is None:
+    if valid_grid is not None:
+      if isinstance(valid_grid, (tuple, list)):
+        valid_grid = jnp.array([valid_grid], dtype=jnp.int32)
+      elif valid_grid.ndim == 1:
+        valid_grid = valid_grid[None, :]
+    else:
       valid_grid = jnp.tile(
           jnp.array([[num_frames, height, width]], dtype=jnp.int32),
           (batch_size, 1),
@@ -1688,6 +1693,10 @@ def __call__(
       Interpolated positional embeddings of shape [batch, num_frames * height * width, hidden_size]
     """
     if video_grid_thw is not None:
+      if isinstance(video_grid_thw, (tuple, list)):
+        video_grid_thw = jnp.array([video_grid_thw], dtype=jnp.int32)
+      elif video_grid_thw.ndim == 1:
+        video_grid_thw = video_grid_thw[None, :]
       batch_size = video_grid_thw.shape[0]
     elif attention_mask is not None:
       batch_size = attention_mask.shape[0]

tests/unit/qwen3_omni_layers_test.py

@@ -37,7 +37,7 @@
     Qwen3OmniMoeVisionRotaryEmbedding as JaxQwen3OmniMoeVisionRotaryEmbedding,
 )
 from maxtext.layers.decoders import deepstack_process
-from maxtext.layers.encoders import AudioEncoder
+from maxtext.layers.encoders import AudioEncoder, VisionEncoder
 from maxtext.multimodal.processor_qwen3_omni import (
     maybe_pad_video_values_to_max_grid,
     preprocess_video,
@@ -868,6 +868,100 @@ def test_padded_video_valid_outputs_match_unpadded_batched(self):
         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)
+    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],
+        out_hidden_size_for_vit=24,
+        vision_encoder_block=common_types.VisionEncoderBlockType.QWEN3_OMNI,
+    )
+    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)
+
+    # Instantiate VisionEncoder (wraps encoder + projector)
+    vision_encoder = VisionEncoder(config=config, mesh=self.mesh, rngs=nnx.Rngs(42))
+
+    # Run individual unpadded
+    raw_output_1, _ = vision_encoder(raw_video_1)
+    raw_output_2, _ = vision_encoder(raw_video_2)
+
+    # Run batched padded
+    batched_padded_output, _ = vision_encoder(
+        batched_padded_video,
+        input_masks=batched_video_mask,
+        video_grid_thw=batched_grid_thw,
+    )
+
+    # Downsample masks to projected tokens
+    projected_mask_1 = mm_processor.downsample_video_mask_to_tokens(video_mask_1, config)
+    projected_mask_2 = mm_processor.downsample_video_mask_to_tokens(video_mask_2, config)
+
+    # Verify equivalence for batch item 0
+    np.testing.assert_allclose(
+        np.asarray(raw_output_1).reshape(-1, config.out_hidden_size_for_vit),
+        np.asarray(batched_padded_output[0])[np.asarray(projected_mask_1[0]).astype(bool)],
+        rtol=1e-4,
+        atol=1e-4,
+    )
+
+    # Verify equivalence for batch item 1
+    np.testing.assert_allclose(
+        np.asarray(raw_output_2).reshape(-1, config.out_hidden_size_for_vit),
+        np.asarray(batched_padded_output[1])[np.asarray(projected_mask_2[0]).astype(bool)],
+        rtol=1e-4,
+        atol=1e-4,
+    )
+
 
 class TestDeepstackProcess(unittest.TestCase):
   """Tests for deepstack_process.