Codestin Search App

Deeven-Seru · 2026-09-03T04:36:21Z

This PR fixes the zero-day division by zero NaN vulnerability described in #23551.

By adding backend.epsilon() to the denominator when normalizing predictions across all backends (Legacy, TF, Torch, JAX, NumPy, OpenVINO), we prevent catastrophic gradient poisoning when predictions sum to exactly zero.

This issue was originally identified and patched in TensorFlow (tensorflow/tensorflow#123494) and is being ported upstream here per the maintainers' request.

Fixes #23551

Contributor Agreement

I am a human, and not a bot.
I will be responsible for responding to review comments in a timely manner.
I will work with the maintainers to push this PR forward until submission.

gemini-code-assist

Code Review

This pull request updates the categorical cross-entropy and sparse categorical cross-entropy implementations across multiple backends (JAX, NumPy, OpenVINO, TensorFlow, Torch, and legacy) to add backend epsilon to the denominator during normalization, preventing potential division-by-zero errors. The review feedback suggests improving consistency and robustness by storing the epsilon value in a variable cast to the output's data type (using convert_to_tensor) instead of calling backend.epsilon() multiple times.

codecov-commenter · 2026-09-03T04:47:39Z

Codecov Report

❌ Patch coverage is 44.57831% with 46 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.52%. Comparing base (e2e6815) to head (bc939c8).
⚠️ Report is 49 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/nn.py	23.91%	35 Missing ⚠️
keras/src/backend/tensorflow/nn.py	0.00%	5 Missing ⚠️
keras/src/backend/torch/nn.py	44.44%	4 Missing and 1 partial ⚠️
keras/src/backend/numpy/nn.py	66.66%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (e2e6815) and HEAD (bc939c8). Click for more details.

HEAD has 18 uploads less than BASE

Flag BASE (e2e6815) HEAD (bc939c8)

keras 9 3

keras-tpu 1 0

keras-jax 3 0

keras-cpu 5 3

keras-tensorflow 2 0

keras-torch 2 1

keras-gpu 3 0

Additional details and impacted files

@@             Coverage Diff             @@
##           master   #23552       +/-   ##
===========================================
- Coverage   84.93%   73.52%   -11.41%     
===========================================
  Files         468      478       +10     
  Lines       71147    71678      +531     
  Branches    11788    11838       +50     
===========================================
- Hits        60426    52702     -7724     
- Misses       7717    16050     +8333     
+ Partials     3004     2926       -78

Flag	Coverage Δ
keras	`73.42% <44.57%> (-11.32%)`	⬇️
keras-cpu	`73.42% <44.57%> (-10.70%)`	⬇️
keras-gpu	`?`
keras-jax	`?`
keras-numpy	`54.06% <16.86%> (+0.06%)`	⬆️
keras-openvino	`59.73% <37.34%> (+0.03%)`	⬆️
keras-tensorflow	`?`
keras-torch	`59.50% <19.27%> (-0.27%)`	⬇️
keras-tpu	`?`

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

hertschuh

Thanks for looking into this!

Deeven-Seru · 2026-09-04T04:39:12Z

done with requested modification 👍

hertschuh

The JAX tests are failing, apparently, there are still NaNs in the result. Can you look into this?

Deeven-Seru · 2026-09-04T18:52:06Z

Hi @hertschuh , trying to figure out a fix for this, any suggestions?

hertschuh · 2026-09-04T23:06:13Z

Hi @hertschuh , trying to figure out a fix for this, any suggestions?

Can you try if divide_no_nan fixes it with JAX?

Deeven-Seru · 2026-09-05T06:34:22Z

@hertschuh PTAL :)

This prevents catastrophic gradient poisoning when predictions sum to exactly zero by adding `backend.epsilon()` to the denominator across legacy, PyTorch, TF, OpenVINO and NumPy backends using `convert_to_tensor(backend.epsilon(), dtype=output.dtype)`. For JAX backend, we utilize `divide_no_nan` to explicitly protect against subnormal float16 hardware FTZ behaviors that cause 0.0/0.0 NaNs. Ported from tensorflow/tensorflow#123494 Fixes keras-team#23551

hertschuh · 2026-09-08T18:10:35Z

@Deeven-Seru

@hertschuh PTAL :)

Oh, I was thinking you would do divide_no_nan everywhere for all backends for consistency. Does that not work?

Deeven-Seru · 2026-09-09T03:10:20Z

@Deeven-Seru

@hertschuh PTAL :)

Oh, I was thinking you would do divide_no_nan everywhere for all backends for consistency. Does that not work?

I tried Using divide_no_nan but in NumPy eagerly evaluates the division before masking, which still throws a RuntimeWarning: divide by zero on execution. Furthermore, PyTorch and TF divide_no_nan implementations allocate multiple intermediate mask tensors, making them slower than a single highly-optimized maximum kernel. So maximum preserves faster continuous gradients for those backends, making divide_no_nan strictly a JAX-specific compiler workaround.

hertschuh · 2026-09-10T01:03:04Z

@Deeven-Seru

@hertschuh PTAL :)

Oh, I was thinking you would do divide_no_nan everywhere for all backends for consistency. Does that not work?

