Make confusion metrics compilable. (#21775)

By removing the use of `ops.nonzero` which returns an array of non-predetermined size.

Follow-up to https://github.com/keras-team/keras/pull/21765/files#r2462099871

Fixes #19376

Author: hertschuh

keras/src/backend/jax/numpy.py

@@ -3,7 +3,6 @@
 
 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
@@ -1002,11 +1001,6 @@ 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.py

@@ -654,7 +654,7 @@ def _find_max_under_constraint(self, constrained, dependent, predicate):
         Args:
             constrained: Over these values the constraint is specified. A rank-1
                 tensor.
-            dependent: From these values the maximum that satiesfies the
+            dependent: From these values the maximum that satisfies the
                 constraint is selected. Values in this tensor and in
                 `constrained` are linked by having the same threshold at each
                 position, hence this tensor must have the same shape.
@@ -664,11 +664,12 @@ def _find_max_under_constraint(self, constrained, dependent, predicate):
         Returns:
             maximal dependent value, if no value satisfies the constraint 0.0.
         """
-        feasible = ops.nonzero(predicate(constrained, self.value))
-        feasible_exists = ops.greater(ops.size(feasible), 0)
-        max_dependent = ops.max(ops.take(dependent, feasible), initial=0)
-
-        return ops.where(feasible_exists, max_dependent, 0.0)
+        feasible = predicate(constrained, self.value)
+        # Mask values based on whether they satisfy the constraint and take max.
+        return ops.max(
+            ops.multiply(dependent, ops.cast(feasible, dependent.dtype)),
+            initial=0,
+        )
 
 
 @keras_export("keras.metrics.SensitivityAtSpecificity")