[Docs] Optimize parameter description to declare allowed type (2/N) (#138152)

zeshengzong · pytorchmergebot · commit 7b863230ea5c · 2024-10-18T11:18:19.000Z
Inspired by issue #137422 and #103847 Optimize method parameter types in docs to given users a more clear about what expected to pass to methods. Previous PR: - [x] #137956 Pull Request resolved: #138152 Approved by: https://github.com/albanD
diff --git a/torch/_torch_docs.py b/torch/_torch_docs.py
@@ -311,7 +311,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.index_copy,
     r"""
-index_copy(input, dim, index, source, *, out=None) -> Tensor
+index_copy(input: Tensor, dim: int, index: Tensor, source: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 See :meth:`~Tensor.index_add_` for function description.
 """,
@@ -320,7 +320,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.index_reduce,
     r"""
-index_reduce(input, dim, index, source, reduce, *, include_self=True, out=None) -> Tensor
+index_reduce(input: Tensor, dim: int, index: Tensor, source: Tensor, reduce: str, *, include_self: bool = True, out: Optional[Tensor]) -> Tensor # noqa: B950
 
 See :meth:`~Tensor.index_reduce_` for function description.
 """,
@@ -578,12 +578,15 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.adjoint,
     r"""
-adjoint(Tensor) -> Tensor
+adjoint(input: Tensor) -> Tensor
 Returns a view of the tensor conjugated and with the last two dimensions transposed.
 
 ``x.adjoint()`` is equivalent to ``x.transpose(-2, -1).conj()`` for complex tensors and
 to ``x.transpose(-2, -1)`` for real tensors.
 
+Args:
+    {input}
+
 Example::
     >>> x = torch.arange(4, dtype=torch.float)
     >>> A = torch.complex(x, x).reshape(2, 2)
@@ -732,7 +735,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.allclose,
     r"""
-allclose(input, other, rtol=1e-05, atol=1e-08, equal_nan=False) -> bool
+allclose(input: Tensor, other: Tensor, rtol: float = 1e-05, atol: float = 1e-08, equal_nan: bool = False) -> bool
 
 This function checks if :attr:`input` and :attr:`other` satisfy the condition:
 
@@ -766,7 +769,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.all,
     r"""
-all(input) -> Tensor
+all(input: Tensor) -> Tensor
 
 Tests if all elements in :attr:`input` evaluate to `True`.
 
@@ -821,7 +824,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.any,
     r"""
-any(input) -> Tensor
+any(input: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 Tests if any element in :attr:`input` evaluates to `True`.
 
@@ -876,7 +879,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.angle,
     r"""
-angle(input, *, out=None) -> Tensor
+angle(input: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 Computes the element-wise angle (in radians) of the given :attr:`input` tensor.
 
@@ -946,7 +949,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.as_tensor,
     r"""
-as_tensor(data, dtype=None, device=None) -> Tensor
+as_tensor(data: Any, dtype: Optional[dtype] = None, device: Optional[DeviceLikeType]) -> Tensor
 
 Converts :attr:`data` into a tensor, sharing data and preserving autograd
 history if possible.
@@ -1025,7 +1028,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.arcsin,
     r"""
-arcsin(input, *, out=None) -> Tensor
+arcsin(input: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 Alias for :func:`torch.asin`.
 """,
@@ -1061,7 +1064,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.arcsinh,
     r"""
-arcsinh(input, *, out=None) -> Tensor
+arcsinh(input: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 Alias for :func:`torch.asinh`.
 """,
@@ -1070,7 +1073,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.atan,
     r"""
-atan(input, *, out=None) -> Tensor
+atan(input: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 Returns a new tensor with the arctangent of the elements of :attr:`input`.
 
@@ -1097,7 +1100,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.arctan,
     r"""
-arctan(input, *, out=None) -> Tensor
+arctan(input: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 Alias for :func:`torch.atan`.
 """,
@@ -1106,7 +1109,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.atan2,
     r"""
-atan2(input, other, *, out=None) -> Tensor
+atan2(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 Element-wise arctangent of :math:`\text{{input}}_{{i}} / \text{{other}}_{{i}}`
 with consideration of the quadrant. Returns a new tensor with the signed angles
@@ -1138,15 +1141,15 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.arctan2,
     r"""
-arctan2(input, other, *, out=None) -> Tensor
+arctan2(input: Tensor, other: Tensor, *, out: Optional[Tensor]) -> Tensor
 Alias for :func:`torch.atan2`.
 """,
 )
 
 add_docstr(
     torch.atanh,
     r"""
-atanh(input, *, out=None) -> Tensor
+atanh(input: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 Returns a new tensor with the inverse hyperbolic tangent of the elements of :attr:`input`.
 
@@ -1178,7 +1181,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.arctanh,
     r"""
-arctanh(input, *, out=None) -> Tensor
+arctanh(input: Tensor, *, out: Optional[Tensor]) -> Tensor
 
 Alias for :func:`torch.atanh`.
 """,
@@ -1187,7 +1190,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.asarray,
     r"""
-asarray(obj, *, dtype=None, device=None, copy=None, requires_grad=False) -> Tensor
+asarray(obj: Any, *, dtype: Optional[dtype], device: Optional[DeviceLikeType], copy: Optional[bool] = None, requires_grad: bool = False) -> Tensor # noqa: B950
 
 Converts :attr:`obj` to a tensor.
 
@@ -1352,7 +1355,7 @@ def merge_dicts(*dicts):
 add_docstr(
     torch.bernoulli,
     r"""
-bernoulli(input, *, generator=None, out=None) -> Tensor
+bernoulli(input: Tensor, *, generator: Optional[Generator], out: Optional[Tensor]) -> Tensor
 
 Draws binary random numbers (0 or 1) from a Bernoulli distribution.
 
