Fix mixed dtypes issue in nn.dot_product_attention. (#21558)

Author: james77777778

keras/src/backend/jax/nn.py

@@ -1241,6 +1241,12 @@ def dot_product_attention(
             f"Received: query.shape={query.shape}, key.shape={key.shape}, "
             f"value.shape={value.shape}."
         )
+    compute_dtype = backend.result_type(query.dtype, key.dtype, value.dtype)
+    query = cast(query, compute_dtype)
+    key = cast(key, compute_dtype)
+    value = cast(value, compute_dtype)
+    if bias is not None:
+        bias = convert_to_tensor(bias, dtype=compute_dtype)
 
     # Check platform
     platform = jax.devices()[0].platform

keras/src/backend/numpy/nn.py

@@ -1163,6 +1163,13 @@ def dot_product_attention(
             f"Received: query.shape={query.shape}, key.shape={key.shape}, "
             f"value.shape={value.shape}."
         )
+    compute_dtype = backend.result_type(query.dtype, key.dtype, value.dtype)
+    query = cast(query, compute_dtype)
+    key = cast(key, compute_dtype)
+    value = cast(value, compute_dtype)
+    if bias is not None:
+        bias = convert_to_tensor(bias, dtype=compute_dtype)
+
     _, _, _, H = key.shape
     scale = (1.0 / np.sqrt(H)) if scale is None else scale
     return _dot_product_attention_xla(

keras/src/backend/openvino/nn.py

@@ -486,3 +486,19 @@ def ctc_decode(
 
 def psnr(x1, x2, max_val):
     raise NotImplementedError("`psnr` is not supported with openvino backend")
+
+
+def dot_product_attention(
+    query,
+    key,
+    value,
+    bias=None,
+    mask=None,
+    scale=None,
+    is_causal=False,
+    flash_attention=None,
+    attn_logits_soft_cap=None,
+):
+    raise NotImplementedError(
+        "`dot_product_attention` is not supported with openvino backend"
+    )

keras/src/backend/tensorflow/nn.py

@@ -1065,6 +1065,13 @@ def dot_product_attention(
             f"Received: query.shape={query.shape}, key.shape={key.shape}, "
             f"value.shape={value.shape}."
         )
+    compute_dtype = backend.result_type(query.dtype, key.dtype, value.dtype)
+    query = cast(query, compute_dtype)
+    key = cast(key, compute_dtype)
+    value = cast(value, compute_dtype)
+    if bias is not None:
+        bias = convert_to_tensor(bias, dtype=compute_dtype)
+
     H = tf.shape(key)[-1]
     scale = (1.0 / tf.sqrt(tf.cast(H, "float32"))) if scale is None else scale
     return _dot_product_attention_xla(

keras/src/backend/torch/nn.py

@@ -1043,6 +1043,11 @@ def dot_product_attention(
             f"Received: query.shape={query.shape}, key.shape={key.shape}, "
             f"value.shape={value.shape}."
         )
+    compute_dtype = backend.result_type(query.dtype, key.dtype, value.dtype)
+    query = cast(query, compute_dtype)
+    key = cast(key, compute_dtype)
+    value = cast(value, compute_dtype)
+
     mask = mask if mask is None else convert_to_tensor(mask, dtype="bool")
     if mask is not None:
         # Explicit set `is_causal` to `False` when `mask` is not `None`.

keras/src/ops/nn.py

@@ -2641,7 +2641,8 @@ def compute_output_spec(
         mask=None,
         scale=None,
     ):
-        return KerasTensor(query.shape, dtype=query.dtype)
+        dtype = backend.result_type(query.dtype, key.dtype, value.dtype)
+        return KerasTensor(query.shape, dtype=dtype)
 
 
 @keras_export(

keras/src/ops/nn_test.py

@@ -3114,14 +3114,20 @@ def test_ctc_decode(self, dtype):
         self.assertEqual(standardize_dtype(decoded.dtype), "int32")
         self.assertEqual(standardize_dtype(scores.dtype), expected_dtype)
 
-    @parameterized.named_parameters(named_product(dtype=FLOAT_DTYPES))
-    def test_dot_product_attention(self, dtype):
+    @parameterized.named_parameters(
+        named_product(
+            dtypes=list(combinations(FLOAT_DTYPES, 2))
+            + [(dtype, dtype) for dtype in FLOAT_DTYPES]
+        )
+    )
+    def test_dot_product_attention(self, dtypes):
         # TODO: Get expected output from jax if `jax.nn.dot_product_attention`
         # is available.
-        query = knp.ones((2, 3, 3, 8), dtype=dtype)
-        key = knp.ones((2, 3, 3, 8), dtype=dtype)
-        value = knp.ones((2, 3, 3, 8), dtype=dtype)
-        expected_dtype = dtype
+        query_dtype, key_value_dtype = dtypes
+        query = knp.ones((2, 3, 3, 8), dtype=query_dtype)
+        key = knp.ones((2, 3, 3, 8), dtype=key_value_dtype)
+        value = knp.ones((2, 3, 3, 8), dtype=key_value_dtype)
+        expected_dtype = backend.result_type(*dtypes)
 
         self.assertDType(
             knn.dot_product_attention(query, key, value), expected_dtype