Fix bitwise left_shift and right_shift result dtype... (#21034)

when second argument is a constant int.

Previously, a `convert_to_tensor` was applied to the second argument, making it an `int32` or `int64`. The result dtype would take into account this dtype, which could upgrade the dtype of the result.

The expectation is that if the second argument is a constant, the result dtype is the same as the first argument. This is already supported correctly by all underlying backend implementations.

Author: hertschuh

keras/src/backend/jax/numpy.py

@@ -439,7 +439,8 @@ def bitwise_xor(x, y):
 
 def bitwise_left_shift(x, y):
     x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
+    if not isinstance(y, int):
+        y = convert_to_tensor(y)
     return jnp.left_shift(x, y)
 
 
@@ -449,7 +450,8 @@ def left_shift(x, y):
 
 def bitwise_right_shift(x, y):
     x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
+    if not isinstance(y, int):
+        y = convert_to_tensor(y)
     return jnp.right_shift(x, y)
 
 

keras/src/backend/numpy/numpy.py

@@ -355,7 +355,8 @@ def bitwise_xor(x, y):
 
 def bitwise_left_shift(x, y):
     x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
+    if not isinstance(y, int):
+        y = convert_to_tensor(y)
     return np.left_shift(x, y)
 
 
@@ -365,7 +366,8 @@ def left_shift(x, y):
 
 def bitwise_right_shift(x, y):
     x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
+    if not isinstance(y, int):
+        y = convert_to_tensor(y)
     return np.right_shift(x, y)
 
 

keras/src/backend/tensorflow/numpy.py

@@ -1007,10 +1007,11 @@ def bitwise_xor(x, y):
 
 def bitwise_left_shift(x, y):
     x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
-    dtype = dtypes.result_type(x.dtype, y.dtype)
-    x = tf.cast(x, dtype)
-    y = tf.cast(y, dtype)
+    if not isinstance(y, int):
+        y = convert_to_tensor(y)
+        dtype = dtypes.result_type(x.dtype, y.dtype)
+        x = tf.cast(x, dtype)
+        y = tf.cast(y, dtype)
     return tf.bitwise.left_shift(x, y)
 
 
@@ -1020,10 +1021,11 @@ def left_shift(x, y):
 
 def bitwise_right_shift(x, y):
     x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
-    dtype = dtypes.result_type(x.dtype, y.dtype)
-    x = tf.cast(x, dtype)
-    y = tf.cast(y, dtype)
+    if not isinstance(y, int):
+        y = convert_to_tensor(y)
+        dtype = dtypes.result_type(x.dtype, y.dtype)
+        x = tf.cast(x, dtype)
+        y = tf.cast(y, dtype)
     return tf.bitwise.right_shift(x, y)
 
 

keras/src/backend/torch/numpy.py

@@ -479,7 +479,8 @@ def bitwise_xor(x, y):
 
 def bitwise_left_shift(x, y):
     x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
+    if not isinstance(y, int):
+        y = convert_to_tensor(y)
     return torch.bitwise_left_shift(x, y)
 
 
@@ -489,7 +490,8 @@ def left_shift(x, y):
 
 def bitwise_right_shift(x, y):
     x = convert_to_tensor(x)
-    y = convert_to_tensor(y)
+    if not isinstance(y, int):
+        y = convert_to_tensor(y)
     return torch.bitwise_right_shift(x, y)
 
 

keras/src/ops/numpy.py

@@ -1417,7 +1417,10 @@ def call(self, x, y):
         return backend.numpy.bitwise_left_shift(x, y)
 
     def compute_output_spec(self, x, y):
-        dtype = dtypes.result_type(x.dtype, y.dtype)
+        if isinstance(y, int):
+            dtype = x.dtype
+        else:
+            dtype = dtypes.result_type(x.dtype, y.dtype)
         return KerasTensor(x.shape, dtype=dtype)
 
 
@@ -1451,7 +1454,10 @@ def call(self, x, y):
         return backend.numpy.left_shift(x, y)
 
     def compute_output_spec(self, x, y):
-        dtype = dtypes.result_type(x.dtype, y.dtype)
+        if isinstance(y, int):
+            dtype = x.dtype
+        else:
+            dtype = dtypes.result_type(x.dtype, y.dtype)
         return KerasTensor(x.shape, dtype=dtype)
 
 
@@ -1483,7 +1489,10 @@ def call(self, x, y):
         return backend.numpy.bitwise_right_shift(x, y)
 
     def compute_output_spec(self, x, y):
-        dtype = dtypes.result_type(x.dtype, y.dtype)
+        if isinstance(y, int):
+            dtype = x.dtype
+        else:
+            dtype = dtypes.result_type(x.dtype, y.dtype)
         return KerasTensor(x.shape, dtype=dtype)
 
 
@@ -1517,7 +1526,10 @@ def call(self, x, y):
         return backend.numpy.right_shift(x, y)
 
     def compute_output_spec(self, x, y):
-        dtype = dtypes.result_type(x.dtype, y.dtype)
+        if isinstance(y, int):
+            dtype = x.dtype
+        else:
+            dtype = dtypes.result_type(x.dtype, y.dtype)
         return KerasTensor(x.shape, dtype=dtype)
 
 

keras/src/ops/numpy_test.py

@@ -6229,16 +6229,16 @@ def test_bitwise_xor(self, dtypes):
         self.assertDType(knp.BitwiseXor().symbolic_call(x1, x2), expected_dtype)
 
     @parameterized.named_parameters(
-        named_product(dtypes=itertools.combinations(INT_DTYPES, 2))
+        named_product(dtypes=itertools.product(INT_DTYPES, INT_DTYPES + [None]))
     )
     def test_bitwise_left_shift(self, dtypes):
         import jax.numpy as jnp
 
         dtype1, dtype2 = dtypes
         x1 = knp.ones((1,), dtype=dtype1)
-        x2 = knp.ones((1,), dtype=dtype2)
+        x2 = knp.ones((1,), dtype=dtype2) if dtype2 else 1
         x1_jax = jnp.ones((1,), dtype=dtype1)
-        x2_jax = jnp.ones((1,), dtype=dtype2)
+        x2_jax = jnp.ones((1,), dtype=dtype2) if dtype2 else 1
         expected_dtype = standardize_dtype(jnp.left_shift(x1_jax, x2_jax).dtype)
 
         self.assertDType(knp.bitwise_left_shift(x1, x2), expected_dtype)
@@ -6249,16 +6249,16 @@ def test_bitwise_left_shift(self, dtypes):
     # left_shift is same as bitwise_left_shift
 
     @parameterized.named_parameters(
-        named_product(dtypes=itertools.combinations(INT_DTYPES, 2))
+        named_product(dtypes=itertools.product(INT_DTYPES, INT_DTYPES + [None]))
     )
     def test_bitwise_right_shift(self, dtypes):
         import jax.numpy as jnp
 
         dtype1, dtype2 = dtypes
         x1 = knp.ones((1,), dtype=dtype1)
-        x2 = knp.ones((1,), dtype=dtype2)
+        x2 = knp.ones((1,), dtype=dtype2) if dtype2 else 1
         x1_jax = jnp.ones((1,), dtype=dtype1)
-        x2_jax = jnp.ones((1,), dtype=dtype2)
+        x2_jax = jnp.ones((1,), dtype=dtype2) if dtype2 else 1
         expected_dtype = standardize_dtype(
             jnp.right_shift(x1_jax, x2_jax).dtype
         )