Fix TF int64 promotion issue. (#21679)

* Fix TF int64 promotion issue.

* Update.

* Fix skipif.

* Refine the comments.

* Update.

Author: james77777778

keras/src/backend/common/dtypes.py

@@ -232,16 +232,10 @@ def _resolve_weak_type(dtype, precision="32"):
         return f"float{precision}"
 
 
-BIT64_TO_BIT16_DTYPE = {
-    "int32": "int16",
-    "int64": "int16",
-    "uint32": "uint16",
-    "uint64": "uint16",
-    "float32": "float16",
-    "float64": "float16",
-}
 BIT64_TO_BIT32_DTYPE = {
-    "int64": "int32",
+    # Since TF variables require int64 to be placed on the GPU, we exclusively
+    # enable the int64 dtype for TF.
+    "int64": "int64" if config.backend() == "tensorflow" else "int32",
     "uint64": "uint32",
     "float64": "float32",
     "complex128": "complex64",
@@ -277,8 +271,8 @@ def _lattice_result_type(*args):
     if out_weak_type:
         out_dtype = _resolve_weak_type(out_dtype, precision=precision)
 
-    # Force to be 32-bit dtype when encountering 64-bit dtype.
-    # TODO(hongyu): Add a config to enable 64-bit dtypes.
+    # Force to be 32-bit dtype when encountering 64-bit dtype. This is to
+    # be aligned with JAX's default behavior.
     out_dtype = BIT64_TO_BIT32_DTYPE.get(out_dtype, out_dtype)
     return out_dtype
 

keras/src/backend/common/dtypes_test.py

@@ -1,5 +1,6 @@
 from unittest.mock import patch
 
+import pytest
 from absl.testing import parameterized
 
 from keras.src import backend
@@ -27,6 +28,13 @@ class DtypesTest(test_case.TestCase):
     ] + [None]
     if backend.backend() == "torch":
         ALL_DTYPES = [x for x in ALL_DTYPES if x not in ("uint16", "uint32")]
+    elif backend.backend() == "tensorflow":
+        # TODO(hongyu): Re-enable uint32 tests once we determine how to handle
+        # dtypes.result_type(uint32, int*) -> int64 promotion.
+        # Since TF variables require int64 to be placed on the GPU, we
+        # exclusively enable the int64 dtype for TF. However, JAX does not
+        # natively support int64, which prevents us from comparing the dtypes.
+        ALL_DTYPES = [x for x in ALL_DTYPES if x not in ("uint32",)]
     elif backend.backend() == "openvino":
         ALL_DTYPES = [x for x in ALL_DTYPES if x not in ("complex64",)]
 
@@ -55,6 +63,29 @@ def test_result_type_with_tensor(self, dtype1, dtype2):
         expected = jnp.result_type(x1_jax, x2_jax).name
         self.assertEqual(out, expected)
 
+    @parameterized.named_parameters(
+        named_product(
+            dtype=[
+                "int8",
+                "int16",
+                "int32",
+                "int64",
+                "uint8",
+                "uint16",
+                "uint32",
+            ]
+        )
+    )
+    @pytest.mark.skipif(
+        backend.backend() != "tensorflow", reason="TensorFlow only"
+    )
+    def test_result_type_with_int64(self, dtype):
+        # https://github.com/keras-team/keras/issues/21677
+        x1 = ops.ones((1,), dtype="int64")
+        x2 = ops.ones((1,), dtype=dtype)
+        out = backend.result_type(x1.dtype, x2.dtype)
+        self.assertEqual(out, "int64")
+
     def test_result_type_with_none(self):
         import jax.numpy as jnp
 

keras/src/backend/common/variables_test.py

@@ -811,6 +811,14 @@ class VariableOpsDTypeTest(test_case.TestCase):
             x for x in ALL_DTYPES if x not in ("uint16", "uint32", "complex64")
         ]
         INT_DTYPES = [x for x in INT_DTYPES if x not in ("uint16", "uint32")]
+    elif backend.backend() == "tensorflow":
+        # TODO(hongyu): Re-enable uint32 tests once we determine how to handle
+        # dtypes.result_type(uint32, int*) -> int64 promotion.
+        # Since TF variables require int64 to be placed on the GPU, we
+        # exclusively enable the int64 dtype for TF. However, JAX does not
+        # natively support int64, which prevents us from comparing the dtypes.
+        ALL_DTYPES = [x for x in ALL_DTYPES if x not in ("uint32",)]
+        INT_DTYPES = [x for x in INT_DTYPES if x not in ("uint32",)]
     elif backend.backend() == "openvino":
         ALL_DTYPES = [x for x in ALL_DTYPES if x not in ("complex64",)]
     NON_COMPLEX_DTYPES = [

keras/src/ops/numpy_test.py

@@ -5774,6 +5774,14 @@ class NumpyDtypeTest(testing.TestCase):
     if backend.backend() == "torch":
         ALL_DTYPES = [x for x in ALL_DTYPES if x not in ("uint16", "uint32")]
         INT_DTYPES = [x for x in INT_DTYPES if x not in ("uint16", "uint32")]
+    elif backend.backend() == "tensorflow":
+        # TODO(hongyu): Re-enable uint32 tests once we determine how to handle
+        # dtypes.result_type(uint32, int*) -> int64 promotion.
+        # Since TF variables require int64 to be placed on the GPU, we
+        # exclusively enable the int64 dtype for TF. However, JAX does not
+        # natively support int64, which prevents us from comparing the dtypes.
+        ALL_DTYPES = [x for x in ALL_DTYPES if x not in ("uint32",)]
+        INT_DTYPES = [x for x in INT_DTYPES if x not in ("uint32",)]
 
     @parameterized.named_parameters(
         named_product(dtypes=itertools.combinations(ALL_DTYPES, 2))