[SOT] Split NumpyVariable for separate array/number handling (#71216)

Author: DrRyanHuang

python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py

@@ -1064,7 +1064,7 @@ def is_not_func(var: VariableBase, other: VariableBase):
                 binary_fn,
                 (
                     "TensorVariable",
-                    "TensorVariable | SymbolicVariable | ConstantVariable | NumpyVariable",
+                    "TensorVariable | SymbolicVariable | ConstantVariable | NumpyNumberVariable",
                 ),
                 partial(
                     lambda magic_name, var, other: var.graph.call_tensor_method(
@@ -1092,7 +1092,7 @@ def tensor_mod_dispatcher(
                 Dispatcher.register(
                     binary_fn,
                     (
-                        "SymbolicVariable | ConstantVariable | NumpyVariable",
+                        "SymbolicVariable | ConstantVariable | NumpyNumberVariable",
                         "TensorVariable",
                     ),
                     partial(

python/paddle/jit/sot/opcode_translator/executor/variables/__init__.py

@@ -27,6 +27,8 @@
     GlobalVariable,
     ModuleVariable,
     NullVariable,
+    NumpyArrayVariable,
+    NumpyNumberVariable,
     NumpyVariable,
     ObjectVariable,
     ParameterVariable,

python/paddle/jit/sot/opcode_translator/executor/variables/basic.py

@@ -1202,49 +1202,71 @@ def main_info(self) -> dict[str, Any]:
     def get_py_value(self, allow_tensor=False) -> Any:
         return self.value
 
+    @staticmethod
+    def format_dtype(dtype: np.dtype):
+        return f"np.{dtype}"
+
+    @staticmethod
+    def format_number(number: np.number):
+        return f"{NumpyVariable.format_dtype(number.dtype)}({number.item()})"
+
+    def make_stringified_guard(self) -> None:
+        raise NotImplementedError
+
+    @VariableFactory.register_from_value()
+    def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
+        if isinstance(value, (np.number)):
+            return NumpyNumberVariable(value, graph, tracker)
+        if isinstance(value, (np.ndarray)):
+            return NumpyArrayVariable(value, graph, tracker)
+        return None
+
+
+class NumpyNumberVariable(NumpyVariable):
     @check_guard
     def make_stringified_guard(self) -> list[StringifiedExpression]:
         frame_value_tracer = self.tracker.trace_value_from_frame()
         obj_free_var_name = f"__{self.id}"
 
-        def format_dtype(dtype: np.dtype):
-            return f"np.{dtype}"
+        dtype_guard = StringifiedExpression(
+            f"{{}}.dtype == {NumpyVariable.format_dtype(self.get_py_value().dtype)}",
+            [frame_value_tracer],
+            union_free_vars(frame_value_tracer.free_vars, {"np": np}),
+        )
+
+        return [
+            dtype_guard,
+            StringifiedExpression(
+                f"{{}} == {NumpyVariable.format_number(self.get_py_value())}",
+                [frame_value_tracer],
+                union_free_vars(frame_value_tracer.free_vars, {"np": np}),
+            ),
+        ]
+
 
-        def format_number(number: np.number):
-            return f"{format_dtype(number.dtype)}({number.item()})"
+class NumpyArrayVariable(NumpyVariable):
+    @check_guard
+    def make_stringified_guard(self) -> list[StringifiedExpression]:
+        frame_value_tracer = self.tracker.trace_value_from_frame()
+        obj_free_var_name = f"__{self.id}"
 
         dtype_guard = StringifiedExpression(
-            f"{{}}.dtype == {format_dtype(self.get_py_value().dtype)}",
+            f"{{}}.dtype == {NumpyVariable.format_dtype(self.get_py_value().dtype)}",
             [frame_value_tracer],
             union_free_vars(frame_value_tracer.free_vars, {"np": np}),
         )
-        if isinstance(self.get_py_value(), np.number):
-            return [
-                dtype_guard,
-                StringifiedExpression(
-                    f"{{}} == {format_number(self.get_py_value())}",
-                    [frame_value_tracer],
-                    union_free_vars(frame_value_tracer.free_vars, {"np": np}),
-                ),
-            ]
-        else:
-            return [
-                dtype_guard,
-                StringifiedExpression(
-                    f"({{}} == {obj_free_var_name}).all()",
-                    [frame_value_tracer],
-                    union_free_vars(
-                        frame_value_tracer.free_vars,
-                        {obj_free_var_name: self.get_py_value()},
-                    ),
-                ),
-            ]
 
-    @VariableFactory.register_from_value()
-    def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
-        if isinstance(value, (np.ndarray, np.number)):
-            return NumpyVariable(value, graph, tracker)
-        return None
+        return [
+            dtype_guard,
+            StringifiedExpression(
+                f"({{}} == {obj_free_var_name}).all()",
+                [frame_value_tracer],
+                union_free_vars(
+                    frame_value_tracer.free_vars,
+                    {obj_free_var_name: self.get_py_value()},
+                ),
+            ),
+        ]
 
 
 class NullVariable(VariableBase):