Skip to content

Fix handling of SaS metrics for Jax #21765

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 24, 2025
+20 −0
Merged

Fix handling of SaS metrics for Jax #21765

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions keras/src/backend/jax/numpy.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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]
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This fix causes nonzero to return an array of the same size as the input (it pads the output) when jitted. So this will behave completely different when jitted and not jitted. We cannot do this.

In the context of SaS metrics for instance.

keras/keras/src/metrics/confusion_metrics.py

Lines 667 to 671 in 18f79d6

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)

ops.size will always be the same as the size as self.value therefore the condition feasible_exists will always be True.

Instead, we should apply the predicate and do a reduction.

feasible = predicate(constrained, self.value)
feasible_exists = keras.ops.any(feasible)

Then, we need to find the max too.

Either we expose the size option (and implement it for all backends) and use it properly in confusion_metrics.py, or we reimplement the max without nonzero and with ops that are compilable (fixed size intermediate values).


return jnp.nonzero(x)


Expand Down
14 changes: 14 additions & 0 deletions keras/src/metrics/confusion_metrics_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down