fix(litertlm): minor code issues from review

- Skip populate_function_gemma_metadata when llm_model_type is an explicit
  override, mirroring litert-torch's override behavior.
- Handle NumPy integer axes (np.int64/32) in _patched_amax.
- Update Falcon xfail reason to current torch.fx
  _ModuleNotInstalledAsSubmoduleError.
- Update Gemma3n xfail reason's shape numbers to current
  [3,8,8,16] strides (1024,8,1,64) -> [192,16].

Author: pctablet505

keras_hub/src/models/falcon/falcon_causal_lm_test.py

@@ -178,8 +178,9 @@ def test_litert_export(self):
     @pytest.mark.xfail(
         strict=True,
         reason=(
-            "torch.export tracing failure: broadcast shape mismatch "
-            "[1,2,1,8] vs [1,2,8,12] at export time (not a runtime gap); "
+            "torch.export tracing failure: "
+            "torch.fx._ModuleNotInstalledAsSubmoduleError in FalconAttention "
+            "(module not registered as a submodule during tracing); "
             "under diagnosis"
         ),
     )

keras_hub/src/models/gemma3n/gemma3n_causal_lm_test.py

@@ -437,10 +437,10 @@ def test_gqa_fit(self):
         strict=True,
         reason=(
             "Gemma3n LiteRT-LM export fails in litert-torch conversion: "
-            "aten.view shape/stride mismatch ([3,64,8] strides (512,1,64) -> "
-            "[192,8]) during forward_prefill decomposition. Pre-existing "
-            "export bug, not a numeric-parity gap; blocks multimodal numeric "
-            "wiring (WS3.3). Under diagnosis."
+            "aten.view shape/stride mismatch ([3,8,8,16] strides "
+            "(1024,8,1,64) -> [192,16]) during forward_prefill decomposition. "
+            "Pre-existing export bug, not a numeric-parity gap; blocks "
+            "multimodal numeric wiring (WS3.3). Under diagnosis."
         ),
     )
     def test_litertlm_export(self):

keras_hub/src/tests/test_case.py

@@ -1183,6 +1183,8 @@ def _verify_litertlm_multimodal_numerics(
             # call (no `sampler_config`), matching what the family tests
             # under verification actually export.
             sampler_config=None,
+            # The harness never passes an `llm_model_type` override.
+            model_type_overridden=False,
         )
         prefill_inputs = _export._build_prefill_inputs(plan)[prefill_seq_len]
 

keras_hub/src/utils/litertlm/export.py

@@ -438,6 +438,7 @@ class ExportPlan:
     tokens_per_image: int | None
     separate_vision_encoder: bool
     sampler_config: object | None
+    model_type_overridden: bool
 
 
 def _build_prefill_inputs(plan):
@@ -800,6 +801,7 @@ def _assemble_bundle(
         vision_cfg=plan.vision_cfg,
         audio_cfg=plan.audio_cfg,
         sampler_config=plan.sampler_config,
+        model_type_overridden=plan.model_type_overridden,
     )
 
     litert_lm_builder = _import_litert_lm_builder()
@@ -1226,6 +1228,7 @@ def export_to_litertlm(
         tokens_per_image=tokens_per_image,
         separate_vision_encoder=separate_vision_encoder,
         sampler_config=sampler_config,
+        model_type_overridden=llm_model_type is not None,
     )
 
     with _cpu_default_device_scope():
@@ -1502,6 +1505,7 @@ def _build_llm_metadata(
     vision_cfg=None,
     audio_cfg=None,
     sampler_config=None,
+    model_type_overridden=False,
 ):
     """Serialize an ``LlmMetadata`` protobuf to *path*."""
     # The protobuf lives under an internal-looking subpackage of
