[SOT] Improve dispatch for operator.is_ (#73066)

0x3878f · web-flow · commit 4877401b95bc · 2025-06-10T03:02:43.000+08:00
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py b/python/paddle/jit/sot/opcode_translator/executor/variable_dispatch.py
@@ -983,16 +983,17 @@ def str_endswith(
     ),
 )
 
+
 # VariableBase
-Dispatcher.register(
-    operator.is_,
-    ("VariableBase", "VariableBase"),
-    lambda var, other: ConstantVariable(
+@Dispatcher.register_decorator(operator.is_)
+def is_func(var: VariableBase, other: VariableBase):
+    if var.get_py_type() is not other.get_py_type():
+        return ConstantVariable(False, var.graph, DummyTracker([var, other]))
+    return ConstantVariable(
         var.get_py_value() is other.get_py_value(),
         var.graph,
-        tracker=DummyTracker([var, other]),
-    ),
-)
+        DummyTracker([var, other]),
+    )
 
 
 @Dispatcher.register_decorator(operator.is_not)
diff --git a/test/sot/test_builtin_dispatch.py b/test/sot/test_builtin_dispatch.py
@@ -173,6 +173,11 @@ def test_builtin_type_check_eq():
     return eq_results, ne_results
 
 
+@check_no_breakgraph
+def test_is(x, y):
+    return x is y
+
+
 class TestBuiltinDispatch(TestCaseBase):
     def test_dispatch_len(self):
         self.assert_results(dispatch_len, paddle.to_tensor([1, 2, 3]))
@@ -343,6 +348,20 @@ def test_dispatch_all(self):
         self.assert_results(test_all, d_false)
         self.assert_results(test_all_iter, l_true_and_false)
 
+    def test_dispatch_is(self):
+        x = paddle.ones(shape=[1, 2])
+        y = paddle.ones(shape=[1, 2])
+        # TODO(wangmingkai02): support comparison of same tensor object
+        # self.assert_results(test_is, x, x)
+        # self.assert_results(test_is, [x], [x])
+        self.assert_results(test_is, x, y)
+        self.assert_results(test_is, x, None)
+        self.assert_results(test_is, [x], x)
+        self.assert_results(test_is, None, x)
+        self.assert_results(test_is, [x], None)
+        self.assert_results(test_is, None, [x])
+        self.assert_results(test_is, None, None)
+
 
 def run_getattr(x: paddle.Tensor):
     attr = 'dtype'
