[Typing][B-38,B-39] Add type annotations for python/paddle/autograd/{autograd,saved_tensors_hooks}.py (#67179)

enkilee · web-flow · commit 7aa37b94c824 · 2024-08-10T02:45:52.000+08:00
diff --git a/python/paddle/autograd/autograd.py b/python/paddle/autograd/autograd.py
@@ -14,10 +14,14 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
+from typing import TYPE_CHECKING, overload
 
 import paddle
 from paddle.base import framework
 
+if TYPE_CHECKING:
+    from paddle import Tensor
+
 
 def as_tensors(xs):
     if isinstance(xs, framework.Variable):
@@ -64,7 +68,12 @@ class Jacobian:
 
     """
 
-    def __init__(self, ys, xs, is_batched=False):
+    def __init__(
+        self,
+        ys: Tensor,
+        xs: Tensor,
+        is_batched: bool = False,
+    ) -> None:
         if not is_batched:
             if not 0 <= len(xs.shape) <= 1:
                 raise ValueError(
@@ -91,7 +100,7 @@ def __init__(self, ys, xs, is_batched=False):
             self._jacobian = _JacobianBatchFirst(ys, xs)
 
     @property
-    def shape(self):
+    def shape(self) -> list[int]:
         """The shape of flattened Jacobian matrix."""
         return self._jacobian.shape
 
@@ -448,11 +457,43 @@ def _multi_index(indexes, shape):
     return tuple(positive_indexes)
 
 
+@overload
+def jacobian(
+    ys: Tensor,
+    xs: Tensor,
+    batch_axis: int | None = ...,
+) -> Jacobian: ...
+
+
+@overload
+def jacobian(
+    ys: Sequence[Tensor],
+    xs: Sequence[Tensor],
+    batch_axis: int | None = ...,
+) -> tuple[tuple[Jacobian, ...], ...]: ...
+
+
+@overload
+def jacobian(
+    ys: Tensor,
+    xs: Sequence[Tensor],
+    batch_axis: int | None = ...,
+) -> tuple[Jacobian, ...]: ...
+
+
+@overload
 def jacobian(
-    ys: paddle.Tensor | tuple[paddle.Tensor, ...],
-    xs: paddle.Tensor | tuple[paddle.Tensor, ...],
-    batch_axis: int | None = None,
-) -> tuple[tuple[Jacobian, ...], ...] | tuple[Jacobian, ...] | Jacobian:
+    ys: Sequence[Tensor],
+    xs: Tensor,
+    batch_axis: int | None = ...,
+) -> tuple[Jacobian, ...]: ...
+
+
+def jacobian(
+    ys,
+    xs,
+    batch_axis=None,
+):
     r"""
     Computes the Jacobian of the dependent variable ``ys`` versus the independent
     variable ``xs``.
@@ -542,11 +583,27 @@ def jacobian(
     return _jacobian
 
 
+@overload
+def hessian(
+    ys: Tensor,
+    xs: Tensor,
+    batch_axis: int | None = ...,
+) -> Hessian: ...
+
+
+@overload
+def hessian(
+    ys: Tensor,
+    xs: Sequence[Tensor],
+    batch_axis: int | None = ...,
+) -> tuple[tuple[Hessian, ...], ...]: ...
+
+
 def hessian(
-    ys: paddle.Tensor,
-    xs: paddle.Tensor | tuple[paddle.Tensor, ...],
-    batch_axis: int | None = None,
-) -> tuple[tuple[Hessian, ...], ...] | Hessian:
+    ys,
+    xs,
+    batch_axis=None,
+):
     r"""
     Computes the Jacobian of the dependent variable ``ys`` versus the independent
     variable ``xs``.
diff --git a/python/paddle/autograd/saved_tensors_hooks.py b/python/paddle/autograd/saved_tensors_hooks.py
@@ -11,9 +11,16 @@
 # 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.
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 from paddle.base import core
 
+if TYPE_CHECKING:
+    from collections.abc import Callable
+
+    from paddle import Tensor
 __all__ = []
 
 
@@ -103,14 +110,18 @@ class saved_tensors_hooks:
             >>> y.sum().backward()
     """
 
-    def __init__(self, pack_hook, unpack_hook):
+    def __init__(
+        self,
+        pack_hook: Callable[[Tensor], Any | None],
+        unpack_hook: Callable[[Any], Tensor | None],
+    ) -> None:
         self.pack_hook = pack_hook
         self.unpack_hook = unpack_hook
 
-    def __enter__(self):
+    def __enter__(self) -> None:
         core.eager.register_saved_tensors_hooks(
             self.pack_hook, self.unpack_hook
         )
 
-    def __exit__(self, *args):
+    def __exit__(self, *args: object) -> None:
         core.eager.reset_saved_tensors_hooks()
