[BE][Ez]: Fully type nn.utils.clip_grad (#154801)

Skylion007 · pytorchmergebot · commit fcc682be4bda · 2025-07-09T14:27:51.000Z
Full types clip_grad and exposed typing annotations that were hidden by a bad decorator Pull Request resolved: #154801 Approved by: https://github.com/jansel
diff --git a/torch/nn/utils/clip_grad.py b/torch/nn/utils/clip_grad.py
@@ -4,8 +4,8 @@
 import types
 import typing
 import warnings
-from typing import cast, Optional, Union
-from typing_extensions import deprecated
+from typing import Callable, cast, Optional, TypeVar, Union
+from typing_extensions import deprecated, ParamSpec, TypeAlias
 
 import torch
 from torch import Tensor
@@ -19,13 +19,16 @@
 __all__: list[str] = []
 
 
-_tensor_or_tensors = Union[
+_tensor_or_tensors: TypeAlias = Union[  # noqa: PYI042
     torch.Tensor,
     typing.Iterable[torch.Tensor],  # noqa: UP006 - needed until XLA's patch is updated
 ]
 
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
 
-def _no_grad(func):
+
+def _no_grad(func: Callable[_P, _R]) -> Callable[_P, _R]:
     """
     This wrapper is needed to avoid a circular import when using @torch.no_grad on the exposed functions
     clip_grad_norm_ and clip_grad_value_ themselves.
