support scaled dot product attention

Author: danielenricocahall

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -264,7 +264,6 @@ NNOpsDynamicShapeTest::test_categorical_crossentropy
 NNOpsDynamicShapeTest::test_multi_hot_dtype_
 NNOpsCorrectnessTest::test_conv_transpose_
 NNOpsCorrectnessTest::test_ctc_decode
-NNOpsCorrectnessTest::test_dot_product_attention_
 NNOpsCorrectnessTest::test_multi_hot_
 NNOpsCorrectnessTest::test_binary_crossentropy
 NNOpsCorrectnessTest::test_categorical_crossentropy
@@ -282,7 +281,6 @@ NNOpsCorrectnessTest::test_rms_normalization_10.0
 NNOpsDtypeTest::test_ctc_decode
 NNOpsDtypeTest::test_glu_
 NNOpsDtypeTest::test_polar_
-NNOpsDtypeTest::test_dot_product_attention_
 NNOpsDynamicShapeTest::test_glu
 NNOpsBehaviorTest::test_invalid_strategy_ctc_decode
 NNOpsBehaviorTest::test_logit_recovery_binary_crossentropy

keras/src/backend/openvino/nn.py

@@ -693,9 +693,47 @@ def dot_product_attention(
     flash_attention=None,
     attn_logits_soft_cap=None,
 ):
-    raise NotImplementedError(
-        "`dot_product_attention` is not supported with openvino backend"
+    if bias is not None:
+        raise NotImplementedError(
+            "`dot_product_attention` with `bias` is not supported "
+            "with openvino backend"
+        )
+    if flash_attention is not None:
+        raise NotImplementedError(
+            "`dot_product_attention` with `flash_attention` is not supported "
+            "with openvino backend"
+        )
+    if attn_logits_soft_cap is not None:
+        raise NotImplementedError(
+            "`dot_product_attention` with `attn_logits_soft_cap` is not "
+            "supported with openvino backend"
+        )
+    query = get_ov_output(query)
+    key = get_ov_output(key)
+    value = get_ov_output(value)
+    if query.get_element_type() != key.get_element_type():
+        ov_type = OPENVINO_DTYPES[backend.floatx()]
+        query = ov_opset.convert(query, ov_type)
+        key = ov_opset.convert(key, ov_type)
+    if value.get_element_type() != query.get_element_type():
+        ov_type = OPENVINO_DTYPES[backend.floatx()]
+        value = ov_opset.convert(value, ov_type)
+    axes_const = ov_opset.constant([0, 2, 1, 3], Type.i32).output(0)
+
+    query = ov_opset.transpose(query, axes_const)
+    key = ov_opset.transpose(key, axes_const)
+    value = ov_opset.transpose(value, axes_const)
+    mask = get_ov_output(mask) if mask is not None else None
+    scale = (
+        get_ov_output(scale, query.get_element_type())
+        if scale is not None
+        else None
+    )
+    dpa = ov_opset.scaled_dot_product_attention(
+        query, key, value, attention_mask=mask, scale=scale, causal=is_causal
     )
+    dpa = ov_opset.transpose(dpa, axes_const)
+    return OpenVINOKerasTensor(dpa.output(0))
 
 
 def unfold(input, kernel_size, dilation=1, padding=0, stride=1):

keras/src/ops/nn_test.py

@@ -2448,7 +2448,7 @@ def test_dot_product_attention(
             mask = mask[None, None, ...]
             mask = np.tile(mask, (2, 4, 1, 1))
         if bias is not None:
-            if backend.backend() == "torch":
+            if backend.backend() in ("torch", "openvino"):
                 self.skipTest(
                     "torch does not support `bias` with `dot_product_attention`"
                 )