在scaled_dot_product_attention函数中，改为如果输入是3D，输出也是3D

paconvert/api_matcher.py

@@ -4799,13 +4799,31 @@ def generate_code(self, kwargs):
         query = kwargs.pop("query")
         key = kwargs.pop("key")
         value = kwargs.pop("value")
+
         API_TEMPLATE = textwrap.dedent(
             """
-            {}({}.transpose([0, 2, 1, 3]), {}.transpose([0, 2, 1, 3]), {}.transpose([0, 2, 1, 3]), {}).transpose([0, 2, 1, 3])
+            {} = list(range({}.ndim))
+            {}[-2], {}[-3] = {}[-3], {}[-2]
+            {}({}.transpose({}), {}.transpose({}), {}.transpose({}), {}).transpose({})
             """
         )
+        perm = get_unique_name("perm")
         code = API_TEMPLATE.format(
-            self.get_paddle_api(), query, key, value, self.kwargs_to_str(kwargs)
+            perm,
+            query,
+            perm,
+            perm,
+            perm,
+            perm,
+            self.get_paddle_api(),
+            query,
+            perm,
+            key,
+            perm,
+            value,
+            perm,
+            self.kwargs_to_str(kwargs),
+            perm,
         )
         return code
 

tests/test_scaled_dot_product_attention.py

@@ -182,3 +182,17 @@ def test_case_9():
         """
     )
     obj.run(pytorch_code, ["result"], rtol=1e-5)
+
+
+def test_case_10():
+    pytorch_code = textwrap.dedent(
+        """
+        import numpy as np
+        import torch
+        np.random.seed(100)
+        x = np.random.rand(8, 128, 64)
+        query = torch.tensor(x, dtype=torch.float16)
+        result = torch.nn.functional.scaled_dot_product_attention(query, query, query)
+        """
+    )
+    obj.run(pytorch_code, ["result"], rtol=1e-5)