Update to use Python 3.9 syntax (pytorch#4909)

Summary:
Pull Request resolved: pytorch#4909

X-link: facebookresearch/FBGEMM#1933

Pull Request resolved: pytorch#4905

Reviewed By: cthi

Differential Revision: D82979037

Pulled By: q10

Author: cyyever

fbgemm_gpu/bench/batched_unary_embeddings_benchmark.py

@@ -8,7 +8,6 @@
 
 import functools
 from math import sqrt
-from typing import List, Tuple
 
 import click
 import fbgemm_gpu
@@ -33,7 +32,7 @@ def generate_unary_feature(
     num_embeddings: int,
     # pyre-fixme[24]: Generic type `list` expects 1 type parameter, use
     #  `typing.List[<element type>]` to avoid runtime subscripting errors.
-) -> Tuple[List, List, List]:
+) -> tuple[list, list, list]:
     lengths = []
     offsets = []
     indices = []
@@ -53,7 +52,7 @@ def generate_unary_feature(
 
 
 class MyModule(torch.nn.Module):
-    def __init__(self, num_tasks: int, hash_sizes: List[int]) -> None:
+    def __init__(self, num_tasks: int, hash_sizes: list[int]) -> None:
         super().__init__()
         self.num_tasks = num_tasks
         self.hash_sizes = hash_sizes
@@ -73,7 +72,7 @@ def __init__(self, num_tasks: int, hash_sizes: List[int]) -> None:
                 self.emb_modules.append(emb)
 
     def forward(
-        self, offsets: List[torch.Tensor], indices: List[torch.Tensor]
+        self, offsets: list[torch.Tensor], indices: list[torch.Tensor]
     ) -> torch.Tensor:
         tt_list = []
         for n in range(self.num_tasks):

fbgemm_gpu/bench/bench_utils.py

@@ -10,7 +10,6 @@
 import logging
 import threading
 import time
-from typing import List, Tuple
 
 import torch
 
@@ -32,7 +31,7 @@ def benchmark_torch_function(  # noqa: C901
     name: str = "",
     num_threads: int = 1,
     copy_f_for_multi_thread_test: bool = False,
-) -> Tuple[float, torch.Tensor]:
+) -> tuple[float, torch.Tensor]:
     logging.debug(f"Start to benchmark {name}...")
     if device != "cpu" and device != "" and device != "cuda":
         torch.cuda.set_device(device)
@@ -68,7 +67,7 @@ def benchmark_torch_function(  # noqa: C901
             dtype=torch.float,
             device=device,
         )
-        duration_ms_list: List[float] = []
+        duration_ms_list: list[float] = []
 
         f_list = [f]
         # make deepcopy of f if necessary

fbgemm_gpu/bench/histogram_binning_calibration_benchmark.py

@@ -8,7 +8,7 @@
 
 import logging
 import time
-from typing import Callable, Tuple
+from typing import Callable
 
 import click
 import torch
@@ -25,9 +25,9 @@
 
 
 def benchmark_hbc_function(
-    func: Callable[[Tensor], Tuple[Tensor, Tensor]],
+    func: Callable[[Tensor], tuple[Tensor, Tensor]],
     input: Tensor,
-) -> Tuple[float, Tensor]:
+) -> tuple[float, Tensor]:
     if input.is_cuda:
         torch.cuda.synchronize()
         start_event = torch.cuda.Event(enable_timing=True)
@@ -118,7 +118,7 @@ def cli(
         [num_bins * (num_segments + 1)], dtype=torch.float64
     ).fill_(0.0)
 
-    def fbgemm_hbc_cpu(input: Tensor) -> Tuple[Tensor, Tensor]:
+    def fbgemm_hbc_cpu(input: Tensor) -> tuple[Tensor, Tensor]:
         return torch.ops.fbgemm.histogram_binning_calibration(
             input,
             bin_num_examples,
@@ -130,7 +130,7 @@ def fbgemm_hbc_cpu(input: Tensor) -> Tuple[Tensor, Tensor]:
             0.9995,
         )
 
