From be476519fb708f07ecea7456fe867e26c8fdf98e Mon Sep 17 00:00:00 2001 From: Mateusz Mikolajczyk Date: Mon, 18 Dec 2023 11:32:54 +0100 Subject: [PATCH] [PT FE] ScaledDotProductAttention fix issue with other float dtypes (#21638) * Fix scale input * Improve torch scaled dot tests --- .../src/op/scaled_dot_product_attention.cpp | 2 +- .../test_scaled_dot_product_attention.py | 19 ++++++++++--------- 2 files changed, 11 insertions(+), 10 deletions(-) diff --git a/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp b/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp index ed8ecc4b0f8..5a9e0d1f95c 100644 --- a/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp +++ b/src/frontends/pytorch/src/op/scaled_dot_product_attention.cpp @@ -32,7 +32,7 @@ std::shared_ptr translate_scaled_dot_product_attention_common(const No inputs.push_back(context.mark_node(std::make_shared(zero, query))); } if (!context.input_is_none(6)) - inputs.push_back(context.get_input(6)); + inputs.push_back(context.mark_node(std::make_shared(context.get_input(6), query))); return context.mark_node(std::make_shared(inputs, is_causal)); } diff --git a/tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py b/tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py index 67e6322ef0d..d442b9e86a4 100644 --- a/tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py +++ b/tests/layer_tests/pytorch_tests/test_scaled_dot_product_attention.py @@ -9,20 +9,20 @@ from pytorch_layer_test_class import PytorchLayerTest class TestScaledDotProductAttention(PytorchLayerTest): - def _prepare_input(self): - return (np.random.randn(1, 2, 8, 4).astype(np.float32), np.random.randn(1, 2, 8, 4).astype(np.float32), np.random.randn(1, 2, 8, 4).astype(np.float32)) + def _prepare_input(self, dtype): + return (np.random.randn(1, 2, 8, 4).astype(dtype), np.random.randn(1, 2, 8, 4).astype(dtype), np.random.randn(1, 2, 8, 4).astype(dtype)) - def create_model(self, mask, is_causal, scale): + def create_model(self, mask, is_causal, scale, dtype): import torch.nn.functional as F import torch class aten_scaled_dot_product_atten(torch.nn.Module): - def __init__(self, mask=False, is_causal=False, scale=False) -> None: + def __init__(self, mask=False, is_causal=False, scale=False, dtype=np.float32) -> None: super().__init__() self.mask = None if not mask else torch.from_numpy( - np.random.randint(0, 2, (8, 8)).astype(np.float32)) + np.random.randint(0, 2, (8, 8)).astype(dtype)) self.is_causal = is_causal if is_causal and mask: self.mask.to(torch.bool) @@ -36,13 +36,14 @@ class TestScaledDotProductAttention(PytorchLayerTest): ref_net = None - return aten_scaled_dot_product_atten(mask, is_causal, scale), ref_net, 'aten::scaled_dot_product_attention' + return aten_scaled_dot_product_atten(mask, is_causal, scale, dtype), ref_net, 'aten::scaled_dot_product_attention' @pytest.mark.nightly @pytest.mark.precommit @pytest.mark.precommit_fx_backend @pytest.mark.parametrize(['mask', 'is_causal'], [(False, False), (False, True), (True, True), (True, False)]) @pytest.mark.parametrize('scale', [False, True]) - def test_scaled_dot_product_atten(self, ie_device, precision, ir_version, mask, is_causal, scale): - self._test(*self.create_model(mask, is_causal, scale), - ie_device, precision, ir_version) + @pytest.mark.parametrize("dtype", (np.float32, np.float64)) + def test_scaled_dot_product_atten(self, ie_device, precision, ir_version, mask, is_causal, scale, dtype): + self._test(*self.create_model(mask, is_causal, scale, dtype), + ie_device, precision, ir_version, kwargs_to_prepare_input={"dtype": dtype})