Update nn.py to use jnp.add rather than + operand. (#19323)

jenriver · web-flow · commit 3b92b678f7e1 · 2024-04-14T17:28:16.000-07:00
diff --git a/keras/backend/jax/nn.py b/keras/backend/jax/nn.py
@@ -583,7 +583,7 @@ def batch_normalization(
         offset = jnp.reshape(offset, shape)
         res = res + offset
 
-    return x * inv + res
+    return jnp.add(x * inv, res)
 
 
 def ctc_loss(
