Fix for a change in JAX's promotion rules for floor and ceil. (ke…

…ras-team#19946)

JAX has an unreleased change jax-ml/jax#21441 that make it so that `int`s and `bool`s are no longer promoted to floats by `floor` and `ceil`.

Our current implementation follows the Numpy promotion rule, i.e. `int`s and `bool`s are promoted to floats. Therefore:
- updated unit tests, which use `jax.numpy` as the reference.
- updated our JAX implementation of `floor` and `ceil` (note that `ceil` would already cast).

keras/src/backend/jax/numpy.py

@@ -381,7 +381,8 @@ def ceil(x):
         dtype = config.floatx()
     else:
         dtype = dtypes.result_type(x.dtype, float)
-    return cast(jnp.ceil(x), dtype)
+    x = cast(x, dtype)
+    return jnp.ceil(x)
 
 
 def clip(x, x_min, x_max):
@@ -558,7 +559,10 @@ def flip(x, axis=None):
 def floor(x):
     x = convert_to_tensor(x)
     if standardize_dtype(x.dtype) == "int64":
-        x = cast(x, config.floatx())
+        dtype = config.floatx()
+    else:
+        dtype = dtypes.result_type(x.dtype, float)
+    x = cast(x, dtype)
     return jnp.floor(x)
 
 

keras/src/ops/numpy_test.py

@@ -12,6 +12,7 @@
 from keras.src import backend
 from keras.src import testing
 from keras.src.backend.common import dtypes
+from keras.src.backend.common import is_int_dtype
 from keras.src.backend.common import standardize_dtype
 from keras.src.backend.common.keras_tensor import KerasTensor
 from keras.src.ops import numpy as knp
@@ -5794,7 +5795,8 @@ def test_ceil(self, dtype):
         x = knp.array(value, dtype=dtype)
         x_jax = jnp.array(value, dtype=dtype)
         expected_dtype = standardize_dtype(jnp.ceil(x_jax).dtype)
-        if dtype == "int64":
+        # Here, we follow Numpy's rule, not JAX's; ints are promoted to floats.
+        if dtype == "bool" or is_int_dtype(dtype):
             expected_dtype = backend.floatx()
 
         self.assertEqual(standardize_dtype(knp.ceil(x).dtype), expected_dtype)
@@ -6377,7 +6379,8 @@ def test_floor(self, dtype):
         x = knp.ones((1,), dtype=dtype)
         x_jax = jnp.ones((1,), dtype=dtype)
         expected_dtype = standardize_dtype(jnp.floor(x_jax).dtype)
-        if dtype == "int64":
+        # Here, we follow Numpy's rule, not JAX's; ints are promoted to floats.
+        if dtype == "bool" or is_int_dtype(dtype):
             expected_dtype = backend.floatx()
 
         self.assertEqual(standardize_dtype(knp.floor(x).dtype), expected_dtype)