tests: fix L4/L5 e2e harness for Gemma4 per_layer_inputs + BOOL audio mask

After #296 moved Gemma4's per-layer input computation from the text
decoder into the embedding sub-model, the embedding model now emits a
second output (per_layer_inputs) and the decoder accepts it as a
required input. The e2e harness only wired the first embedding output
(inputs_embeds) into the decoder, so every multi-model Gemma4 path
(text-only on multi-model, VL prefill, VL generation, speech-language
prefill, speech-language generation) failed with:

  ValueError: Required inputs (['per_layer_inputs']) are missing from
  input feed (['inputs_embeds', ...])

Fix by passing any extra embedding outputs through to the decoder by
name. This is generic — for models without per_layer_inputs the extra
loop iteration just no-ops.

Separately, the speech-language audio encoder builds Gemma4's
input_features_mask as tensor(bool), but the harness unconditionally
cast feature-extractor outputs to float32 (and constructed the fallback
all-True mask as np.bool_ that then crashed ort_easy's DLPack path,
which has no bool type code). Fix by:

- Honoring the session's declared input dtype for each audio-encoder
  input (BOOL stays BOOL, FLOAT becomes float32, etc.).
- Routing bool numpy arrays through OrtValue.ortvalue_from_numpy
  directly in OnnxModelSession, bypassing ort_easy's DLPack-first path.

Verified locally on H200:
  L4 text-generation/gemma-4-e2b           PASS
  L4 image-text-to-text/gemma-4-e2b-it     PASS
  L4 speech-language/gemma-4-e2b-it-audio  PASS
  L5 image-text-to-text/gemma-4-e2b-it     PASS
  L5 speech-language/gemma-4-e2b-it-audio  PASS
(L5 text-generation/gemma-4-e2b is gated by integration markers and
skipped under fast runs.)

ruff check + format pass.

Signed-off-by: justinchuby <11205048+justinchuby@users.noreply.github.com>

Author: justinchuby

src/mobius/_testing/ort_inference.py

@@ -66,14 +66,18 @@ def _ort_value_to_numpy(value: ort.OrtValue) -> np.ndarray:
 
 
 def _numpy_to_ort_value(value: np.ndarray) -> ort.OrtValue:
-    """Convert a NumPy array to an OrtValue, supporting ml_dtypes (bf16, etc).
+    """Convert a NumPy array to an OrtValue, supporting ml_dtypes and bool.
 
     ``onnxruntime_easy.ort_value`` prefers the DLPack path, which is broken