I tried Using divide_no_nan but in NumPy eagerly evaluates the division before masking, which still throws a RuntimeWarning: divide by zero on execution. Furthermore, PyTorch and TF divide_no_nan implementations allocate multiple intermediate mask tensors, making them slower than a single highly-optimized maximum kernel. So maximum preserves faster continuous gradients for those backends, making divide_no_nan strictly a JAX-specific compiler workaround.

Hmm... Ok, looking at the math, I like the epsilon approach better actually. But now I don't like that JAX is inconsistent in the approach.

Sorry for the back and forth, but can you revert back the JAX implementation to the epsilon approach? Let's debug why it's not passing the tests.

…ossentropy via custom_vjp

Deeven-Seru · 2026-09-10T08:29:38Z

the tests are failing due to float16 underflow I guess , I tried to resolve it by adding custom_vjp to compute in float32

mathewOracle · 2026-09-10T10:27:42Z

 )
 from keras.src.backend.jax.core import cast
 from keras.src.backend.jax.core import convert_to_tensor
+from keras.src.backend.jax.numpy import divide_no_nan


unused import — not referenced anywhere in this diff, grep confirms it

Iteratively working on thr pr

mathewOracle · 2026-09-10T10:28:33Z

+    return -jnp.sum(target * log_prob, axis=axis)
+
+
+categorical_crossentropy_vjp = jax.custom_vjp(


this isn't mentioned in the PR description

I will change the description once the fix is confirmed

this isn't mentioned in the PR description

google-ml-butler Bot added the size:S label Sep 3, 2026

google-ml-butler Bot assigned gbaned Sep 3, 2026

gemini-code-assist Bot reviewed Sep 3, 2026

View reviewed changes

Deeven-Seru force-pushed the fix-categorical-crossentropy-nan branch 2 times, most recently from e66e3ae to 95f6d7e Compare September 3, 2026 06:44

hertschuh requested changes Sep 4, 2026

View reviewed changes

Comment thread test_nan.py Outdated

hertschuh added the stat:awaiting response from contributor label Sep 4, 2026

Deeven-Seru force-pushed the fix-categorical-crossentropy-nan branch from 95f6d7e to 76af0f0 Compare September 4, 2026 04:37

google-ml-butler Bot removed the stat:awaiting response from contributor label Sep 4, 2026

Deeven-Seru force-pushed the fix-categorical-crossentropy-nan branch 2 times, most recently from d54659c to a0604f2 Compare September 4, 2026 06:22

hertschuh reviewed Sep 4, 2026

View reviewed changes

Comment thread test_nan_jax.py Outdated

hertschuh added the stat:awaiting response from contributor label Sep 4, 2026

google-ml-butler Bot removed the stat:awaiting response from contributor label Sep 4, 2026

hertschuh added the stat:awaiting response from contributor label Sep 4, 2026

google-ml-butler Bot removed the stat:awaiting response from contributor label Sep 5, 2026

Deeven-Seru force-pushed the fix-categorical-crossentropy-nan branch from 91c4074 to 9a560be Compare September 5, 2026 05:29

Deeven-Seru force-pushed the fix-categorical-crossentropy-nan branch from 9a560be to 9825ea7 Compare September 5, 2026 06:39

keerthanakadiri added the stat:awaiting keras-eng Awaiting response from Keras engineer label Sep 7, 2026

hertschuh added stat:awaiting response from contributor and removed stat:awaiting keras-eng Awaiting response from Keras engineer labels Sep 8, 2026

google-ml-butler Bot removed the stat:awaiting response from contributor label Sep 9, 2026

hertschuh added the stat:awaiting response from contributor label Sep 10, 2026

Fix float16 NaN in categorical_crossentropy and sparse_categorical_cr…

bc939c8

…ossentropy via custom_vjp

google-ml-butler Bot removed the stat:awaiting response from contributor label Sep 10, 2026

mathewOracle suggested changes Sep 10, 2026

View reviewed changes

keerthanakadiri added stat:awaiting keras-eng Awaiting response from Keras engineer stat:awaiting response from contributor and removed stat:awaiting keras-eng Awaiting response from Keras engineer labels Sep 11, 2026

		return -jnp.sum(target * log_prob, axis=axis)


		categorical_crossentropy_vjp = jax.custom_vjp(

Uh oh!

Conversation

Deeven-Seru commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Contributor Agreement

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov-commenter commented Sep 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Deeven-Seru commented Sep 4, 2026

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Deeven-Seru commented Sep 4, 2026

Uh oh!

hertschuh commented Sep 4, 2026

Uh oh!

Deeven-Seru commented Sep 5, 2026

Uh oh!

hertschuh commented Sep 8, 2026

Uh oh!

Deeven-Seru commented Sep 9, 2026

Uh oh!

hertschuh commented Sep 10, 2026

Uh oh!

Deeven-Seru commented Sep 10, 2026

Uh oh!

mathewOracle Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

Deeven-Seru Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

mathewOracle Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

Deeven-Seru Sep 10, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

Deeven-Seru commented Sep 3, 2026 •

edited

Loading

codecov-commenter commented Sep 3, 2026 •

edited

Loading