amend

Author: hengtaoguo

src/maxtext/layers/embeddings.py

@@ -1407,7 +1407,7 @@ def __call__(
 
     Args:
       inputs: Input tensor of shape [B, S, N, head_dim] (batch, sequence, heads, head_dim)
-             where T=num_frames, H=height, W=width (all static)
+             where sequence length S = num_frames * height * width (static maximum).
       num_frames: Number of temporal frames (static)
       height: Height in patches (static)
       width: Width in patches (static)

tests/unit/qwen3_omni_layers_test.py

@@ -53,6 +53,7 @@
     Qwen3OmniMoeVisionProjector as JaxQwen3OmniMoeVisionProjector,
 )
 from maxtext.multimodal import processor as mm_processor
+from maxtext.multimodal import utils as mm_utils
 from tests.utils.multimodal_test_utils import (
     assert_all_close_jax_torch,
     copy_attention_weights_to_maxtext,
@@ -869,6 +870,67 @@ def test_padded_video_projected_outputs_match_unpadded_batched(self):
         atol=1e-4,
     )
 
+    # --- Test Merging into Text Sequence ---
+    emb_dim = config.out_hidden_size_for_vit
+
+    # Example 1: 2 placeholders at indices 5, 6.
+    text_embeddings_1 = jnp.arange(1 * 20 * emb_dim, dtype=jnp.float32).reshape(1, 20, emb_dim)
+    placeholder_mask_1 = jnp.zeros((1, 20), dtype=jnp.int32)
+    placeholder_mask_1 = placeholder_mask_1.at[:, 5:7].set(1)
+
+    # Example 2: 4 placeholders at indices 5, 6, 7, 8.
+    text_embeddings_2 = jnp.arange(1 * 20 * emb_dim, dtype=jnp.float32).reshape(1, 20, emb_dim) + 1000.0
+    placeholder_mask_2 = jnp.zeros((1, 20), dtype=jnp.int32)
+    placeholder_mask_2 = placeholder_mask_2.at[:, 5:9].set(1)
+
+    # Batch text inputs
+    batched_text_embeddings = jnp.concatenate([text_embeddings_1, text_embeddings_2], axis=0)
+    batched_placeholder_mask = jnp.concatenate([placeholder_mask_1, placeholder_mask_2], axis=0)
+    batched_projected_mask = jnp.concatenate([projected_mask_1, projected_mask_2], axis=0)
+
+    # Run batched merge
+    batched_merged = mm_utils.merge_mm_embeddings(
+        text_embeddings=batched_text_embeddings,
+        multimodal_embeddings=batched_padded_output,
+        mask=batched_placeholder_mask,
+        token_masks=batched_projected_mask,
+    )
+
+    # Run individual unpadded merges
+    token_mask_1 = jnp.ones((1, 2), dtype=jnp.int32)
+    merged_1 = mm_utils.merge_mm_embeddings(
+        text_embeddings=text_embeddings_1,
+        multimodal_embeddings=raw_output_1,
+        mask=placeholder_mask_1,
+        token_masks=token_mask_1,
+    )
+
+    # Example 2 has 4 valid blocks
+    token_mask_2 = jnp.ones((1, 4), dtype=jnp.int32)
+    merged_2 = mm_utils.merge_mm_embeddings(
+        text_embeddings=text_embeddings_2,
+        multimodal_embeddings=raw_output_2,
+        mask=placeholder_mask_2,
+        token_masks=token_mask_2,
+    )
+
+    # Verify that batched merge matches individual unpadded merges
+    # Batch item 0
+    np.testing.assert_allclose(
+        np.asarray(batched_merged[0]),
+        np.asarray(merged_1[0]),
+        rtol=1e-4,
+        atol=1e-4,
+    )
+
+    # Batch item 1
+    np.testing.assert_allclose(
+        np.asarray(batched_merged[1]),
+        np.asarray(merged_2[0]),
+        rtol=1e-4,
+        atol=1e-4,
+    )
+
 
 class TestDeepstackProcess(unittest.TestCase):
   """Tests for deepstack_process.