转换规则 No. 20 add torch.lu matcher

paconvert/api_mapping.json

@@ -8785,5 +8785,18 @@
     "kwargs_change": {
       "eps": "epsilon"
     }
+  },
+  "torch.lu": {
+    "Matcher": "LuMatcher",
+    "paddle_api": "paddle.linalg.lu",
+    "args_list": [
+      "A",
+      "pivot",
+      "get_infos",
+      "out"
+    ],
+    "kwargs_change": {
+      "A": "x"
+    }
   }
 }

paconvert/api_matcher.py

@@ -3647,6 +3647,47 @@ def generate_code(self, kwargs):
         return GenericMatcher.generate_code(self, kwargs)
 
 
+class LuMatcher(BaseMatcher):
+    def generate_code(self, kwargs):
+        out_v = kwargs.pop("out") if "out" in kwargs else None
+
+        if out_v:
+
+            out_3_var = "get_infos" in kwargs and kwargs["get_infos"] != "(False)"
+            new_kwargs = {}
+            new_kwargs["x"] = kwargs.pop("A")
+            new_kwargs.update(kwargs)
+
+            if out_3_var:
+                API_TEMPLATE = textwrap.dedent(
+                    """
+                    tmp_lu, tmp_p, tmp_info = {}({})
+                    paddle.assign(tmp_lu, {}[0]), paddle.assign(tmp_p, {}[1]), paddle.assign(tmp_info, {}[2])
+                    """
+                )
+                code = API_TEMPLATE.format(
+                    self.get_paddle_api(),
+                    self.kwargs_to_str(new_kwargs),
+                    out_v,
+                    out_v,
+                    out_v,
+                )
+            else:
+                API_TEMPLATE = textwrap.dedent(
+                    """
+                    tmp_lu, tmp_p = {}({})
+                    paddle.assign(tmp_lu, {}[0]), paddle.assign(tmp_p, {}[1])
+                    """
+                )
+                code = API_TEMPLATE.format(
+                    self.get_paddle_api(), self.kwargs_to_str(new_kwargs), out_v, out_v
+                )
+
+            return code
+
+        return GenericMatcher.generate_code(self, kwargs)
+
+
 class RandomSplitMatcher(BaseMatcher):
     def generate_code(self, kwargs):
         API_TEMPLATE = textwrap.dedent(

tests/test_lu.py

@@ -0,0 +1,176 @@
+# Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import textwrap
+
+from apibase import APIBase
+
+obj = APIBase("torch.lu")
+
+
+def test_case_1():
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        A = torch.tensor(
+            [
+                [
+                    [0.3591, -0.0479, -0.2174],
+                    [-0.6957, -1.4667, 1.4384],
+                    [0.0735, 0.1147, 0.0513],
+                ],
+                [
+                    [-1.2565, -2.1263, 0.8075],
+                    [-0.3665, -3.3540, -0.9417],
+                    [-0.1299, -0.0689, -0.6207],
+                ],
+            ]
+        )
+        A_LU, pivots = torch.lu(A)
+        """
+    )
+    obj.run(pytorch_code, ["A_LU", "pivots"])
+
+
+def test_case_2():
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        A = torch.tensor(
+            [
+                [
+                    [0.3591, -0.0479, -0.2174],
+                    [-0.6957, -1.4667, 1.4384],
+                    [0.0735, 0.1147, 0.0513],
+                ],
+                [
+                    [-1.2565, -2.1263, 0.8075],
+                    [-0.3665, -3.3540, -0.9417],
+                    [-0.1299, -0.0689, -0.6207],
+                ],
+            ]
+        )
+        A_LU, pivots, info = torch.lu(A, get_infos=True)
+        """
+    )
+    obj.run(pytorch_code, ["A_LU", "pivots", "info"])
+
+
+def test_case_3():
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        A = torch.tensor(
+            [
+                [
+                    [0.3591, -0.0479, -0.2174],
+                    [-0.6957, -1.4667, 1.4384],
+                    [0.0735, 0.1147, 0.0513],
+                ],
+                [
+                    [-1.2565, -2.1263, 0.8075],
+                    [-0.3665, -3.3540, -0.9417],
+                    [-0.1299, -0.0689, -0.6207],
+                ],
+            ]
+        )
+        A_LU, pivots, info = torch.lu(A, pivot=True, get_infos=True)
+        """
+    )
+    obj.run(pytorch_code, ["A_LU", "pivots", "info"])
+
+
+def test_case_4():
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        A = torch.tensor(
+            [
+                [
+                    [0.3591, -0.0479, -0.2174],
+                    [-0.6957, -1.4667, 1.4384],
+                    [0.0735, 0.1147, 0.0513],
+                ],
+                [
+                    [-1.2565, -2.1263, 0.8075],
+                    [-0.3665, -3.3540, -0.9417],
+                    [-0.1299, -0.0689, -0.6207],
+                ],
+            ]
+        )
+        A_LU = torch.empty_like(A)
+        pivots = torch.empty((2, 3), dtype=torch.int32)
+        info = torch.empty((2, ), dtype=torch.int32)
+        torch.lu(A, pivot=True, get_infos=True, out=(A_LU, pivots, info))
+        """
+    )
+    obj.run(pytorch_code, ["A_LU", "pivots", "info"])
+
+
+def test_case_5():
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        A = torch.tensor(
+            [
+                [
+                    [0.3591, -0.0479, -0.2174],
+                    [-0.6957, -1.4667, 1.4384],
+                    [0.0735, 0.1147, 0.0513],
+                ],
+                [
+                    [-1.2565, -2.1263, 0.8075],
+                    [-0.3665, -3.3540, -0.9417],
+                    [-0.1299, -0.0689, -0.6207],
+                ],
+            ]
+        )
+        A_LU = torch.empty_like(A)
+        pivots = torch.empty((2, 3), dtype=torch.int32)
+        torch.lu(A, pivot=True, get_infos=False, out=(A_LU, pivots))
+        """
+    )
+    obj.run(pytorch_code, ["A_LU", "pivots"])
+
+
+def test_case_6():
+    pytorch_code = textwrap.dedent(
+        """
+        import torch
+
+        A = torch.tensor(
+            [
+                [
+                    [0.3591, -0.0479, -0.2174],
+                    [-0.6957, -1.4667, 1.4384],
+                    [0.0735, 0.1147, 0.0513],
+                ],
+                [
+                    [-1.2565, -2.1263, 0.8075],
+                    [-0.3665, -3.3540, -0.9417],
+                    [-0.1299, -0.0689, -0.6207],
+                ],
+            ]
+        )
+        A_LU = torch.empty_like(A)
+        pivots = torch.empty((2, 3), dtype=torch.int32)
+        LU = torch.lu(A, pivot=True, get_infos=False, out=(A_LU, pivots))
+        """
+    )
+    obj.run(pytorch_code, ["A_LU", "pivots", "LU"])