add UnbindScaledDotProductModel torch2 quantization test

Author: ruro

tests/cross_fw/test_templates/helpers.py

@@ -440,6 +440,15 @@ def forward(self, query, key, value):
         return nn.functional.scaled_dot_product_attention(query, key, value)
 
 
+class UnbindScaledDotProductAttentionModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        query, key, value = x.unbind(0)
+        return nn.functional.scaled_dot_product_attention(query, key, value)
+
+
 class DepthwiseConvTestModel(nn.Module):
     INPUT_SIZE = [1, 2, 4, 4]
 

tests/torch2/data/function_hook/quantization/test_quantized_graphs/unbind_scaled_dot_product_attention_model.dot

@@ -0,0 +1,19 @@
+strict digraph {
+x [id=0, type="nncf_model_input", metatype=PTInputNoopMetatype];
+"/unbind/0" [id=1, type=unbind, metatype=PTSplitMetatype];
+"__nncf_hooks.pre_hooks./scaled_dot_product_attention/0__0.0._scale_param_storage" [id=2, type="nncf_model_const", metatype=PTConstNoopMetatype];
+"pre_hook__-scaled_dot_product_attention-0__0[0]/symmetric_quantize/0" [id=3, type="symmetric_quantize", metatype=UnknownMetatype];
+"__nncf_hooks.pre_hooks./scaled_dot_product_attention/0__1.0._scale_param_storage" [id=4, type="nncf_model_const", metatype=PTConstNoopMetatype];
+"pre_hook__-scaled_dot_product_attention-0__1[0]/symmetric_quantize/0" [id=5, type="symmetric_quantize", metatype=UnknownMetatype];
+"/scaled_dot_product_attention/0" [id=6, type="scaled_dot_product_attention", metatype=PTScaledDotProductAttentionMetatype];
+output [id=7, type="nncf_model_output", metatype=PTOutputNoopMetatype];
+x -> "/unbind/0" [dtype=float, shape="(3, 1, 8, 16)", out_port_id=0, in_port_id=0];
+"/unbind/0" -> "pre_hook__-scaled_dot_product_attention-0__0[0]/symmetric_quantize/0" [dtype=float, shape="(1, 8, 16)", out_port_id=0, in_port_id=0];
+"/unbind/0" -> "pre_hook__-scaled_dot_product_attention-0__1[0]/symmetric_quantize/0" [dtype=float, shape="(1, 8, 16)", out_port_id=1, in_port_id=0];
+"/unbind/0" -> "/scaled_dot_product_attention/0" [dtype=float, shape="(1, 8, 16)", out_port_id=2, in_port_id=2];
+"__nncf_hooks.pre_hooks./scaled_dot_product_attention/0__0.0._scale_param_storage" -> "pre_hook__-scaled_dot_product_attention-0__0[0]/symmetric_quantize/0" [dtype=float, shape="(1,)", out_port_id=0, in_port_id=4];
+"pre_hook__-scaled_dot_product_attention-0__0[0]/symmetric_quantize/0" -> "/scaled_dot_product_attention/0" [dtype=float, shape="(1, 8, 16)", out_port_id=0, in_port_id=0];
+"__nncf_hooks.pre_hooks./scaled_dot_product_attention/0__1.0._scale_param_storage" -> "pre_hook__-scaled_dot_product_attention-0__1[0]/symmetric_quantize/0" [dtype=float, shape="(1,)", out_port_id=0, in_port_id=4];
+"pre_hook__-scaled_dot_product_attention-0__1[0]/symmetric_quantize/0" -> "/scaled_dot_product_attention/0" [dtype=float, shape="(1, 8, 16)", out_port_id=0, in_port_id=1];
+"/scaled_dot_product_attention/0" -> output [dtype=float, shape="(1, 8, 16)", out_port_id=0, in_port_id=0];
+}

tests/torch2/function_hook/quantization/test_quantized_graphs.py

@@ -25,6 +25,7 @@
 from tests.cross_fw.test_templates.helpers import EmbeddingModel
 from tests.cross_fw.test_templates.helpers import RoPEModel
 from tests.cross_fw.test_templates.helpers import ScaledDotProductAttentionModel
+from tests.cross_fw.test_templates.helpers import UnbindScaledDotProductAttentionModel
 from tests.torch import test_models
 from tests.torch.quantization.test_algo_quantization import SharedLayersModel
 from tests.torch.test_compressed_graph import ModelDesc
@@ -45,6 +46,14 @@
         ),
         {},
     ),
+    (
+        ModelDesc(
+            "unbind_scaled_dot_product_attention_model",
+            UnbindScaledDotProductAttentionModel,
+            {"x": [3, 1, 8, 16]},
+        ),
+        {},
+    ),
     (ModelDesc("shared_model", SharedLayersModel, [1, 1, 5, 6]), {}),
     (ModelDesc("alexnet", test_models.AlexNet, [1, 3, 32, 32]), {}),
     (ModelDesc("lenet", test_models.LeNet, [1, 3, 32, 32]), {}),