Recognize re-exported PyTorch Tensor in efficiency lints

Summary:
PyTorch exposes its public `torch.Tensor` type through an internal re-export in common stub layouts, so the efficiency lints could see a resolved tensor class while still rejecting it because the defining name was not exactly `torch.Tensor`. This made the opt-in lint group look inert in real projects even though the synthetic tests passed.

Treat the canonical resolved implementation target for `torch.Tensor` as a PyTorch tensor while keeping the match narrow enough to avoid assuming every `Tensor` class under the torch package is the same type. The regression tests cover both PyTorch’s internal re-export and a user module that re-exports the public tensor alias, which documents that alias chains should resolve back to the canonical tensor identity.

Fixes #4062

Reviewed By: yangdanny97

Differential Revision: D111287027

fbshipit-source-id: 3a44a800b16da7beb0d2a1d31c5042ca4071d3ff

Author: grievejia

pyrefly/lib/alt/expr.rs

@@ -3099,9 +3099,15 @@ impl<'a, Ans: LookupAnswer> AnswersSolver<'a, Ans> {
     }
 
     fn is_pytorch_tensor_type(ty: &Type) -> bool {
+        fn is_torch_tensor_class(cls: &ClassType) -> bool {
+            let module = cls.class_object().module_name();
+            let module = module.as_str();
+            cls.name().as_str() == "Tensor" && matches!(module, "torch" | "torch._tensor")
+        }
+
         match ty {
-            Type::ClassType(cls) => cls.has_qname("torch", "Tensor"),
-            Type::ShapedArray(shaped_array) => shaped_array.base_class.has_qname("torch", "Tensor"),
+            Type::ClassType(cls) => is_torch_tensor_class(cls),
+            Type::ShapedArray(shaped_array) => is_torch_tensor_class(&shaped_array.base_class),
             _ => false,
         }
     }

pyrefly/lib/test/pytorch_efficiency_lint.rs

@@ -93,6 +93,41 @@ def zeros(*size: int, device: device | None = None) -> Tensor: ...
     e
 }
 
+fn reexported_torch_env_with_lint() -> TestEnv {
+    let mut e = TestEnv::new().enable_pytorch_efficiency_lint_error();
+    e.add(
+        "torch._tensor",
+        r#"
+class Tensor:
+    def item(self) -> int | float: ...
+    def to(self, device: object) -> "Tensor": ...
+    def cuda(self) -> "Tensor": ...
+"#,
+    );
+    e.add(
+        "torch",
+        r#"
+from torch._tensor import Tensor as Tensor
+
+class device: ...
+
+def zeros(*size: int, device: device | None = None) -> Tensor: ...
+"#,
+    );
+    e
+}
+
+fn user_reexported_torch_env_with_lint() -> TestEnv {
+    let mut e = reexported_torch_env_with_lint();
+    e.add(
+        "mytorch",
+        r#"
+from torch import Tensor as Tensor
+"#,
+    );
+    e
+}
+
 testcase!(
     test_tensor_item_call,
     env_with_lint(),
@@ -129,6 +164,33 @@ def f(x: torch.Tensor, device: torch.device) -> None:
 "#,
 );
 
+testcase!(
+    test_reexported_torch_tensor_lints,
+    reexported_torch_env_with_lint(),
+    r#"
+import torch
+
+def f(x: torch.Tensor, device: torch.device) -> None:
+    x.item()  # E: `Tensor.item()` causes implicit GPU-to-CPU synchronization
+    x.cuda()  # E: `Tensor.cuda()` hard-codes the target device
+    print(x)  # E: printing a `Tensor` causes implicit GPU-to-CPU synchronization
+    torch.zeros(2, 3).to(device)  # E: `torch.zeros(...).to(device)` creates the tensor on CPU first, then copies it
+"#,
+);
+
+testcase!(
+    test_user_reexported_torch_tensor_lints,
+    user_reexported_torch_env_with_lint(),
+    r#"
+import mytorch
+
+def f(x: mytorch.Tensor) -> None:
+    x.item()  # E: `Tensor.item()` causes implicit GPU-to-CPU synchronization
+    x.cuda()  # E: `Tensor.cuda()` hard-codes the target device
+    print(x)  # E: printing a `Tensor` causes implicit GPU-to-CPU synchronization
+"#,
+);
+
 testcase!(
     test_non_tensor_item_call_ok,
     env_with_lint(),