diff --git a/tensor-shapes/pyrefly-torch-stubs/test/test_public_attributes.py b/tensor-shapes/pyrefly-torch-stubs/test/test_public_attributes.py index 16d457ce65..ecbea56a29 100644 --- a/tensor-shapes/pyrefly-torch-stubs/test/test_public_attributes.py +++ b/tensor-shapes/pyrefly-torch-stubs/test/test_public_attributes.py @@ -3,11 +3,15 @@ # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. -from typing import assert_type +from typing import assert_type, TYPE_CHECKING import torch from shape_extensions import assert_shape +if TYPE_CHECKING: + + def identity(value: int) -> int: ... + def test_public_attributes() -> None: assert_type(torch.complex64, torch.dtype) @@ -23,3 +27,7 @@ def test_public_attributes() -> None: torch.set_grad_enabled, torch.set_default_dtype, ) + + if TYPE_CHECKING: + assert_type(torch.compile(identity)(1), int) + assert_type(torch.compile()(identity)(1), int) diff --git a/tensor-shapes/pyrefly-torch-stubs/torch-stubs/__init__.pyi b/tensor-shapes/pyrefly-torch-stubs/torch-stubs/__init__.pyi index 32a3485e8f..f4ff40ed4f 100644 --- a/tensor-shapes/pyrefly-torch-stubs/torch-stubs/__init__.pyi +++ b/tensor-shapes/pyrefly-torch-stubs/torch-stubs/__init__.pyi @@ -650,7 +650,6 @@ from torch.autograd import ( enable_grad as enable_grad, set_grad_enabled as set_grad_enabled, ) -from torch.compiler import compile as compile from torch.func import vmap as vmap from torch.functional import ( align_tensors as align_tensors, @@ -4395,6 +4394,14 @@ def cross[B: IntTuple]( ... # Context managers +# Define the top-level API directly because older Torch versions lack torch.compiler. +@overload +def compile[F: Callable[..., Any]](model: F, **kwargs: Any) -> F: ... +@overload +def compile[**P, R]( + model: None = None, **kwargs: Any +) -> Callable[[Callable[P, R]], Callable[P, R]]: ... + class no_grad: """Context manager and decorator that disables gradient tracking.