[SOT] Use same list when create dynamic axes meta to ensure SymbolicVariable are same (#71932)

SigureMo · web-flow · commit bd3696489668 · 2025-03-27T10:08:40.000+08:00
diff --git a/python/paddle/jit/sot/infer_meta.py b/python/paddle/jit/sot/infer_meta.py
@@ -36,7 +36,13 @@
 from paddle.framework import use_pir_api
 from paddle.utils import flatten, is_sequence
 
-from .utils import Cache, Singleton, map_if_extend, meta_str
+from .utils import (
+    Cache,
+    Singleton,
+    map_if_extend,
+    meta_str,
+    update_list_inplace,
+)
 
 DynamicSymbolT = TypeVar("DynamicSymbolT")
 SOT_INFER_META_INNER_VAR = "___SOT_INFER_META_INNER_VAR"
@@ -139,8 +145,12 @@ def with_dynamic_axes(self, dynamic_axes: list[int]) -> MetaInfo:
             SymbolicInt() if i in dynamic_axes else dim
             for i, dim in enumerate(self.shape)
         ]
+        # NOTE(SigureMo): Ensure output meta.shape is same list object as
+        # self.shape to avoid create two different data proxy for tensor.shape.
+        # It will caused create a new SymbolicVariable when it's a dynamic dim.
+        self.shape = update_list_inplace(self.shape, shape)
         return MetaInfo(
-            shape,
+            self.shape,
             self.dtype,
             self.stop_gradient,
             self.name,
diff --git a/python/paddle/jit/sot/opcode_translator/executor/dispatcher.py b/python/paddle/jit/sot/opcode_translator/executor/dispatcher.py
@@ -105,7 +105,6 @@ def type(self) -> tuple[type[Any], ...]:
         return convert_annotation_to_type(self.annotation)
 
     def match_arg(self, arg: Any) -> bool:
-        # TODO: support VAR_KEYWORD
         if self.kind == inspect.Parameter.VAR_POSITIONAL:
             is_tuple = isinstance(arg, tuple)
             return is_tuple and all(isinstance(a, self.type) for a in arg)
diff --git a/python/paddle/jit/sot/utils/__init__.py b/python/paddle/jit/sot/utils/__init__.py
@@ -105,4 +105,5 @@
     no_eval_frame,
     printable,
     switch_symbol_registry,
+    update_list_inplace,
 )
diff --git a/python/paddle/jit/sot/utils/utils.py b/python/paddle/jit/sot/utils/utils.py
@@ -430,6 +430,14 @@ def do_until_stop_iteration(fn: Callable[[], T]) -> list[T]:
     return res
 
 
+def update_list_inplace(
+    original_list: list[T], new_contents: list[T]
+) -> list[T]:
+    original_list.clear()
+    original_list.extend(new_contents)
+    return original_list
+
+
 def get_obj_stable_repr(obj) -> str:
     if hasattr(obj, '__qualname__'):
         return obj.__qualname__

Original file line number	Diff line number	Diff line change
`@@ -105,4 +105,5 @@`
`105`	`105`	`no_eval_frame,`
`106`	`106`	`printable,`
`107`	`107`	`switch_symbol_registry,`
	`108`	`+ update_list_inplace,`
`108`	`109`	`)`