-    for ml_dtypes scalars (NumPy's __dlpack__ rejects non-standard dtypes).
-    Route ml_dtypes arrays through ``ortvalue_from_numpy_with_onnx_type``
-    instead.
+    for ml_dtypes scalars (NumPy's __dlpack__ rejects non-standard dtypes)
+    and for bool (DLPack has no native bool type code). Route those arrays
+    through the explicit numpy-to-OrtValue APIs instead.
     """
     if isinstance(value, np.ndarray):
+        if value.dtype == np.bool_:
+            # ORT supports tensor(bool) natively but DLPack does not, so
+            # bypass ort_easy's DLPack-first path.
+            return ort.OrtValue.ortvalue_from_numpy(np.ascontiguousarray(value))
         try:
             import ml_dtypes
         except ImportError:  # pragma: no cover

tests/e2e_golden_test.py

@@ -666,6 +666,10 @@ def _run_vision_language_prefill(
                     )
                 else:
                     dec_feeds[name] = np.arange(seq_len, dtype=np.int64).reshape(1, -1)
+            elif name in emb_out:
+                # Gemma4 (when hidden_size_per_layer_input > 0) emits a second
+                # embedding output ``per_layer_inputs`` that the decoder needs.
+                dec_feeds[name] = emb_out[name]
         outputs = dec_session.run(dec_feeds)
     finally:
         dec_session.close()
@@ -862,6 +866,12 @@ def _run_vl_generation(
         else:
             next_decode_pos = prompt_seq_len
 
+        # Wire extra embedding outputs the decoder expects by name
+        # (e.g. Gemma4 ``per_layer_inputs``).
+        for name in dec_session.input_names:
+            if name not in dec_feeds and name in emb_out:
+                dec_feeds[name] = emb_out[name]
+
         prefill_out = dec_session.run(dec_feeds)
         logits = prefill_out["logits"]
         next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True).astype(np.int64)
@@ -918,6 +928,12 @@ def _run_vl_generation(
                 else:
                     step_feeds["position_ids"] = np.array([[next_decode_pos]], dtype=np.int64)
 
+            # Wire extra embedding outputs the decoder expects by name
+            # (e.g. Gemma4 ``per_layer_inputs``).
+            for name in dec_session.input_names:
+                if name not in step_feeds and name in step_emb_out:
+                    step_feeds[name] = step_emb_out[name]
+
             step_out = dec_session.run(step_feeds)
             logits = step_out["logits"]
             next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True).astype(np.int64)
@@ -1069,6 +1085,11 @@ def _run_text_only_multimodel_prefill(
                 dec_feeds[name] = np.ones((1, seq_len), dtype=np.int64)
             elif name == "position_ids":
                 dec_feeds[name] = np.arange(seq_len, dtype=np.int64).reshape(1, -1)
+            elif name in emb_out:
+                # Gemma4 (when hidden_size_per_layer_input > 0) emits a second
+                # embedding output ``per_layer_inputs`` that the decoder needs.
+                # Wire any extra embedding outputs through by name.
+                dec_feeds[name] = emb_out[name]
         outputs = dec_session.run(dec_feeds)
     finally:
         dec_session.close()
@@ -1200,6 +1221,11 @@ def _run_phi4mm_multimodal_prefill(
             "position_ids": np.arange(seq_len, dtype=np.int64).reshape(1, -1),
             **kv_cache,
         }
+        # Wire any extra embedding outputs the decoder expects by name
+        # (e.g. Gemma4 ``per_layer_inputs``).
+        for name in dec_session.input_names:
+            if name not in dec_feeds and name in emb_out:
+                dec_feeds[name] = emb_out[name]
         outputs = dec_session.run(dec_feeds)
     finally:
         dec_session.close()
@@ -1252,7 +1278,11 @@ def _run_speech_language_prefill(
         audio_feeds: dict[str, np.ndarray] = {}
         for name in audio_session.input_names:
             if name in audio_processed:
-                audio_feeds[name] = audio_processed[name].astype(np.float32)
+                # Cast to the session's declared dtype (e.g. input_features_mask
+                # is BOOL on Gemma4 audio encoder; the HF feature extractor may
+                # emit it as float or int).
+                target_dtype = audio_session.get_input_dtype(name) or np.float32
+                audio_feeds[name] = audio_processed[name].astype(target_dtype)
             elif name == "input_features" and "input_features" in audio_processed:
                 audio_feeds[name] = audio_processed["input_features"].astype(np.float32)
         # Provide all-True mask for single-clip inference when the model
@@ -1264,7 +1294,8 @@ def _run_speech_language_prefill(
             and "input_features" in audio_feeds
         ):
             feats = audio_feeds["input_features"]
-            audio_feeds["input_features_mask"] = np.ones(feats.shape[:2], dtype=np.bool_)
+            mask_dtype = audio_session.get_input_dtype("input_features_mask") or np.bool_
+            audio_feeds["input_features_mask"] = np.ones(feats.shape[:2], dtype=mask_dtype)
         audio_out = audio_session.run(audio_feeds)
     finally:
         audio_session.close()
@@ -1352,6 +1383,10 @@ def _run_speech_language_prefill(
                     ndims = pos_shape[0] if isinstance(pos_shape[0], int) else 3
                     pos = np.tile(pos, (ndims, 1, 1))
                 dec_feeds[name] = pos
+            elif name in emb_out:
+                # Extra embedding outputs the decoder expects by name
+                # (e.g. Gemma4 ``per_layer_inputs``).
+                dec_feeds[name] = emb_out[name]
         outputs = dec_session.run(dec_feeds)
     finally:
         dec_session.close()
@@ -1701,7 +1736,8 @@ def _run_speech_language_generation(
         audio_feeds: dict[str, np.ndarray] = {}
         for name in audio_session.input_names:
             if name in audio_processed:
-                audio_feeds[name] = audio_processed[name].astype(np.float32)
+                target_dtype = audio_session.get_input_dtype(name) or np.float32
+                audio_feeds[name] = audio_processed[name].astype(target_dtype)
             elif name == "input_features" and "input_features" in audio_processed:
                 audio_feeds[name] = audio_processed["input_features"].astype(np.float32)
         # Provide all-True mask for single-clip inference when the model
@@ -1713,7 +1749,8 @@ def _run_speech_language_generation(
             and "input_features" in audio_feeds
         ):
             feats = audio_feeds["input_features"]
-            audio_feeds["input_features_mask"] = np.ones(feats.shape[:2], dtype=np.bool_)
+            mask_dtype = audio_session.get_input_dtype("input_features_mask") or np.bool_
+            audio_feeds["input_features_mask"] = np.ones(feats.shape[:2], dtype=mask_dtype)
         audio_out = audio_session.run(audio_feeds)
     finally:
         audio_session.close()
@@ -1798,6 +1835,12 @@ def _run_speech_language_generation(
                 pos = np.tile(pos, (ndims_pos, 1, 1))
             dec_feeds["position_ids"] = pos
 
+        # Wire extra embedding outputs the decoder expects by name
+        # (e.g. Gemma4 ``per_layer_inputs``).
+        for name in dec_session.input_names:
+            if name not in dec_feeds and name in emb_out:
+                dec_feeds[name] = emb_out[name]
+
         prefill_out = dec_session.run(dec_feeds)
         logits = prefill_out["logits"]
         next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True).astype(np.int64)
@@ -1841,6 +1884,12 @@ def _run_speech_language_generation(
                 else:
                     step_feeds["position_ids"] = np.array([[past_seq_len]], dtype=np.int64)
 
+            # Wire extra embedding outputs the decoder expects by name
+            # (e.g. Gemma4 ``per_layer_inputs``).
+            for name in dec_session.input_names:
+                if name not in step_feeds and name in step_emb_out:
+                    step_feeds[name] = step_emb_out[name]
+
             step_out = dec_session.run(step_feeds)
             logits = step_out["logits"]
             next_token = np.argmax(logits[:, -1, :], axis=-1, keepdims=True).astype(np.int64)