Add Qwen3-Omni layers test to CI

.github/workflows/run_tests_coordinator.yml

@@ -209,7 +209,7 @@ jobs:
             "gpu-integration": "--ignore=tests/post_training",
             "cpu-unit": "--ignore=tests/post_training",
             "cpu-post-training-unit": "",
-            "cpu-torch-reference": "-o addopts= -rf --import-mode=importlib --strict-markers tests/unit/gemma4_layers_test.py tests/unit/gemma4_small_layers_test.py tests/unit/qwen3_next_vs_reference_test.py tests/unit/qwen3_5_layers_test.py"
+            "cpu-torch-reference": "-o addopts= -rf --import-mode=importlib --strict-markers tests/unit/gemma4_layers_test.py tests/unit/gemma4_small_layers_test.py tests/unit/qwen3_next_vs_reference_test.py tests/unit/qwen3_5_layers_test.py tests/unit/qwen3_omni_layers_test.py"
           }')[inputs.flavor] }}
         ${{ inputs.additional_pytest_args }}
 

tests/unit/qwen3_omni_layers_test.py

@@ -25,6 +25,7 @@
 import jax
 import jax.numpy as jnp
 from jax.sharding import Mesh
+import pytest
 from maxtext.configs import pyconfig
 from maxtext.common import common_types
 from maxtext.utils.globals import MAXTEXT_REPO_ROOT
@@ -62,7 +63,6 @@
     copy_patch_embed_weights,
     copy_patch_merger_weights,
     copy_vision_encoder_weights,
-    create_block_diagonal_attention_mask,
     create_random_jax_torch,
 )
 import numpy as np
@@ -777,6 +777,7 @@ def test_hidden_states_unchanged_without_visual_tokens(self):
     np.testing.assert_allclose(np.array(result), hidden_np, rtol=1e-6, atol=1e-6)
 
 
+@pytest.mark.skip(reason="Requires decord, which may not be installed in remote CI runners.")
 class TestQwen3OmniPreprocessing(unittest.TestCase):
   """Test MaxText Qwen3 Omni preprocessor against HuggingFace reference."""
 
@@ -1003,17 +1004,12 @@ def _test_encoder_layer_with_batch_size(self, batch_size):
 
     jax_input, torch_input_3d = create_random_jax_torch(batch_size, seq_len, hidden_size)
 
-    # PyTorch forward pass - expects 2D input (total_seq_len, hidden_dim) with cu_seqlens
-    torch_input_2d = torch_input_3d.reshape(-1, hidden_size)
-
-    # Create cu_seqlens for PyTorch (cumulative sequence lengths for each batch)
-    # For batch_size=2, seq_len=12: [0, 12, 24] indicates two sequences of length 12 each
-    cu_seqlens = torch.tensor([i * seq_len for i in range(batch_size + 1)], dtype=torch.int32)
-
-    attention_mask = create_block_diagonal_attention_mask(cu_seqlens, torch_input_2d.dtype)
-
-    torch_output_1d = torch_layer(torch_input_2d, cu_seqlens=cu_seqlens, attention_mask=attention_mask)[0]
-    torch_output = torch_output_1d.reshape(batch_size, seq_len, hidden_size)
+    # PyTorch audio layers take 2D packed input. Run each batch item separately
+    # to match MaxText's batched attention without relying on HF's old mask API.
+    cu_seqlens = torch.tensor([0, seq_len], dtype=torch.int32)
+    torch_output = torch.stack(
+        [torch_layer(torch_input_3d[i], cu_seqlens=cu_seqlens)[0] for i in range(batch_size)], dim=0
+    )
 
     jax_output = maxtext_layer(jax_input, deterministic=True)
 
@@ -1144,18 +1140,15 @@ def test_audio_encoder_matches_torch(self):
     torch_after_pos = torch_conv_out + torch_pos_emb
 
     # Run through encoder layers + layernorm (but not projector)
-    # Process all chunks together
+    # Process chunks separately, matching MaxText's (batch * chunks, seq, hidden) attention shape.
     seq_len_per_chunk = torch_after_pos.shape[1]
-    cu_seqlens = torch.tensor([i * seq_len_per_chunk for i in range(num_chunks + 1)], dtype=torch.int32)
-    attention_mask = create_block_diagonal_attention_mask(cu_seqlens, torch_after_pos.dtype)
-
-    # Flatten: (num_chunks, seq_len_per_chunk, hidden) -> (num_chunks*seq_len_per_chunk, hidden)
-    hidden_state = torch_after_pos.reshape(-1, torch_after_pos.shape[-1])
+    cu_seqlens = torch.tensor([0, seq_len_per_chunk], dtype=torch.int32)
+    hidden_state = torch_after_pos
     for layer in torch_model.layers:
-      hidden_state = layer(hidden_state, cu_seqlens=cu_seqlens, attention_mask=attention_mask)[0]
+      hidden_state = torch.stack([layer(chunk, cu_seqlens=cu_seqlens)[0] for chunk in hidden_state], dim=0)
     hidden_state = torch_model.ln_post(hidden_state)
 
-    # Reshape back: (num_chunks*seq_len_per_chunk, hidden) -> (batch=1, num_chunks*seq_len_per_chunk, hidden)
+    # Reshape back: (num_chunks, seq_len_per_chunk, hidden) -> (batch=1, num_chunks*seq_len_per_chunk, hidden)
     torch_output = hidden_state.reshape(1, num_chunks * seq_len_per_chunk, -1)
 
     # MaxText forward