fix sas metrics in jax fit (keras-team#21765)

Author: danielenricocahall

keras/src/backend/jax/numpy.py

@@ -3,6 +3,7 @@
 
 import jax.experimental.sparse as jax_sparse
 import jax.numpy as jnp
+from jax import core
 from jax import export as jax_export
 
 from keras.src.backend import config
@@ -1001,6 +1002,11 @@ def ndim(x):
 
 
 def nonzero(x):
+    if isinstance(x, core.Tracer):
+        # needed because this is called for several metric calculations,
+        # which will supply tracer values during `fit` execution
+        return jnp.nonzero(x, size=core.get_aval(x).size)[0]
+
     return jnp.nonzero(x)
 
 

keras/src/metrics/confusion_metrics_test.py

@@ -787,6 +787,20 @@ def test_invalid_num_thresholds(self):
         ):
             metrics.SensitivityAtSpecificity(0.4, num_thresholds=-1)
 
+    @pytest.mark.requires_trainable_backend
+    def test_handles_sas_metrics(self):
+        # Test for https://github.com/keras-team/keras/issues/19376
+        model = models.Sequential(
+            [
+                layers.Input((1,)),
+                layers.Dense(1),
+            ]
+        )
+        sas = metrics.SpecificityAtSensitivity(0.5, name="sas")
+
+        model.compile(optimizer="adam", loss="crossentropy", metrics=[sas])
+        model.fit(np.ones((5, 1)), np.ones((5, 1)))
+
 
 class SpecificityAtSensitivityTest(testing.TestCase):
     def test_config(self):