-    def fbgemm_hbc_by_feature_cpu(input: Tensor) -> Tuple[Tensor, Tensor]:
+    def fbgemm_hbc_by_feature_cpu(input: Tensor) -> tuple[Tensor, Tensor]:
         return torch.ops.fbgemm.histogram_binning_calibration_by_feature(
             input,
             segment_values,
@@ -146,7 +146,7 @@ def fbgemm_hbc_by_feature_cpu(input: Tensor) -> Tuple[Tensor, Tensor]:
             0.9995,
         )
 
-    def fbgemm_generic_hbc_by_feature_cpu(input: Tensor) -> Tuple[Tensor, Tensor]:
+    def fbgemm_generic_hbc_by_feature_cpu(input: Tensor) -> tuple[Tensor, Tensor]:
         return torch.ops.fbgemm.generic_histogram_binning_calibration_by_feature(
             input,
             segment_values,
@@ -186,7 +186,7 @@ def fbgemm_generic_hbc_by_feature_cpu(input: Tensor) -> Tuple[Tensor, Tensor]:
             bin_num_examples_gpu: Tensor = bin_num_examples.cuda()
             bin_num_positives_gpu: Tensor = bin_num_positives.cuda()
 
-            def fbgemm_hbc_gpu(input: Tensor) -> Tuple[Tensor, Tensor]:
+            def fbgemm_hbc_gpu(input: Tensor) -> tuple[Tensor, Tensor]:
                 return torch.ops.fbgemm.histogram_binning_calibration(
                     input,
                     bin_num_examples_gpu,
@@ -206,7 +206,7 @@ def fbgemm_hbc_gpu(input: Tensor) -> Tuple[Tensor, Tensor]:
                 by_feature_bin_num_positives.cuda()
             )
 
-            def fbgemm_hbc_by_feature_gpu(input: Tensor) -> Tuple[Tensor, Tensor]:
+            def fbgemm_hbc_by_feature_gpu(input: Tensor) -> tuple[Tensor, Tensor]:
                 return torch.ops.fbgemm.histogram_binning_calibration_by_feature(
                     input,
                     segment_values_gpu,
@@ -226,7 +226,7 @@ def fbgemm_hbc_by_feature_gpu(input: Tensor) -> Tuple[Tensor, Tensor]:
 
             def fbgemm_generic_hbc_by_feature_gpu(
                 input: Tensor,
-            ) -> Tuple[Tensor, Tensor]:
+            ) -> tuple[Tensor, Tensor]:
                 return (
                     torch.ops.fbgemm.generic_histogram_binning_calibration_by_feature(
                         input,

fbgemm_gpu/bench/jagged_tensor_benchmark.py

@@ -12,7 +12,6 @@
 import logging
 import random
 from dataclasses import dataclass
-from typing import List, Tuple
 
 import click
 import fbgemm_gpu
@@ -505,7 +504,7 @@ def masked_select_jagged_1d(
 
     def ref(
         values: torch.Tensor, lengths: torch.Tensor, mask: torch.Tensor
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         masked_values_ref = values[mask]
         cum_count = torch.cumsum(mask, 0)
         cum_count = torch.cat((cum_count, torch.tensor([0])))
@@ -653,9 +652,9 @@ def keyed_jagged_index_select_dim1(
         ref_inputs.append((key_values, key_lengths, indices, key_weights))
 
     def keyed_jagged_index_select_dim1_ref(
-        inputs: List[torch.Tensor],
+        inputs: list[torch.Tensor],
         has_weights: bool,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         outputs = []
         output_weights = []
         for key_values, key_lengths, indices, _ in inputs:
@@ -758,11 +757,11 @@ def jagged_slice_ref(
         offsets: torch.Tensor,
         start: torch.Tensor,
         max_L: int,
-    ) -> Tuple[torch.Tensor, torch.Tensor]:
+    ) -> tuple[torch.Tensor, torch.Tensor]:
         end_offsets_ = max_L + start + offsets[:-1]
         end_offsets = torch.where(end_offsets_ > offsets[1:], offsets[1:], end_offsets_)
         start_offsets = start + offsets[:-1]
-        indices_to_select: List[torch.Tensor] = []
+        indices_to_select: list[torch.Tensor] = []
         for i in range(end_offsets.size(0)):
             indices_to_select.append(
                 torch.arange(start_offsets[i].item(), end_offsets[i].item())

fbgemm_gpu/bench/merge_embeddings_benchmark.py

@@ -10,7 +10,6 @@
 
 import logging
 import signal
-from typing import List, Tuple
 
 import click
 import fbgemm_gpu
@@ -59,7 +58,7 @@ def get_table_batched_offsets_from_dense(
     merged_indices: torch.Tensor,
     # pyre-fixme[2]: Parameter must be annotated.
     gpu_num,
-) -> Tuple[torch.Tensor, torch.Tensor]:
+) -> tuple[torch.Tensor, torch.Tensor]:
     (T, B, L) = merged_indices.size()
     lengths = np.ones((T, B)) * L
     flat_lengths = lengths.flatten()
@@ -80,7 +79,7 @@ def generate_requests(
     E: int,
     # inter-batch indices reuse rate
     reuse: float = 0.0,
-) -> List[Tuple[torch.IntTensor, torch.IntTensor, None]]:
+) -> list[tuple[torch.IntTensor, torch.IntTensor, None]]:
     rs = []
     for gpu_num in range(num_gpus):
         all_indices = torch.randint(

fbgemm_gpu/bench/quantize_ops_benchmark.py

@@ -9,8 +9,9 @@
 import functools
 import logging
 import random
+from collections.abc import Iterable
 from contextlib import nullcontext
-from typing import Iterable, Optional, Union
+from typing import Optional, Union
 
 import click
 import fbgemm_gpu

fbgemm_gpu/bench/sparse_ops_benchmark.py

@@ -11,7 +11,6 @@
 import logging
 import math
 import random
-from typing import List
 
 import click
 import fbgemm_gpu
@@ -681,42 +680,42 @@ def index_select_bench(
     optim_group: torch.optim.Optimizer = torch.optim.SGD(gis_inputs, lr=0.1)
 
     def index_select_fwd_ref(
-        inputs: List[torch.Tensor], indices: List[torch.Tensor]
-    ) -> List[torch.Tensor]:
+        inputs: list[torch.Tensor], indices: list[torch.Tensor]
+    ) -> list[torch.Tensor]:
         outputs = []
         for input, index in zip(inputs, indices):
             optim_index.zero_grad()
             outputs.append(torch.index_select(input, 0, index))
         return outputs
 
     def index_select_bwd_ref(
-        outputs: List[torch.Tensor], grads: List[torch.Tensor]
+        outputs: list[torch.Tensor], grads: list[torch.Tensor]
     ) -> None:
         for output, grad in zip(outputs, grads):
             optim_index.zero_grad()
             output.backward(grad, retain_graph=True)
 
     def batch_index_select_fwd(
-        concat_inputs: List[torch.Tensor],
-        concat_indices: List[int],
-        input_num_indices: List[int],
-        input_rows: List[int],
-        input_columns: List[int],
+        concat_inputs: list[torch.Tensor],
+        concat_indices: list[int],
+        input_num_indices: list[int],
+        input_rows: list[int],
+        input_columns: list[int],
     ) -> torch.autograd.Variable:
         optim_batch.zero_grad()
         return torch.ops.fbgemm.batch_index_select_dim0(
             concat_inputs, concat_indices, input_num_indices, input_rows, input_columns
         )
 
     def group_index_select_fwd(
-        gis_inputs: List[torch.Tensor], indices: List[int]
+        gis_inputs: list[torch.Tensor], indices: list[int]
     ) -> torch.autograd.Variable:
         optim_group.zero_grad()
         return torch.ops.fbgemm.group_index_select_dim0(gis_inputs, indices)
 
     def batch_group_index_select_bwd(
         output: torch.autograd.Variable,
-        grads: List[torch.Tensor],
+        grads: list[torch.Tensor],
         optim: torch.optim.Optimizer,
     ) -> torch.autograd.Variable:
         optim.zero_grad()

fbgemm_gpu/bench/tbe/split_table_batched_embeddings_benchmark.py

@@ -12,7 +12,7 @@
 import os
 import tempfile
 from contextlib import nullcontext
-from typing import Any, Callable, Dict, Optional
+from typing import Any, Callable, Optional
 
 import click
 import numpy as np
@@ -231,7 +231,7 @@ def device(  # noqa C901
         pooling_mode = PoolingMode.NONE
         do_pooling = False
 
-    common_split_args: Dict[str, Any] = {
+    common_split_args: dict[str, Any] = {
         "weights_precision": weights_precision,
         "stochastic_rounding": stoc,
         "output_dtype": output_dtype,
@@ -1384,7 +1384,7 @@ def vbe(
         else EmbeddingLocation.HOST
     )
 
-    common_split_args: Dict[str, Any] = {
+    common_split_args: dict[str, Any] = {
         "weights_precision": embconfig.weights_dtype,
         "stochastic_rounding": embconfig.stochastic_rounding,
         "output_dtype": embconfig.output_dtype,

fbgemm_gpu/bench/tbe/tbe_cache_benchmark.py

@@ -8,7 +8,6 @@
 
 import logging
 import random
-from typing import List, Tuple
 
 import click
 import numpy as np
@@ -104,7 +103,7 @@ def create_embedding_specs(
     cached_tables_ratio: float,
     num_embeddings: int,
     embedding_dims: int,
-) -> List[Tuple[str, int, int, SparseType, EmbeddingLocation]]:
+) -> list[tuple[str, int, int, SparseType, EmbeddingLocation]]:
     """
     Returns embedding specs to be used with IntNBitTableBatchedEmbeddingBagsCodegen.
     """
@@ -157,7 +156,7 @@ def create_embedding_specs(
 
 def create_request(
     num_tables: int, num_embeddings: int, batch: int, avg_pooling_factor: int
-) -> Tuple[Tensor, Tensor]:
+) -> tuple[Tensor, Tensor]:
     """
     Returns [indices, offsets], which are inputs of embedding bags.
     """

fbgemm_gpu/bench/tbe/tbe_inference_benchmark.py

@@ -16,7 +16,7 @@
 import statistics
 from contextlib import nullcontext
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Optional
 
 import click
 import numpy as np
@@ -1398,14 +1398,14 @@ def nbit_uvm_compare_direct_mapped(
     )
 
     if mixed:
-        Ds: List[int] = [
+        Ds: list[int] = [
             round_up(np.random.randint(low=int(0.5 * D), high=int(1.5 * D)), 4)
             for _ in range(T)
         ]
         # pyre-fixme[9]: D has type `int`; used as `floating[typing.Any]`.
         D = np.average(Ds)
     else:
-        Ds: List[int] = [D] * T
+        Ds: list[int] = [D] * T
 
     _requests_uvm = generate_requests(
         iters,
@@ -1417,7 +1417,7 @@ def nbit_uvm_compare_direct_mapped(
         alpha=alpha,
         weighted=weighted,
     )
-    requests_uvm: List[TBERequest] = [
+    requests_uvm: list[TBERequest] = [
         TBERequest(req.indices.int(), req.offsets.int(), req.per_sample_weights)
         for req in _requests_uvm
     ]
@@ -1429,7 +1429,7 @@ def nbit_uvm_compare_direct_mapped(
         + param_size_multiplier * B * sum(Ds[:T]) * L
     )
 
-    stats: Dict[str, Any] = {
+    stats: dict[str, Any] = {
         "B": B,
         "T": T,
         "E": E,