Work around Keras 3.14.1 JAX backend dtype mismatch in slice/slice_update

Author: pctablet505

conftest.py

@@ -48,6 +48,42 @@ def pytest_addoption(parser):
 
 
 def pytest_configure(config):
+    # Monkey-patch JAX slice/slice_update to normalize mixed-dtype
+    # start_indices. Keras >= 3.14.1 can pass Python ints (int64) alongside
+    # JAX int32 arrays to jax.lax.dynamic_slice, which rejects them.
+    if keras.config.backend() == "jax":
+        import jax
+        import jax.numpy as jnp
+        from keras.src.backend.jax import core as jax_core
+
+        _original_jax_slice = jax_core.slice
+        _original_jax_slice_update = jax_core.slice_update
+
+        def _normalize_start_indices(start_indices):
+            arrays = [
+                jnp.asarray(idx)
+                if not isinstance(idx, jax.Array)
+                else idx
+                for idx in start_indices
+            ]
+            dtypes = {a.dtype for a in arrays}
+            if len(dtypes) > 1:
+                arrays = [jnp.astype(a, jnp.int32) for a in arrays]
+            return arrays
+
+        def _patched_slice(inputs, start_indices, shape):
+            start_indices = _normalize_start_indices(start_indices)
+            return _original_jax_slice(inputs, start_indices, shape)
+
+        def _patched_slice_update(inputs, start_indices, updates):
+            start_indices = _normalize_start_indices(start_indices)
+            return _original_jax_slice_update(
+                inputs, start_indices, updates
+            )
+
+        jax_core.slice = _patched_slice
+        jax_core.slice_update = _patched_slice_update
+
     # Monkey-patch training methods for OpenVINO backend
     if keras.config.backend() == "openvino":
         keras.Model.fit = lambda *args, **kwargs: pytest.skip(