Migrate from Tuple -> tuple in torch/_inductor (pytorch#144264)

Pull Request resolved: pytorch#144264
Approved by: https://github.com/eellison

Author: bobrenjc93

torch/_inductor/__init__.py

@@ -130,7 +130,7 @@ def aoti_compile_and_package(
 def _aoti_compile_and_package_inner(
     gm: torch.nn.Module,
     # flat_example_inputs: List[Any],
-    args: Tuple[Any],
+    args: tuple[Any],
     kwargs: Optional[Dict[str, Any]] = None,
     *,
     load_and_run: bool = False,
@@ -198,7 +198,7 @@ def aoti_load_package(path: Union[str, io.BytesIO]) -> Any:  # type: ignore[type
 
 def aot_compile(
     gm: torch.fx.GraphModule,
-    args: Tuple[Any],
+    args: tuple[Any],
     kwargs: Optional[Dict[str, Any]] = None,
     *,
     options: Optional[Dict[str, Any]] = None,

torch/_inductor/aoti_eager.py

@@ -2,7 +2,7 @@
 import logging
 import os
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional
 from unittest import mock
 
 import torch
@@ -85,7 +85,7 @@ def supported_builtin_dtype_torch_dtype() -> Dict[type, torch.dtype]:
     return {int: torch.int32, float: torch.float, bool: torch.bool}
 
 
-def supported_scalar_types() -> Tuple[type, ...]:
+def supported_scalar_types() -> tuple[type, ...]:
     type_to_torch_dtype = supported_builtin_dtype_torch_dtype()
     return tuple(type_to_torch_dtype.keys())
 
@@ -170,7 +170,7 @@ def aoti_compile_with_persistent_cache(
     device_type: str,
     dynamic: bool,
     f: Callable[..., Any],
-    args: Tuple[Any],
+    args: tuple[Any],
     kwargs: Dict[str, Any],
     *,
     dynamic_shapes: Optional[Dict[str, Any]] = None,

torch/_inductor/autoheuristic/artifacts/_MMRankingA100.py

@@ -238,7 +238,7 @@ def fill_choices(self) -> None:
     def get_name(self) -> str:
         return 'mm'
 
-    def get_best_choices(self, context: AHContext) -> Optional[List[Tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> Optional[List[tuple[float, int]]]:
         if context.get_value('arith_intensity') <= 52.6245059967041:
             if context.get_value('n') <= 34.0:
                 if context.get_value('n') <= 18.0:

torch/_inductor/autoheuristic/artifacts/_MMRankingH100.py

@@ -242,7 +242,7 @@ def fill_choices(self) -> None:
     def get_name(self) -> str:
         return 'mm'
 
-    def get_best_choices(self, context: AHContext) -> Optional[List[Tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> Optional[List[tuple[float, int]]]:
         if context.get_value('arith_intensity') <= 29.89772129058838:
             if context.get_value('n') <= 34.0:
                 if context.get_value('n') <= 18.0:

torch/_inductor/autoheuristic/artifacts/_MixedMMA100.py

@@ -62,7 +62,7 @@ def fill_choices(self) -> None:
     def get_name(self) -> str:
         return 'mixed_mm'
 
-    def get_best_choices(self, context: AHContext) -> Optional[List[Tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> Optional[List[tuple[float, int]]]:
         if str(context.get_value('1LEQmLEQ16')) != 'True':
             if context.get_value('m') <= 32.5:
                 if context.get_value('n') <= 6976.0:

torch/_inductor/autoheuristic/artifacts/_MixedMMH100.py

@@ -61,7 +61,7 @@ def fill_choices(self) -> None:
     def get_name(self) -> str:
         return 'mixed_mm'
 
-    def get_best_choices(self, context: AHContext) -> Optional[List[Tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> Optional[List[tuple[float, int]]]:
         if context.get_value('arith_intensity') <= 15.988086223602295:
             if context.get_value('n') <= 25280.0:
                 if context.get_value('n') <= 1344.0:

torch/_inductor/autoheuristic/autoheuristic_utils.py

@@ -1,5 +1,5 @@
 import functools
-from typing import Any, Callable, Dict, List, Tuple
+from typing import Any, Callable, Dict, List
 
 import torch
 
@@ -64,7 +64,7 @@ def add_feature(
         self.features.append(AHFeature(name, value, is_categorical=is_categorical))
         self.context_dict[name] = value
 
-    def get_numerical_and_categorical_features(self) -> Tuple[List[str], List[str]]:
+    def get_numerical_and_categorical_features(self) -> tuple[List[str], List[str]]:
         numerical_features = []
         categorical_features = []
         for feature in self.features:
@@ -93,7 +93,7 @@ class AHMetadata:
     def __init__(
         self,
         shared_memory: Any,
-        device_capa: Tuple[int, int],
+        device_capa: tuple[int, int],
         choices: List[Choice],
         name: str,
     ) -> None:
@@ -327,7 +327,7 @@ def mat2_is_contig_fn(data: Any) -> bool:
     return [mat1_is_contig_op, mat2_is_contig_op]
 
 
-def context_add_strides(context: AHContext, name: str, stride: Tuple[int, ...]) -> None:
+def context_add_strides(context: AHContext, name: str, stride: tuple[int, ...]) -> None:
     for i, s in enumerate(stride):
         context.add_feature(f"{name}_stride_{i}", s)
 

torch/_inductor/autoheuristic/learnedheuristic_interface.py

@@ -1,4 +1,4 @@
-from typing import List, Optional, Tuple
+from typing import List, Optional
 
 from torch._inductor.autoheuristic.autoheuristic_utils import (
     AHContext,
@@ -88,5 +88,5 @@ def get_decisions_ranked(self, context: AHContext) -> Optional[List[str]]:
         choices = [choice for choice in choices if choice is not None]
         return choices
 
-    def get_best_choices(self, context: AHContext) -> Optional[List[Tuple[float, int]]]:
+    def get_best_choices(self, context: AHContext) -> Optional[List[tuple[float, int]]]:
         return []

torch/_inductor/codecache.py

@@ -42,7 +42,6 @@
     NoReturn,
     Optional,
     Sequence,
-    Tuple,
     TYPE_CHECKING,
     TypeVar,
     Union,
@@ -375,7 +374,7 @@ def code_hash(code: Union[str, bytes], extra: str = "") -> str:
 
 def get_path(
     basename: str, extension: str, specified_dir: str = ""
-) -> Tuple[str, str, str]:
+) -> tuple[str, str, str]:
     if specified_dir:
         if os.path.isabs(specified_dir):
             subdir = specified_dir
@@ -403,7 +402,7 @@ def write(
     extra: str = "",
     hash_type: str = "code",
     specified_dir: str = "",
-) -> Tuple[str, str]:
+) -> tuple[str, str]:
     # use striped content to compute hash so we don't end up with different
     # hashes just because the content begins/ends with different number of
     # spaces.
@@ -527,7 +526,7 @@ def __init__(
 
     def _reduce_fake_tensor(
         self, t: Tensor
-    ) -> Tuple[Callable[[T], T], Tuple[TensorMetadata]]:
+    ) -> tuple[Callable[[T], T], tuple[TensorMetadata]]:
         """
         Custom reducer to pickle FakeTensors.
         """
@@ -537,7 +536,7 @@ def _reduce_fake_tensor(
     def _reduce_tensor(
         self,
         t: Tensor,
-    ) -> Tuple[Callable[[T], T], Tuple[Union[TensorMetadata, TensorMetadataAndValues]]]:
+    ) -> tuple[Callable[[T], T], tuple[Union[TensorMetadata, TensorMetadataAndValues]]]:
         """
         Custom reducer to pickle Tensors.  If we see tensors, we know they're constants
         stored as attributes on the GraphModule.
@@ -570,7 +569,7 @@ def _reduce_tensor(
         # Otherwise, we just include the metadata.
         return (_ident, (metadata,))
 
-    def _reduce_symint(self, s: SymInt) -> Tuple[Callable[[T], T], Tuple[str]]:
+    def _reduce_symint(self, s: SymInt) -> tuple[Callable[[T], T], tuple[str]]:
         """
         Custom reducer to pickle SymInts.
         """
@@ -588,7 +587,7 @@ def _reduce_unsupported(self, s: Any) -> NoReturn:
 
     def _reduce_graph_module(
         self, gm: torch.fx.GraphModule
-    ) -> Tuple[Any, Tuple[Dict[str, Any], str]]:
+    ) -> tuple[Any, tuple[Dict[str, Any], str]]:
         """
         Custom reducer for graph module to handle irrelevant data for user
         defined triton kernels
@@ -860,7 +859,7 @@ def compiled_fx_graph_hash(
     example_inputs: Sequence[InputType],
     fx_kwargs: _CompileFxKwargs,
     inputs_to_check: Sequence[int],
-) -> Tuple[str, List[str]]:
+) -> tuple[str, List[str]]:
     """
     Generate a unique hash of the FX graph for caching.
     """
@@ -980,7 +979,7 @@ def _lookup_graph(
         local: bool,
         remote_cache: Optional[RemoteCache[JsonDataTy]],
         constants: CompiledFxGraphConstants,
-    ) -> Tuple[Optional[CompiledFxGraph], Dict[str, Any]]:
+    ) -> tuple[Optional[CompiledFxGraph], Dict[str, Any]]:
         """
         Lookup a compiled graph in the cache by key. On a hit, return the
         deserialized CompiledFxGraph object. On a miss, return None.
@@ -1222,7 +1221,7 @@ def prepare_key(
         fx_kwargs: _CompileFxKwargs,
         inputs_to_check: Sequence[int],
         remote: bool,
-    ) -> Tuple[Optional[Tuple[str, List[str]]], Dict[str, Any]]:
+    ) -> tuple[Optional[tuple[str, List[str]]], Dict[str, Any]]:
         """
         Checks that the inductor input is cacheable, then computes
         and returns the cache key for the input.
@@ -1274,7 +1273,7 @@ def load_with_key(
         remote_cache: Optional[RemoteCache[JsonDataTy]],
         is_backward: bool,
         constants: CompiledFxGraphConstants,
-    ) -> Tuple[Optional[CompiledFxGraph], Dict[str, Any]]:
+    ) -> tuple[Optional[CompiledFxGraph], Dict[str, Any]]:
         """
         Lookup the graph with the given key, and return results and metadata.
         Doesn't do any logging on its own, because AOTAutograd handles a cache miss
@@ -1347,7 +1346,7 @@ def run_command_and_check(cmd_: str) -> None:
 
 
 @functools.lru_cache(None)
-def split_aot_inductor_output_path(path: str) -> Tuple[str, str]:
+def split_aot_inductor_output_path(path: str) -> tuple[str, str]:
     """Returns the path where the AOT Inductor compiled kernels are stored."""
     if path.endswith(".so"):
         return os.path.split(path)
@@ -2720,18 +2719,18 @@ class PyCodeCache:
     # than once, but attach different attributes, i.e., due to different
     # constant values.
     modules: List[ModuleType] = []
-    linemaps: Dict[str, List[Tuple[Any, ...]]] = {}
+    linemaps: Dict[str, List[tuple[Any, ...]]] = {}
 
     @classmethod
-    def write(cls, source_code: str, extra: str = "") -> Tuple[str, str]:
+    def write(cls, source_code: str, extra: str = "") -> tuple[str, str]:
         return write(source_code, "py", extra=extra)
 
     @classmethod
     def load(
         cls,
         source_code: str,
         extra: str = "",
-        linemap: Optional[List[Tuple[int, str]]] = None,
+        linemap: Optional[List[tuple[int, str]]] = None,
         attrs: Optional[Dict[str, Any]] = None,
     ) -> ModuleType:
         key, path = write(source_code, "py", extra=extra)
@@ -2742,7 +2741,7 @@ def load_by_key_path(
         cls,
         key: str,
         path: str,
-        linemap: Optional[List[Tuple[int, str]]] = None,
+        linemap: Optional[List[tuple[int, str]]] = None,
         attrs: Optional[Dict[str, Any]] = None,
     ) -> ModuleType:
         if linemap is None:
@@ -3043,7 +3042,7 @@ class CacheEntry:
     _SOURCE_CODE_SUFFIX = "cu"
 
     @classmethod
-    def write(cls, source_code: str, dst_file_ext: str) -> Tuple[str, str]:
+    def write(cls, source_code: str, dst_file_ext: str) -> tuple[str, str]:
         """
         Writes source code into a file with dst_file_ext as the file extension.
         Returns the hash key of source code, and the path to the file.
@@ -3060,7 +3059,7 @@ def write(cls, source_code: str, dst_file_ext: str) -> Tuple[str, str]:
     @classmethod
     def compile(
         cls, source_code: str, dst_file_ext: str, extra_args: Optional[List[str]] = None
-    ) -> Tuple[str, str, str]:
+    ) -> tuple[str, str, str]:
         """
         Compiles CUDA source_code into a file with dst_file_ext extension.
         Returns a tuple of dst_file_path, hash_key, source_code_path
@@ -3099,7 +3098,7 @@ def compile(
         return (cls.cache[key].output_path, key, input_path)
 
     @classmethod
-    def load(cls, source_code: str, dst_file_ext: str) -> Tuple[DLLWrapper, str, str]:
+    def load(cls, source_code: str, dst_file_ext: str) -> tuple[DLLWrapper, str, str]:
         """
         Compiles source code and loads the generated .so file.
         Returns a tuple of DLLWrapper, hash_key, source_code_path
@@ -3129,7 +3128,7 @@ class CacheEntry:
     _logged_compiler_version = False
 
     @classmethod
-    def write(cls, source_code: str, dst_file_ext: str) -> Tuple[str, str]:
+    def write(cls, source_code: str, dst_file_ext: str) -> tuple[str, str]:
         """
         Writes source code into a file with dst_file_ext as the file extension.
         Returns the hash key of source code, and the path to the file.
@@ -3146,7 +3145,7 @@ def write(cls, source_code: str, dst_file_ext: str) -> Tuple[str, str]:
     @classmethod
     def compile(
         cls, source_code: str, dst_file_ext: str, extra_args: Optional[List[str]] = None
-    ) -> Tuple[str, str, str]:
+    ) -> tuple[str, str, str]:
         """
         Compiles source_code into a file with dst_file_ext extension,
         using the compile command specific for the ROCm platform.
@@ -3194,7 +3193,7 @@ def compile(
         return (cls.cache[key].output_path, key, input_path)
 
     @classmethod
-    def load(cls, source_code: str, dst_file_ext: str) -> Tuple[DLLWrapper, str, str]:
+    def load(cls, source_code: str, dst_file_ext: str) -> tuple[DLLWrapper, str, str]:
         """
         Compiles source code and loads the generated .so file.
         Returns a tuple of DLLWrapper, hash_key, source_code_path

torch/_inductor/codegen/block_analysis.py

@@ -1,7 +1,7 @@
 import collections
 import functools
 import textwrap
-from typing import List, Optional, Tuple
+from typing import List, Optional
 
 import sympy
 from sympy import Expr, Symbol
@@ -49,7 +49,7 @@ def match_mod_div_block_expr(
         index_var: Symbol,
         numel: Expr,
         num_dims: int,
-    ) -> Optional[Tuple[List[Expr], List[Expr], List[Expr]]]:
+    ) -> Optional[tuple[List[Expr], List[Expr], List[Expr]]]:
         """
         Matches modular indexing expressions, converting them to implied block dimensions and strides.
         See triton.py for more information.