Fix numpy tests for window ops, angle, identity, eye. (#21596)

JAX is introducing a change whereby some ops, in particular window ops, return a result of the default dtype for the current configuration (in particular based on the `enable_x64` flag).

For cross-backend consistency, we want the window ops to return `floatx` regardless of how JAX is configured.

Additionally, the unit tests need to be changed to not use the JAX returned dtype as the expected dtype.

Author: hertschuh

keras/src/backend/jax/numpy.py

@@ -39,17 +39,17 @@ def add(x1, x2):
 
 def bartlett(x):
     x = convert_to_tensor(x)
-    return jnp.bartlett(x)
+    return cast(jnp.bartlett(x), config.floatx())
 
 
 def hamming(x):
     x = convert_to_tensor(x)
-    return jnp.hamming(x)
+    return cast(jnp.hamming(x), config.floatx())
 
 
 def hanning(x):
     x = convert_to_tensor(x)
-    return jnp.hanning(x)
+    return cast(jnp.hanning(x), config.floatx())
 
 
 def heaviside(x1, x2):
@@ -60,7 +60,7 @@ def heaviside(x1, x2):
 
 def kaiser(x, beta):
     x = convert_to_tensor(x)
-    return jnp.kaiser(x, beta)
+    return cast(jnp.kaiser(x, beta), config.floatx())
 
 
 def bincount(x, weights=None, minlength=0, sparse=False):
@@ -497,7 +497,7 @@ def right_shift(x, y):
 
 def blackman(x):
     x = convert_to_tensor(x)
-    return jnp.blackman(x)
+    return cast(jnp.blackman(x), config.floatx())
 
 
 def broadcast_to(x, shape):

keras/src/ops/numpy_test.py

@@ -5765,11 +5765,8 @@ def test_add_python_types(self, dtype):
 
     @parameterized.named_parameters(named_product(dtype=ALL_DTYPES))
     def test_bartlett(self, dtype):
-        import jax.numpy as jnp
-
         x = knp.ones((), dtype=dtype)
-        x_jax = jnp.ones((), dtype=dtype)
-        expected_dtype = standardize_dtype(jnp.bartlett(x_jax).dtype)
+        expected_dtype = backend.floatx()
 
         self.assertEqual(
             standardize_dtype(knp.bartlett(x).dtype), expected_dtype
@@ -5781,11 +5778,8 @@ def test_bartlett(self, dtype):
 
     @parameterized.named_parameters(named_product(dtype=ALL_DTYPES))
     def test_blackman(self, dtype):
-        import jax.numpy as jnp
-
         x = knp.ones((), dtype=dtype)
-        x_jax = jnp.ones((), dtype=dtype)
-        expected_dtype = standardize_dtype(jnp.blackman(x_jax).dtype)
+        expected_dtype = backend.floatx()
 
         self.assertEqual(
             standardize_dtype(knp.blackman(x).dtype), expected_dtype
@@ -5797,11 +5791,8 @@ def test_blackman(self, dtype):
 
     @parameterized.named_parameters(named_product(dtype=ALL_DTYPES))
     def test_hamming(self, dtype):
-        import jax.numpy as jnp
-
         x = knp.ones((), dtype=dtype)
-        x_jax = jnp.ones((), dtype=dtype)
-        expected_dtype = standardize_dtype(jnp.hamming(x_jax).dtype)
+        expected_dtype = backend.floatx()
 
         self.assertEqual(
             standardize_dtype(knp.hamming(x).dtype), expected_dtype
@@ -5813,11 +5804,8 @@ def test_hamming(self, dtype):
 
     @parameterized.named_parameters(named_product(dtype=ALL_DTYPES))
     def test_hanning(self, dtype):
-        import jax.numpy as jnp
-
         x = knp.ones((), dtype=dtype)
-        x_jax = jnp.ones((), dtype=dtype)
-        expected_dtype = standardize_dtype(jnp.hanning(x_jax).dtype)
+        expected_dtype = backend.floatx()
 
         self.assertEqual(
             standardize_dtype(knp.hanning(x).dtype), expected_dtype
@@ -5829,13 +5817,9 @@ def test_hanning(self, dtype):
 
     @parameterized.named_parameters(named_product(dtype=ALL_DTYPES))
     def test_kaiser(self, dtype):
-        import jax.numpy as jnp
-
         x = knp.ones((), dtype=dtype)
         beta = knp.ones((), dtype=dtype)
-        x_jax = jnp.ones((), dtype=dtype)
-        beta_jax = jnp.ones((), dtype=dtype)
-        expected_dtype = standardize_dtype(jnp.kaiser(x_jax, beta_jax).dtype)
+        expected_dtype = backend.floatx()
 
         self.assertEqual(
             standardize_dtype(knp.kaiser(x, beta).dtype), expected_dtype
@@ -7268,13 +7252,17 @@ def test_eye(self, dtype):
         import jax.numpy as jnp
 
         expected_dtype = standardize_dtype(jnp.eye(3, dtype=dtype).dtype)
+        if dtype is None:
+            expected_dtype = backend.floatx()
 
         self.assertEqual(
             standardize_dtype(knp.eye(3, dtype=dtype).dtype),
             expected_dtype,
         )
 
         expected_dtype = standardize_dtype(jnp.eye(3, 4, 1, dtype=dtype).dtype)
+        if dtype is None:
+            expected_dtype = backend.floatx()
 
         self.assertEqual(
             standardize_dtype(knp.eye(3, 4, k=1, dtype=dtype).dtype),
@@ -7506,6 +7494,8 @@ def test_identity(self, dtype):
         import jax.numpy as jnp
 
         expected_dtype = standardize_dtype(jnp.identity(3, dtype=dtype).dtype)
+        if dtype is None:
+            expected_dtype = backend.floatx()
 
         self.assertEqual(
             standardize_dtype(knp.identity(3, dtype=dtype).dtype),
@@ -9141,7 +9131,7 @@ def test_angle(self, dtype):
         x = knp.ones((1,), dtype=dtype)
         x_jax = jnp.ones((1,), dtype=dtype)
         expected_dtype = standardize_dtype(jnp.angle(x_jax).dtype)
-        if dtype == "int64":
+        if dtype == "bool" or is_int_dtype(dtype):
             expected_dtype = backend.floatx()
 
         self.assertEqual(standardize_dtype(knp.angle(x).dtype), expected_dtype)