@@ -1561,11 +1565,16 @@ def _build_llm_metadata(
         spec.populate_audio_metadata(meta, audio_cfg)
 
     # Populate function-calling fields for specs that declare them (currently
-    # only `FunctionGemmaSpec`, selected via the `llm_model_type` override).
+    # only `FunctionGemmaSpec`, reached via tokenizer auto-detection).
     # A base-class no-op for every other family, so this adds nothing to
     # existing model types -- the same convention as the vision/audio hooks
-    # above.
-    spec.populate_function_gemma_metadata(meta)
+    # above. Skipped when `llm_model_type` is an explicit caller override:
+    # litert-torch skips its model-specific metadata builder whenever
+    # `litert_lm_model_type_override` is set
+    # (`export_hf/core/litert_lm_builder.py` ~296), so an override yields the
+    # bare `llm_model_type` oneof with no model-specific fields.
+    if not model_type_overridden:
+        spec.populate_function_gemma_metadata(meta)
 
     # Sampler defaults are written only when the caller explicitly passes a
     # `sampler_config`; otherwise the field is omitted so the runtime picks

keras_hub/src/utils/litertlm/model_specs.py

@@ -558,10 +558,13 @@ def populate_function_gemma_metadata(self, meta):
         Default: no-op. Only ``FunctionGemmaSpec`` (the
         ``function_gemma_instruct_270m`` preset) overrides this to fill the
         ``FunctionGemma`` proto's function-calling fields; every other family
-        leaves the field block untouched. Called unconditionally by
-        ``_build_llm_metadata`` right after the ``llm_model_type`` oneof is
-        selected -- the same base-no-op convention ``populate_vision_metadata``
-        / ``populate_audio_metadata`` use.
+        leaves the field block untouched. Called by ``_build_llm_metadata``
+        right after the ``llm_model_type`` oneof is selected -- the same
+        base-no-op convention ``populate_vision_metadata`` /
+        ``populate_audio_metadata`` use -- except when ``llm_model_type`` was
+        an explicit caller override, mirroring litert-torch skipping its
+        model-specific metadata builder on override
+        (``export_hf/core/litert_lm_builder.py`` ~296).
         """
         del meta
 

keras_hub/src/utils/litertlm/traceable_ops.py

@@ -16,6 +16,7 @@
 import contextlib
 import unittest.mock
 
+import numpy as np
 import torch
 from keras.src import backend
 from keras.src.backend.common import dtypes as keras_dtypes
@@ -489,6 +490,10 @@ def _patched_amax(x, axis=None, keepdims=False):
     unchanged, so this is a no-op transform outside the triggering case.
     """
     x = torch_core.convert_to_tensor(x)
+    # NumPy integer axes (e.g. ``np.int64`` from shape arithmetic) are not
+    # Python ``int``; normalize so they take the same traceable path.
+    if isinstance(axis, np.integer):
+        axis = int(axis)
     if axis is not None and isinstance(axis, int) and x.ndim == 4:
         return torch.max(x, dim=axis, keepdim=keepdims).values
     return _ORIGINAL_AMAX(x, axis=axis, keepdims=keepdims)

keras_hub/src/utils/litertlm/traceable_ops_test.py

@@ -439,22 +439,35 @@ def test_amax_matches_original(self):
             (torch.randn(7), 0, True),  # 1-D (fallthrough)
             (torch.randn(2, 2, 3, 4, 5), -1, True),  # 5-D (fallthrough)
             (torch.randn(2, 3, 4, 5), (2, 3), True),  # 4-D tuple (fallthrough)
+            (torch.randn(2, 3, 4, 5), np.int64(-1), True),  # np.int64 axis
+            (torch.randn(2, 3, 4, 5), np.int32(1), True),  # np.int32 axis
         ]
         for x, axis, keepdims in cases:
             with self.subTest(
                 shape=tuple(x.shape), axis=axis, keepdims=keepdims
             ):
-                original = torch_backend_numpy.amax(
-                    x, axis=axis, keepdims=keepdims
-                )
                 patched = traceable_ops._patched_amax(
                     x, axis=axis, keepdims=keepdims
                 )
-                self.assertEqual(tuple(original.shape), tuple(patched.shape))
-                # The reformulation is the identical reduction; values must be
-                # bit-identical, not merely close (a silent numeric drift here
-                # would corrupt attention-sink softmax stabilization).
-                self.assertTrue(torch.equal(original, patched))
+                # The original Keras amax cannot handle NumPy integer axes
+                # (raises "truth value of an empty array is ambiguous"); for
+                # those cases verify only that the patched version works and
+                # matches torch.max directly.
+                if isinstance(axis, np.integer):
+                    expected = torch.max(
+                        x, dim=int(axis), keepdim=keepdims
+                    ).values
+                    self.assertTrue(torch.equal(expected, patched))
+                else:
+                    original = torch_backend_numpy.amax(
+                        x, axis=axis, keepdims=keepdims
+                    )
+                    self.assertEqual(tuple(original.shape), tuple(patched.shape))
+                    # The reformulation is the identical reduction; values must
+                    # be bit-identical, not merely close (a silent numeric
+                    # drift here would corrupt attention-sink softmax
+                    # stabilization).
+                    self.assertTrue(torch.equal(original, patched))
 
     def test_amax_public_ops_max_4d_matches(self):
         # The public keras op that GPT-OSS attention actually calls.