update

Author: zrr1999

sot/opcode_translator/executor/dispatcher.py

@@ -5,8 +5,6 @@
 from functools import cached_property, reduce
 from typing import TYPE_CHECKING, Any, Callable, Dict, Tuple, TypeVar
 
-import paddle
-
 from ...utils import InnerError
 
 if TYPE_CHECKING:

sot/opcode_translator/executor/variable_dispatch.py

@@ -99,7 +99,7 @@
     lambda var: var.bool(),
 )
 
-# VariableBase
+# TensorVariable
 Dispatcher.register(
     paddle.is_tensor,
     ("TensorVariable",),

sot/opcode_translator/executor/variables/basic.py

@@ -15,7 +15,6 @@
     BreakGraphError,
     NameGenerator,
     NotImplementException,
-    NotImplementException,
     log_do,
     paddle_tensor_methods,
 )
@@ -34,21 +33,34 @@
 if TYPE_CHECKING:
     from ..function_graph import FunctionGraph
 
-DTYPE_ABBRS = {
+
+FP_DTYPE_ABBRS = {
     paddle.bfloat16: 'bfloat16',
     paddle.float64: 'float64',
     paddle.float32: 'float32',
     paddle.float16: 'float16',
+}
+
+CP_DTYPE_ABBRS = {
     paddle.complex64: 'complex64',
     paddle.complex128: 'complex128',
+}
+
+INT_DTYPE_ABBRS = {
     paddle.int8: 'int8',
     paddle.int16: 'int16',
     paddle.int32: 'int32',
     paddle.int64: 'int64',
-    paddle.bool: 'bool',
     paddle.uint8: 'uint8',
 }
 
+DTYPE_ABBRS = {
+    **FP_DTYPE_ABBRS,
+    **CP_DTYPE_ABBRS,
+    **INT_DTYPE_ABBRS,
+    paddle.bool: 'bool',
+}
+
 
 class ConstantVariable(VariableBase):
     def __init__(
@@ -174,7 +186,6 @@ def main_info(self) -> dict[str, Any]:
             "dtype": DTYPE_ABBRS[self.meta.dtype],
             "stop_gradient": self.meta.stop_gradient,
             "var_name": self.var_name,
-            "var_name": self.var_name,
         }
 
     def __getitem__(self, key):
@@ -235,28 +246,17 @@ def is_tensor(self):
 
     def is_complex(self):
         dtype = self.meta.dtype
-        is_cp_dtype = dtype == paddle.complex64 or dtype == paddle.complex128
+        is_cp_dtype = dtype in CP_DTYPE_ABBRS
         return ConstantVariable.wrap_literal(is_cp_dtype)
 
     def is_integer(self):
         dtype = self.meta.dtype
-        is_int_dtype = (
-            dtype == paddle.int8
-            or dtype == paddle.uint8
-            or dtype == paddle.int16
-            or dtype == paddle.int32
-            or dtype == paddle.int64
-        )
+        is_int_dtype = dtype in INT_DTYPE_ABBRS
         return ConstantVariable.wrap_literal(is_int_dtype)
 
     def is_floating_point(self):
         dtype = self.meta.dtype
-        is_fp_dtype = (
-            dtype == paddle.float32
-            or dtype == paddle.float64
-            or dtype == paddle.float16
-            or dtype == paddle.bfloat16
-        )
+        is_fp_dtype = dtype in FP_DTYPE_ABBRS
         return ConstantVariable.wrap_literal(is_fp_dtype)
 
     def getattr(self, name: str):
@@ -422,7 +422,36 @@ def get_value(self) -> Any:
         return self.value
 
     def make_stringify_guard(self) -> StringifyExpression:
-        raise NotImplementException("We can not stringify numpy variable")
+        if isinstance(self.get_value(), np.number):
+            assert not isinstance(
+                self.tracker, DummyTracker
+            ), "Can not make guard from dummy tracker"
+
+            frame_value_tracer = self.tracker.trace_value_from_frame()
+            log_do(
+                4,
+                lambda: print(
+                    f"[Guard]: guard_fn for {self}, tracker={self.tracker.__class__.__name__}, value={frame_value_tracer.expr}"
+                ),
+            )
+
+            def format_dtype(dtype: np.dtype):
+                return f"np.{str(dtype)}"
+
+            def format_number(number: np.number):
+                return f"{format_dtype(number.dtype)}({str(number.item())})"
+
+            return StringifyExpression(
+                f"{frame_value_tracer.expr} == {format_number(self.get_value())}",
+                union_free_vars(frame_value_tracer.free_vars, {"np": np}),
+            ) & StringifyExpression(
+                f"{frame_value_tracer.expr}.dtype == {format_dtype(self.get_value().dtype)}",
+                union_free_vars(frame_value_tracer.free_vars, {"np": np}),
+            )
+        else:
+            raise NotImplementException(
+                "We can not stringify numpy variable when value is np.ndarray"
+            )
 
     @VariableFactory.register_from_value()
     def from_value(value: Any, graph: FunctionGraph | None, tracker: Tracker):