Fix gemma4 audio-encoder audio_features export rank (#570)

The gemma4 audio dataflow edge
`audio_encoder.audio_features -> embedding.audio_features` was rank-3 on the
producer side but rank-2 on the consumer side, so onnx-genai pipeline admission
rejected every gemma4 audio (E2B/E4B) package:

    'audio_encoder.audio_features -> embedding.audio_features' has incompatible
    ranks: producer rank 3, consumer rank 2.

HF selects `audio_features[audio_mask]` (a boolean-mask gather) that both drops
padding frames and flattens the batch axis to rank-2
`[num_valid_frames, text_hidden_size]` before scattering rows into the
language-model embeddings.

Move that mask-selection into `_Gemma4AudioEncoderModel.forward`: after the
projector, flatten `[B, T//4, H] -> [-1, H]` and reuse the shared bf16-safe
`_dtype_safe_compress` helper to emit rank-2 `[num_valid_frames, H]`, matching
the `embedding` consumer and the encoder-free `_Gemma4UnifiedAudioEmbedderModel`.
The task builder `_build_audio` no longer duplicates the flatten+Compress. This
covers all gemma4 Conformer-audio variants since they share the module.

Add `TestGemma4AudioEncoderRank` locking the rank-2 contract: producer rank ==
consumer rank == 2 (trailing dim == hidden_size) and a Compress node is present
so the reduction is a real mask-select, not a bare reshape.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Author: justinchuby

src/mobius/models/gemma4.py

@@ -2639,9 +2639,19 @@ class _Gemma4AudioEncoderModel(nn.Module):
 
     Inputs:
     - ``input_features [B, T, input_size]``: mel-spectrogram features
+    - ``input_features_mask [B, T]``: BOOL mask, ``True`` for valid frames
 
     Output:
-    - ``audio_features [B, T//4, text_hidden_size]``: projected audio tokens
+    - ``audio_features [num_valid_frames, text_hidden_size]``: projected audio
+      tokens with padding frames stripped so output rows align 1:1 with audio
+      placeholder tokens.  This mirrors HF, which selects
+      ``audio_features[audio_mask]`` (a rank-2 gather) before scattering the
+      rows into the language-model embeddings.  Emitting rank-2 here keeps the
+      producer rank aligned with the ``embedding`` sub-model's rank-2
+      ``audio_features`` consumer.
+    - ``downsampled_mask [B, T//4]``: BOOL mask over the subsampled time axis,
+      returned for the ``audio_features_mask`` diagnostic output.  ``None`` when
+      no ``input_features_mask`` was supplied.
     """
 
     def __init__(self, config: Gemma4Config):
@@ -2701,7 +2711,23 @@ def forward(
         rms = op.Sqrt(op.Add(mean_sq, eps))
         audio_features = op.CastLike(op.Div(x_f32, rms), audio_features)
         # → projector → [B, T//4, text_hidden_size]
-        return self.projector(op, audio_features), downsampled_mask
+        audio_features = self.projector(op, audio_features)
+        if downsampled_mask is None:
+            # No validity mask (e.g. eager unit call): emit the rank-3 tensor
+            # unchanged.  Export always supplies a mask, so the rank-2 path below
+            # is the one that reaches the ``embedding`` consumer.
+            return audio_features, None
+        # Select valid frames so the output rows align 1:1 with audio placeholder
+        # tokens, matching HF's ``audio_features[audio_mask]``.  Flatten the batch
+        # and subsampled-time axes, then compress by the downsampled mask to get
+        # rank-2 ``[num_valid_frames, text_hidden_size]`` — the rank the
+        # ``embedding`` sub-model consumes.
+        flat_features = op.Reshape(
+            audio_features, op.Constant(value_ints=[-1, self.config.hidden_size])
+        )
+        flat_mask = op.Reshape(downsampled_mask, op.Constant(value_ints=[-1]))
+        selected = _dtype_safe_compress(op, flat_features, flat_mask, axis=0)
+        return selected, downsampled_mask
 
 
 # ---------------------------------------------------------------------------

src/mobius/models/gemma4_test.py

@@ -284,6 +284,86 @@ def test_no_block_sequence_ids_when_causal(self):
         assert "input_ids" not in dec_inputs
 
 
+def _tiny_gemma4_audio_config(**overrides) -> Gemma4Config:
+    """Minimal gemma4 config with a Conformer audio encoder (E2B-style)."""
+    from mobius._configs import Gemma4AudioConfig
+
+    base = dict(
+        image_token_id=255,
+        audio=Gemma4AudioConfig(
+            input_size=32,
+            hidden_size=16,
+            num_layers=1,
+            subsampling_conv_channels=[8, 8],
+            output_proj_dims=16,
+            audio_token_id=254,
+        ),
+    )
+    base.update(overrides)
+    return _tiny_gemma4_config(**base)
+
+
+class TestGemma4AudioEncoderRank:
+    """Lock the audio dataflow-edge rank contract (regression for #570).
+
+    The exported ``audio_encoder`` must emit ``audio_features`` at the SAME rank
+    the ``embedding`` sub-model consumes it (rank-2
+    ``[num_valid_frames, text_hidden_size]``), mirroring HF's
+    ``audio_features[audio_mask]`` selection.  A rank-3 producer against a
+    rank-2 consumer is rejected at onnx-genai pipeline admission.
+    """
+
+    def test_audio_features_producer_consumer_ranks_match(self):
+        from mobius.models.gemma4 import Gemma4Model
+        from mobius.tasks._gemma4 import Gemma4Task
+
+        config = _tiny_gemma4_audio_config()
+        pkg = Gemma4Task().build(Gemma4Model(config), config)
+
+        audio_out = {o.name: o for o in pkg["audio_encoder"].graph.outputs}
+        emb_in = {i.name: i for i in pkg["embedding"].graph.inputs}
+
+        producer = audio_out["audio_features"]
+        consumer = emb_in["audio_features"]
+        assert producer.shape is not None
+        assert consumer.shape is not None
+        producer_rank = len(producer.shape)
+        consumer_rank = len(consumer.shape)
+
+        # Non-vacuous: assert the concrete rank-2 contract on BOTH ends and that
+        # they agree, not merely that export ran.
+        assert producer_rank == 2, (
+            f"audio_encoder audio_features must be rank-2 "
+            f"[num_valid_frames, hidden]; got rank {producer_rank} {producer.shape}"
+        )
+        assert consumer_rank == 2, (
+            f"embedding audio_features input must be rank-2; "
+            f"got rank {consumer_rank} {consumer.shape}"
+        )
+        assert producer_rank == consumer_rank
+        # Trailing dim must be the text hidden size on both ends.
+        assert producer.shape[1] == config.hidden_size
+        assert consumer.shape[1] == config.hidden_size
+
+    def test_audio_encoder_selects_valid_frames_with_compress(self):
+        """The rank reduction is a real mask-select (HF ``audio_features[audio_mask]``).
+
+        Guards against a "fix" that reshapes B*T//4 into rank-2 without dropping
+        padding frames, which would misalign rows with placeholder tokens.
+        """
+        from mobius.models.gemma4 import Gemma4Model
+        from mobius.tasks._gemma4 import Gemma4Task
+
+        config = _tiny_gemma4_audio_config()
+        pkg = Gemma4Task().build(Gemma4Model(config), config)
+
+        audio_graph = pkg["audio_encoder"].graph
+        assert any(n.op_type == "Compress" for n in audio_graph), (
+            "audio_encoder must Compress by the downsampled validity mask so "
+            "output rows align 1:1 with audio placeholder tokens"
+        )
+
+
 def _tiny_gemma4_unified_config(**overrides) -> Gemma4Config:
     """Minimal gemma4_unified config (dense decoder + encoder-free embedders)."""
     from mobius._configs import Gemma4AudioConfig, VisionConfig

src/mobius/tasks/_gemma4.py

@@ -666,19 +666,9 @@ def _build_audio(
         )
         if downsampled_mask is None:
             raise ValueError("Gemma4 audio encoder must return a downsampled validity mask")
-        flattened_features = op.Reshape(
-            audio_features,
-            op.Constant(value_ints=[-1, config.hidden_size]),
-        )
-        flattened_mask = op.Reshape(
-            downsampled_mask,
-            op.Constant(value_ints=[-1]),
-        )
-        flattened_features_f32 = op.Cast(flattened_features, to=ir.DataType.FLOAT)
-        audio_features = op.CastLike(
-            op.Compress(flattened_features_f32, flattened_mask, axis=0),
-            flattened_features,
-        )
+        # The audio encoder module already strips padding frames and flattens the
+        # batch axis, so ``audio_features`` is rank-2 ``[num_valid, hidden_size]``
+        # here — the rank the ``embedding`` sub-model consumes.
         builder.add_output(audio_features, "audio_features")
         builder.add_output(downsampled_mask, "audio_features_mask")