fix(stubs): add torch.compile

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)

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.