add no_grad

DrRyanHuang · DrRyanHuang · commit e198af8c04a5 · 2025-07-11T07:09:40.000Z
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py b/python/paddle/jit/sot/opcode_translator/executor/variables/basic.py
@@ -18,6 +18,7 @@
 import operator
 import sys
 import types
+from contextlib import AbstractContextManager
 from enum import Enum
 from functools import cached_property, reduce
 from typing import TYPE_CHECKING, Any
@@ -2572,3 +2573,77 @@ def from_value(value: object, graph: FunctionGraph, tracker: Tracker):
             class_var.tracker = GetAttrTracker(var, "__class__")
             return var
         return None
+
+
+class ContextManagerVariable(VariableBase):
+    def __init__(
+        self,
+        value: AbstractContextManager,
+        graph: FunctionGraph,
+        tracker: Tracker,
+    ):
+        super().__init__(graph=graph, tracker=tracker)
+        self.value = value
+
+    def get_py_type(self):
+        return self.value.__class__
+
+    def get_py_value(self, allow_tensor=False):
+        return self.value
+
+    @VariableFactory.register_from_value()
+    def from_value(value: object, graph: FunctionGraph, tracker: Tracker):
+        if isinstance(value, AbstractContextManager):
+            var = ContextManagerVariable(
+                value,
+                graph=graph,
+                tracker=tracker,
+            )
+            return var
+        return None
+
+
+class NoGradContextManagerVariable(ContextManagerVariable):
+    def getattr(self, name: str, default=None):
+        from .callable import InternalFunctionVariable, MethodVariable
+
+        if name == "__enter__":
+            from ..function_graph import NoGradContext
+
+            def no_grap_enter(self):
+                self.graph.sir_builder._current_statement_ctxs.append(
+                    NoGradContext()
+                )
+                return self
+
+            variable = InternalFunctionVariable(
+                no_grap_enter, self.graph, DummyTracker([])
+            )
+            return MethodVariable(
+                self, variable, self.graph, GetAttrTracker(self, "__enter__")
+            )
+
+        if name == "__exit__":
+
+            def no_grad_exit(self, exc, value, traceback):
+                self.graph.sir_builder._current_statement_ctxs.pop()
+                return ConstantVariable.wrap_literal(False, graph=self.graph)
+
+            variable = InternalFunctionVariable(
+                no_grad_exit, self.graph, DummyTracker([])
+            )
+            return MethodVariable(
+                self, variable, self.graph, GetAttrTracker(self, "__exit__")
+            )
+        return super().getattr(name, default)
+
+    @VariableFactory.register_from_value(successor="ContextManagerVariable")
+    def from_value(value: object, graph: FunctionGraph, tracker: Tracker):
+        if isinstance(value, paddle.no_grad):
+            var = NoGradContextManagerVariable(
+                value,
+                graph=graph,
+                tracker=tracker,
+            )
+            return var
+        return None
diff --git a/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py b/python/paddle/jit/sot/opcode_translator/executor/variables/callable.py
@@ -1191,6 +1191,22 @@ def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
         return None
 
 
+class InternalFunctionVariable(FunctionVariable):
+    # Handles internal implementations specific for SOT, such as passing hook functions.
+    # Refer to NoGradContextManagerVariable.getattr for usage examples.
+    def __init__(
+        self, fn: Callable[..., Any], graph: FunctionGraph, tracker: Tracker
+    ):
+        super().__init__(fn, graph, tracker)
+        self.fn = fn
+
+    def call_function(self, /, *args, **kwargs):
+        return self.fn(*args, **kwargs)
+
+    def from_value(value: Any, graph: FunctionGraph, tracker: Tracker):
+        raise NotImplementedError
+
+
 class ClassVariable(CallableVariable):
     def __init__(self, class_: type, graph: FunctionGraph, tracker: Tracker):
         super().__init__(graph, tracker)
