Fix TPU tests - for splash attention

keras/src/backend/jax/nn.py

@@ -1353,7 +1353,7 @@ def dot_product_attention(
             )
             # Transpose output back to Keras layout
             return jnp.transpose(output, axes=(0, 2, 1, 3))
-        except Exception:
+        except (jax.errors.ConcretizationTypeError, Exception):
             logging.exception(
                 "Failed to apply Splash kernel for flash attention. "
                 "Falling back to JAX native dot_product_attention."

keras/src/ops/nn_test.py

@@ -1324,6 +1324,37 @@ def test_polar(self):
 
 
 class NNOpsCorrectnessTest(testing.TestCase):
+    def test_dot_product_attention_inside_scan(self):
+        if backend.backend() != "jax":
+            self.skipTest("JAX-specific test")
+
+        import jax
+
+        try:
+            if jax.devices()[0].platform != "tpu":
+                self.skipTest("TPU-specific test")
+        except:
+            self.skipTest("TPU-specific test")
+
+        import jax.numpy as jnp
+
+        from keras.src.backend.jax import nn as jax_nn
+
+        def attention_scan_body(carry, x):
+            query, key, value = x
+            # Use a mask to trigger the issue
+            mask = jnp.ones((1, 4, 8), dtype="bool")
+            out = jax_nn.dot_product_attention(query, key, value, mask=mask)
+            return carry, out
+
+        query = jnp.ones((2, 1, 4, 8))
+        key = jnp.ones((2, 1, 4, 8))
+        value = jnp.ones((2, 1, 4, 8))
+
+        # Scan over the first dimension
+        _, out = jax.lax.scan(attention_scan_body, None, (query, key, value))
+        self.assertEqual(out.shape, (2, 1, 4, 8))
+
     def test_relu(self):
         x = np.array([-1, 0, 1, 2, 3], dtype=np.float32)
         self.assertAllClose(knn.relu(x), [0, 0, 1, 2, 3])