From c2547dfcc339f1c788b13ee3e191fb900e22207a Mon Sep 17 00:00:00 2001
From: Aaron Gokaslan <aaronGokaslan@gmail.com>
Date: Sun, 2 Jun 2024 13:38:33 +0000
Subject: [PATCH] [BE][Ez]: Enable ruff PYI019 (#127684)

Tells pytorch to use typing_extensions.Self when it's able to.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/127684
Approved by: https://github.com/ezyang
---
 pyproject.toml                                |  1 -
 torch/nn/utils/rnn.pyi                        | 49 +++++++------------
 .../_internal/diagnostics/infra/context.py    | 12 +++--
 3 files changed, 26 insertions(+), 36 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 8ec04f77042..01fce4a7d6f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -68,7 +68,6 @@ ignore = [
     "PERF401",
     "PERF403",
     # these ignores are from PYI; please fix!
-    "PYI019",
     "PYI024",
     "PYI036",
     "PYI041",
diff --git a/torch/nn/utils/rnn.pyi b/torch/nn/utils/rnn.pyi
index fed87febe2a..fd033d8888b 100644
--- a/torch/nn/utils/rnn.pyi
+++ b/torch/nn/utils/rnn.pyi
@@ -1,14 +1,5 @@
-from typing import (
-    Any,
-    Iterable,
-    NamedTuple,
-    Optional,
-    overload,
-    Sequence,
-    Tuple,
-    TypeVar,
-    Union,
-)
+from typing import Any, Iterable, NamedTuple, Optional, overload, Sequence, Tuple, Union
+
 from typing_extensions import Self
 
 from torch import Tensor
@@ -24,8 +15,6 @@ class PackedSequence_(NamedTuple):
 
 def bind(optional: Any, fn: Any): ...
 
-_T = TypeVar("_T")
-
 class PackedSequence(PackedSequence_):
     def __new__(
         cls,
@@ -34,39 +23,39 @@ class PackedSequence(PackedSequence_):
         sorted_indices: Optional[Tensor] = ...,
         unsorted_indices: Optional[Tensor] = ...,
     ) -> Self: ...
-    def pin_memory(self: _T) -> _T: ...
-    def cuda(self: _T, *args: Any, **kwargs: Any) -> _T: ...
-    def cpu(self: _T) -> _T: ...
-    def double(self: _T) -> _T: ...
-    def float(self: _T) -> _T: ...
-    def half(self: _T) -> _T: ...
-    def long(self: _T) -> _T: ...
-    def int(self: _T) -> _T: ...
-    def short(self: _T) -> _T: ...
-    def char(self: _T) -> _T: ...
-    def byte(self: _T) -> _T: ...
+    def pin_memory(self: Self) -> Self: ...
+    def cuda(self: Self, *args: Any, **kwargs: Any) -> Self: ...
+    def cpu(self: Self) -> Self: ...
+    def double(self: Self) -> Self: ...
+    def float(self: Self) -> Self: ...
+    def half(self: Self) -> Self: ...
+    def long(self: Self) -> Self: ...
+    def int(self: Self) -> Self: ...
+    def short(self: Self) -> Self: ...
+    def char(self: Self) -> Self: ...
+    def byte(self: Self) -> Self: ...
     @overload
     def to(
-        self: _T,
+        self: Self,
         dtype: _dtype,
         non_blocking: bool = False,
         copy: bool = False,
-    ) -> _T: ...
+    ) -> Self: ...
     @overload
     def to(
-        self: _T,
+        self: Self,
         device: Optional[DeviceLikeType] = None,
         dtype: Optional[_dtype] = None,
         non_blocking: bool = False,
         copy: bool = False,
-    ) -> _T: ...
+    ) -> Self: ...
     @overload
     def to(
-        self: _T,
+        self: Self,
         other: Tensor,
         non_blocking: bool = False,
         copy: bool = False,
-    ) -> _T: ...
+    ) -> Self: ...
     @property
     def is_cuda(self) -> bool: ...
     def is_pinned(self) -> bool: ...
diff --git a/torch/onnx/_internal/diagnostics/infra/context.py b/torch/onnx/_internal/diagnostics/infra/context.py
index 22370850df8..6106a42467c 100644
--- a/torch/onnx/_internal/diagnostics/infra/context.py
+++ b/torch/onnx/_internal/diagnostics/infra/context.py
@@ -21,6 +21,8 @@ from typing import (
     TypeVar,
 )
 
+from typing_extensions import Self
+
 from torch.onnx._internal.diagnostics import infra
 from torch.onnx._internal.diagnostics.infra import formatter, sarif, utils
 from torch.onnx._internal.diagnostics.infra.sarif import version as sarif_version
@@ -92,24 +94,24 @@ class Diagnostic:
         )
         return sarif_result
 
-    def with_location(self: _Diagnostic, location: infra.Location) -> _Diagnostic:
+    def with_location(self: Self, location: infra.Location) -> Self:
         """Adds a location to the diagnostic."""
         self.locations.append(location)
         return self
 
     def with_thread_flow_location(
-        self: _Diagnostic, location: infra.ThreadFlowLocation
-    ) -> _Diagnostic:
+        self: Self, location: infra.ThreadFlowLocation
+    ) -> Self:
         """Adds a thread flow location to the diagnostic."""
         self.thread_flow_locations.append(location)
         return self
 
-    def with_stack(self: _Diagnostic, stack: infra.Stack) -> _Diagnostic:
+    def with_stack(self: Self, stack: infra.Stack) -> Self:
         """Adds a stack to the diagnostic."""
         self.stacks.append(stack)
         return self
 
-    def with_graph(self: _Diagnostic, graph: infra.Graph) -> _Diagnostic:
+    def with_graph(self: Self, graph: infra.Graph) -> Self:
         """Adds a graph to the diagnostic."""
         self.graphs.append(graph)
         return self