tlxd0.1

Author: tlxd

python/paddle/autograd/backward_mode.py

@@ -14,18 +14,24 @@
 
 from __future__ import annotations
 
+from typing import TYPE_CHECKING
+
 import paddle
 from paddle.base import core, framework
 from paddle.base.backward import gradients_with_optimizer  # noqa: F401
 
+if TYPE_CHECKING:
+    from .. import Tensor
+
+
 __all__ = []
 
 
 @framework.dygraph_only
 def backward(
-    tensors: list[paddle.Tensor],
-    grad_tensors: list[paddle.Tensor] | None = None,
-    retain_graph: bool | None = False,
+    tensors: list[Tensor],
+    grad_tensors: list[Tensor, None] | None = None,
+    retain_graph: bool = False,
 ) -> None:
     """
     Compute the backward gradients of given tensors.
@@ -38,7 +44,7 @@ def backward(
             If None, all the gradients of the ``tensors`` is the default value which is filled with 1.0.
             Defaults to None.
 
-        retain_graph(bool|False, optional): If False, the graph used to compute grads will be freed. If you would
+        retain_graph(bool, optional): If False, the graph used to compute grads will be freed. If you would
             like to add more ops to the built graph after calling this method( :code:`backward` ), set the parameter
             :code:`retain_graph` to True, then the grads will be retained. Thus, setting it to False is much more memory-efficient.
             Defaults to False.
@@ -87,8 +93,8 @@ def backward(
     """
 
     def check_tensors(
-        in_out_list: list[paddle.Tensor], name: str
-    ) -> list[paddle.Tensor]:
+        in_out_list: list[Tensor] | tuple[Tensor, ...] | Tensor, name: str
+    ) -> list[Tensor] | tuple[Tensor, ...]:
         assert in_out_list is not None, f"{name} should not be None"
 
         if isinstance(in_out_list, (list, tuple)):