diff --git a/test/sot/test_25_with.py b/test/sot/test_25_with.py
@@ -0,0 +1,209 @@
+# Copyright (c) 2025 paddlepaddle Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from contextlib import contextmanager
+
+from test_case_base import (
+    TestCaseBase,
+)
+
+import paddle
+from paddle import nn
+from paddle.jit.sot.opcode_translator.executor.opcode_executor import (
+    ALREADY_SUPPORTED_EXCEPTION,
+)
+from paddle.jit.sot.psdb import check_no_breakgraph
+from paddle.jit.sot.utils import strict_mode_guard
+
+
+class Manager:
+    def __init__(self):
+        pass
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc, value, traceback):
+        pass
+
+
+class ManagerExitReturnFalse(Manager):
+    def __exit__(self, *args):
+        return False
+
+
+class ManagerExitReturnTrue(Manager):
+    def __exit__(self, *args):
+        return True
+
+
+TEST_WITH_STATEMENT_FLAG = False
+
+
+@contextmanager
+def my_context():
+    global TEST_WITH_STATEMENT_FLAG
+    try:
+        TEST_WITH_STATEMENT_FLAG = True
+        yield
+    finally:
+        TEST_WITH_STATEMENT_FLAG = False
+
+
+@check_no_breakgraph
+def with_manager_normal(x):
+    with Manager() as mgr:
+        x *= 2
+    return x
+
+
+@check_no_breakgraph
+def with_manager_exit_true_raise_error(x):
+    with ManagerExitReturnTrue() as mgr_true:
+        x *= 3
+        raise ValueError("test")
+        x -= 4
+    return x
+
+
+@check_no_breakgraph
+def with_manager_exit_true_zero_division(x):
+    with ManagerExitReturnTrue() as mgr_true:
+        x += 5
+        # TODO(DrRyanHuang): Division by zero (x / 0) will raise an InnerError.
+        # In the future, the actual Exception should be propagated rather than being wrapped as InnerError.
+        1 / 0  # noqa: B018
+        x *= 6
+    return x
+
+
+@check_no_breakgraph
+def with_contextmanager_flag_behavior(x):
+    global TEST_WITH_STATEMENT_FLAG
+    with my_context():
+        if TEST_WITH_STATEMENT_FLAG:
+            x /= 7
+        else:
+            x *= 7
+
+    if not TEST_WITH_STATEMENT_FLAG:
+        x += 8
+    return x
+
+
+@check_no_breakgraph
+def with_manager_exit_false(x):
+    try:
+        with ManagerExitReturnFalse() as mgr_false:
+            x *= 4
+            1 / 0  # noqa: B018
+    except ZeroDivisionError:
+        x /= 4
+    return x
+
+
+# TODO(DrRyanHuang): NoGradContextManagerVariable and UserDefinedContextManagerVariable will be implemented separately in the future.
+# The @strict_mode_guard decorator will be removed here to ensure that fallback is no longer permitted.
+@strict_mode_guard(False)
+def test_no_grad_behavior():
+    x = paddle.rand([1, 2])
+    p = paddle.rand([1, 2])
+    p.stop_gradient = False
+    x.stop_gradient = True
+    with paddle.no_grad():
+        y = (x * p).sum()
+    y.backward()
+    return x.grad, p.grad
+
+
+class TestWithStatement(TestCaseBase):
+    def test_manager_normal(self):
+        t = paddle.to_tensor(-10.0)
+        self.assert_results(with_manager_normal, t)
+
+    def test_manager_exit_true_suppresses(self):
+        t = paddle.to_tensor(-10.0)
+        self.assert_results(with_manager_exit_true_raise_error, t)
+
+    def test_manager_exit_true_zero_division(self):
+        t = paddle.to_tensor(-10.0)
+        self.assert_results(with_manager_exit_true_zero_division, t)
+
+    def test_my_context_flag_behavior(self):
+        t = paddle.to_tensor(-10.0)
+        self.assert_results(with_contextmanager_flag_behavior, t)
+
+    def test_with_manager_exit_false(self):
+        t = paddle.to_tensor(-10.0)
+        self.assert_results(with_manager_exit_false, t)
+
+    def test_no_grad(self):
+        self.assert_results(test_no_grad_behavior)
+
+
+class SimpleNet(nn.Layer):
+    def __init__(self, input_dim, output_dim):
+        super().__init__()
+        self.layer = nn.Linear(input_dim, output_dim)
+
+    def forward(self, x):
+        with paddle.static.amp.fp16_guard():
+            return self._forward(x)
+
+    def _forward(self, x):
+        return self.layer(x)
+
+
+@check_no_breakgraph
+def net_call(x: paddle.Tensor, net: nn.Layer):
+    return net(x)
+
+
+def inner_no_grad_fn(x, y):
+    with paddle.no_grad():
+        return x * y + x**2
+
+
+@check_no_breakgraph
+def no_grad_fn_caller(x, y):
+    z = inner_no_grad_fn(x * y, y)
+    a = x * y + x**3 - 1
+    return z + a
+
+
+class TestPaddleContextManager(TestCaseBase):
+    # TODO(DrRyanHuang): Python 3.11 introduced a new opcode, BEFORE_WITH, which is not supported yet.
+    # Therefore, for versions 3.11 and above, fallback is allowed for now.
+    @strict_mode_guard(ALREADY_SUPPORTED_EXCEPTION)
+    def test_fp16_guard(self):
+        x = paddle.randn([4, 4])
+        model = SimpleNet(4, 8)
+        self.assert_results(net_call, x, model)
+
+    def test_no_grad(self):
+        x = paddle.randn([12])
+        y = paddle.randn([12])
+        x.stop_gradient = False
+        y.stop_gradient = False
+        self.assert_results_with_grad(
+            [x, y],
+            no_grad_fn_caller,
+            x,
+            y,
+        )
+
+
+if __name__ == '__main__':
+    unittest.main()
