[SOT][DynamicShape] Fix SOT SymbolicVariable magic method dispatch (#65235)

SigureMo · web-flow · commit 995c193d8ac5 · 2024-06-18T14:34:37.000+08:00
diff --git a/python/paddle/base/dygraph/tensor_patch_methods.py b/python/paddle/base/dygraph/tensor_patch_methods.py
@@ -989,7 +989,7 @@ def pre_deal_index(self, item):
     def __getitem__(
         self,
         item: TensorIndex,
-    ):
+    ) -> Tensor:
         item = pre_deal_index(self, item)
         return self._getitem_dygraph(item)
 
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py b/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
@@ -27,6 +27,7 @@
     UNARY_OPS,
     magic_method_builtin_dispatch,
 )
+from ...utils.paddle_api_config import get_tensor_methods
 from .dispatch_functions import (
     operator_in,
     operator_is_none,
@@ -973,6 +974,8 @@ def tensor_mod_dispatcher(
 # Symbolic
 for binary_fn in BINARY_OPS:
     for magic_method in magic_method_builtin_dispatch(binary_fn):
+        if magic_method.name not in get_tensor_methods():
+            continue
         # skip all inplace magic method name, we will dispatch it to non-inplace
         # magic methods
         if magic_method.is_inplace:
@@ -997,8 +1000,8 @@ def tensor_mod_dispatcher(
                 binary_fn,
                 ("ConstantVariable", "SymbolicVariable"),
                 partial(
-                    lambda magic_name, var, other: var.graph.call_symbolic_method(
-                        magic_name, var, other
+                    lambda reverse_magic_name, var, other: var.graph.call_symbolic_method(
+                        reverse_magic_name, other, var
                     ),
                     magic_method.name,
                 ),
diff --git a/python/paddle/nn/functional/norm.py b/python/paddle/nn/functional/norm.py
@@ -352,8 +352,9 @@ def layer_norm(
     normalized_ndim = len(normalized_shape)
     begin_norm_axis = input_ndim - normalized_ndim
     if input_ndim < normalized_ndim or (
-        isinstance(normalized_shape[0], int)
-        and input_shape[begin_norm_axis:] != normalized_shape
+        not paddle.utils.is_same_shape(
+            input_shape[begin_norm_axis:], normalized_shape
+        )
     ):
         str_normalized_shape = str(normalized_shape)
         raise ValueError(
diff --git a/python/paddle/utils/__init__.py b/python/paddle/utils/__init__.py
@@ -43,6 +43,7 @@
     get_int_tensor_list,
     get_shape_tensor_inputs,
     hold_mutable_vars,
+    is_same_shape,
     is_sequence,
     map_structure,
     pack_sequence_as,
diff --git a/python/paddle/utils/layers_utils.py b/python/paddle/utils/layers_utils.py
@@ -12,6 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from __future__ import annotations
+
 import copy
 import typing
 from collections import defaultdict
@@ -24,6 +26,7 @@
 from typing_extensions import TypeGuard
 
 import paddle
+from paddle._typing import ShapeLike
 from paddle.pir.core import convert_np_dtype_to_dtype_
 
 from ..base.data_feeder import check_dtype, convert_dtype
@@ -586,3 +589,28 @@ def get_inputs_outputs_in_block(block):
                     inner_outputs.add(out_var_name)
 
     return inner_inputs, inner_outputs
+
+
+def is_same_shape(shape1: ShapeLike, shape2: ShapeLike) -> bool:
+    """
+    Check whether two shapes are the same. Deal with the dynamic shape.
+    """
+    if paddle.in_dynamic_mode():
+        return shape1 == shape2
+
+    def is_tensor(x):
+        return isinstance(x, (paddle.static.Variable, paddle.pir.Value))
+
+    def is_dynamic_axis(axis):
+        return is_tensor(axis) or axis == -1
+
+    if is_tensor(shape1) or is_tensor(shape2):
+        return True
+    if len(shape1) != len(shape2):
+        return False
+    for s1, s2 in zip(shape1, shape2):
+        if is_dynamic_axis(s1) or is_dynamic_axis(s2):
+            continue
+        if s1 != s2:
+            return False
+    return True
