fix type hints for interpolation functions (#157202)

Fixes #129053

Previously interpolate had a bad signature and not correct type hints.
This fixes this issue.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/157202
Approved by: https://github.com/ezyang, https://github.com/albanD

torch/_C/_nn.pyi.in

@@ -67,5 +67,73 @@ def pad_sequence(
     padding_value: float = 0.0,
     padding_side: Literal["left", "right"] = "right",
 ) -> Tensor: ...
+
+# Upsample functions used by torch.nn.functional.interpolate
+def upsample_nearest1d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def upsample_nearest2d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def upsample_nearest3d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def _upsample_nearest_exact1d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def _upsample_nearest_exact2d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def _upsample_nearest_exact3d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def upsample_linear1d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    align_corners: bool,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def _upsample_bilinear2d_aa(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    align_corners: bool,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def upsample_bilinear2d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    align_corners: bool,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def upsample_trilinear3d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    align_corners: bool,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def _upsample_bicubic2d_aa(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    align_corners: bool,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
+def upsample_bicubic2d(
+    input: Tensor,
+    output_size: Sequence[int] | None,
+    align_corners: bool,
+    scale_factors: Sequence[float] | None,
+) -> Tensor: ...
 def flatten_dense_tensors(tensors: list[Tensor]) -> Tensor: ...
 def unflatten_dense_tensors(flat: Tensor, tensors: list[Tensor]) -> list[Tensor]: ...

torch/nn/functional.pyi.in

@@ -396,14 +396,14 @@ def instance_norm(
 __all__ += ["instance_norm"]
 
 def interpolate(
-    input: Any,
-    size: Any | None = ...,
-    scale_factor: Any | None = ...,
+    input: Tensor,
+    size: int | Sequence[int] | None = ...,
+    scale_factor: float | Sequence[float] | None = ...,
     mode: str = ...,
-    align_corners: Any | None = ...,
-    recompute_scale_factor: Any | None = ...,
+    align_corners: bool | None = ...,
+    recompute_scale_factor: bool | None = ...,
     antialias: bool = ...,
-): ...
+) -> Tensor: ...
 
 __all__ += ["interpolate"]