pytorch · XuehaiPan · Jul 1, 2024 · Jul 1, 2024 · Jul 1, 2024 · Jul 1, 2024
diff --git a/.flake8 b/.flake8
@@ -8,6 +8,8 @@ max-line-length = 120
 # E501 is not flexible enough, we're using B950 instead
 ignore =
     E203,E305,E402,E501,E721,E741,F405,F841,F999,W503,W504,C408,E302,W291,E303,
+    # type stub in .py files formatted by black
+    E704,
     # shebang has extra meaning in fbcode lints, so I think it's not worth trying
     # to line this up with executable bit
     EXE001,

diff --git a/.github/scripts/build_triton_wheel.py b/.github/scripts/build_triton_wheel.py
@@ -189,9 +189,9 @@ def main() -> None:
 
     build_triton(
         device=args.device,
-        commit_hash=args.commit_hash
-        if args.commit_hash
-        else read_triton_pin(args.device),
+        commit_hash=(
+            args.commit_hash if args.commit_hash else read_triton_pin(args.device)
+        ),
         version=args.triton_version,
         build_conda=args.build_conda,
         py_version=args.py_version,

diff --git a/.github/scripts/generate_binary_build_matrix.py b/.github/scripts/generate_binary_build_matrix.py
@@ -391,7 +391,7 @@ def generate_wheels_matrix(
                         "container_image": WHEEL_CONTAINER_IMAGES[arch_version],
                         "package_type": package_type,
                         "pytorch_extra_install_requirements": (
-                            PYTORCH_EXTRA_INSTALL_REQUIREMENTS[arch_version]  # fmt: skip
+                            PYTORCH_EXTRA_INSTALL_REQUIREMENTS[arch_version]
                             if os != "linux-aarch64"
                             else ""
                         ),
@@ -414,7 +414,7 @@ def generate_wheels_matrix(
                             "container_image": WHEEL_CONTAINER_IMAGES[arch_version],
                             "package_type": package_type,
                             "pytorch_extra_install_requirements": (
-                                PYTORCH_EXTRA_INSTALL_REQUIREMENTS[arch_version]  # fmt: skip
+                                PYTORCH_EXTRA_INSTALL_REQUIREMENTS[arch_version]
                                 if os != "linux-aarch64"
                                 else ""
                             ),
@@ -463,7 +463,7 @@ def generate_wheels_matrix(
                             ".", "_"
                         ),
                         "pytorch_extra_install_requirements": (
-                            PYTORCH_EXTRA_INSTALL_REQUIREMENTS["12.1"]  # fmt: skip
+                            PYTORCH_EXTRA_INSTALL_REQUIREMENTS["12.1"]
                             if os != "linux"
                             else ""
                         ),

diff --git a/.github/scripts/trymerge.py b/.github/scripts/trymerge.py
@@ -670,7 +670,7 @@ def skip_func(idx: int, candidate: "GitHubPR") -> bool:
         if not open_only or not candidate.is_closed():
             return False
         print(
-            f"Skipping {idx+1} of {len(rev_list)} PR (#{candidate.pr_num}) as its already been merged"
+            f"Skipping {idx + 1} of {len(rev_list)} PR (#{candidate.pr_num}) as its already been merged"
         )
         return True
 
@@ -1419,9 +1419,9 @@ def find_matching_merge_rule(
         pending_checks, failed_checks, _ = categorize_checks(
             checks,
             required_checks,
-            ok_failed_checks_threshold=IGNORABLE_FAILED_CHECKS_THESHOLD
-            if rule.ignore_flaky_failures
-            else 0,
+            ok_failed_checks_threshold=(
+                IGNORABLE_FAILED_CHECKS_THESHOLD if rule.ignore_flaky_failures else 0
+            ),
         )
 
         # categorize_checks assumes all tests are required if required_checks is empty.
@@ -1491,7 +1491,7 @@ def checks_to_str(checks: List[Tuple[str, Optional[str]]]) -> str:
 
 
 def checks_to_markdown_bullets(
-    checks: List[Tuple[str, Optional[str], Optional[int]]]
+    checks: List[Tuple[str, Optional[str], Optional[int]]],
 ) -> List[str]:
     return [
         f"- [{c[0]}]({c[1]})" if c[1] is not None else f"- {c[0]}" for c in checks[:5]
@@ -2202,9 +2202,9 @@ def merge(
                 checks,
                 required_checks
                 + [x for x in checks.keys() if x not in required_checks],
-                ok_failed_checks_threshold=IGNORABLE_FAILED_CHECKS_THESHOLD
-                if ignore_flaky_failures
-                else 0,
+                ok_failed_checks_threshold=(
+                    IGNORABLE_FAILED_CHECKS_THESHOLD if ignore_flaky_failures else 0
+                ),
             )
             # HACK until GitHub will be better about surfacing those
             startup_failures = filter_checks_with_lambda(

diff --git a/.lintrunner.toml b/.lintrunner.toml
@@ -1474,8 +1474,8 @@ init_command = [
     'tools/linter/adapters/pip_init.py',
     '--dry-run={{DRYRUN}}',
     '--no-black-binary',
-    'black==23.12.1',
     'ufmt==2.7.0',
+    'black==24.4.2',
     'usort==1.0.8.post1',
     'isort==5.13.2',
 ]

diff --git a/benchmarks/distributed/rpc/rl/launcher.py b/benchmarks/distributed/rpc/rl/launcher.py
@@ -209,9 +209,8 @@ def main():
         x_axis_variables
     ):  # run benchmark for every x axis variable
         if len(x_axis_variables) > 1:
-            args[
-                args["x_axis_name"]
-            ] = x_axis_variable  # set x axis variable for this benchmark iteration
+            # Set x axis variable for this benchmark iteration
+            args[args["x_axis_name"]] = x_axis_variable
         processes = []
         start_time = time.time()
         for rank in range(args["world_size"]):

diff --git a/benchmarks/dynamo/common.py b/benchmarks/dynamo/common.py
@@ -481,7 +481,7 @@ def print_summary_table(data, print_dataframe=False):
                 print(col.ljust(width), f"mean={data[col].mean():.3f}x")
             elif col in ("accuracy"):
                 pass_rate = (data[col] == "pass").mean()
-                print(col.ljust(width), f"pass_rate={100*pass_rate:.2f}%")
+                print(col.ljust(width), f"pass_rate={100 * pass_rate:.2f}%")
             else:
                 cdata = data[col]
                 print(
@@ -1372,9 +1372,7 @@ def load(cls, model, example_inputs, device):
                 example_kwargs,
             )
             with torch.no_grad():
-                so_path = torch._inductor.aot_compile(
-                    gm, example_args, example_kwargs
-                )  # type: ignore[arg-type]
+                so_path = torch._inductor.aot_compile(gm, example_args, example_kwargs)  # type: ignore[arg-type]
 
             cls.cache[key] = torch._export.aot_load(so_path, device)
 
@@ -1531,12 +1529,10 @@ def _generate_onnx_model_directory(
         return model_path
 
     @abc.abstractmethod
-    def format_pt_inputs(self, pt_inputs: Any) -> Sequence[torch.Tensor]:
-        ...
+    def format_pt_inputs(self, pt_inputs: Any) -> Sequence[torch.Tensor]: ...
 
     @abc.abstractmethod
-    def format_pt_outputs(self, pt_outputs: Any) -> Sequence[torch.Tensor]:
-        ...
+    def format_pt_outputs(self, pt_outputs: Any) -> Sequence[torch.Tensor]: ...
 
     def adapt_pt_inputs_to_onnx(self, pt_inputs) -> Mapping[str, np.ndarray]:
         pt_inputs = self.format_pt_inputs(pt_inputs)
@@ -2194,9 +2190,9 @@ def cast_to(dtype, model, inputs):
         model = model.to(dtype)
 
     inputs = tree_map(
-        lambda x: x.to(dtype)
-        if isinstance(x, torch.Tensor) and x.is_floating_point()
-        else x,
+        lambda x: (
+            x.to(dtype) if isinstance(x, torch.Tensor) and x.is_floating_point() else x
+        ),
         inputs,
     )
     return model, inputs
@@ -2631,9 +2627,11 @@ def deepcopy_and_maybe_parallelize(self, model):
             model = FSDP(
                 model,
                 use_orig_params=True,
-                device_id=torch.cuda.current_device()
-                if self.args.devices[-1] == "cuda"
-                else None,
+                device_id=(
+                    torch.cuda.current_device()
+                    if self.args.devices[-1] == "cuda"
+                    else None
+                ),
                 mixed_precision=mp_policy,
                 limit_all_gathers=True,
                 auto_wrap_policy=self.get_fsdp_auto_wrap_policy(self.args.only),
@@ -2698,9 +2696,11 @@ def record_status(accuracy_status, dynamo_start_stats):
                 self.init_optimizer(name, current_device, model_fp64.parameters())
                 fp64_outputs = self.run_n_iterations(model_fp64, inputs_fp64)
                 fp64_outputs = tree_map(
-                    lambda x: x.to(torch.float64)
-                    if isinstance(x, torch.Tensor) and x.is_floating_point()
-                    else x,
+                    lambda x: (
+                        x.to(torch.float64)
+                        if isinstance(x, torch.Tensor) and x.is_floating_point()
+                        else x
+                    ),
                     fp64_outputs,
                 )
             except Exception:
@@ -3102,9 +3102,9 @@ def warmup(fn, model, example_inputs, mode, niters=10):
                 experiment_kwargs["dynamo_peak_mem"] = dynamo_peak_mem
                 experiment_kwargs["dynamo_stats"] = dynamo_stats
                 if self.args.profile_dynamo_cache_lookup:
-                    experiment_kwargs[
-                        "cache_lookup_latency"
-                    ] = dynamo_cache_lookup_latency
+                    experiment_kwargs["cache_lookup_latency"] = (
+                        dynamo_cache_lookup_latency
+                    )
 
             if experiment.func is speedup_experiment_onnx:
                 experiment = functools.partial(
@@ -3258,9 +3258,9 @@ def warmup(fn, model, example_inputs, mode, niters=5):
                 experiment_kwargs["dynamo_peak_mem"] = dynamo_peak_mem
                 experiment_kwargs["dynamo_stats"] = dynamo_stats
                 if self.args.profile_dynamo_cache_lookup:
-                    experiment_kwargs[
-                        "cache_lookup_latency"
-                    ] = dynamo_cache_lookup_latency
+                    experiment_kwargs["cache_lookup_latency"] = (
+                        dynamo_cache_lookup_latency
+                    )
 
             if experiment.func is coverage_experiment:
                 ok, total = Stats.reset_counters()
@@ -4709,7 +4709,7 @@ def detect_and_mark_batch(t):
         for i, name in enumerate(model_names):
             current_name = name
             if args.progress:
-                print(f"Running model {i+1}/{nmodels}", flush=True)
+                print(f"Running model {i + 1}/{nmodels}", flush=True)
 
             try:
                 timeout = args.timeout

diff --git a/benchmarks/dynamo/join_results.py b/benchmarks/dynamo/join_results.py
@@ -2,6 +2,7 @@
 A tool to merge multiple csv files (generated by torchbench.py/etc) into a single csv file.
 Performs an outer join based on the benchmark name, filling in any missing data with zeros.
 """
+
 import argparse
 import functools
 import operator

diff --git a/benchmarks/dynamo/microbenchmarks/analyze_templates.py b/benchmarks/dynamo/microbenchmarks/analyze_templates.py
@@ -4,6 +4,7 @@
 
 That file can be fed into this script to generate the minimizes total, weighted matmul time as a function of allowed templates.
 """
+
 import json
 
 import click

diff --git a/benchmarks/dynamo/microbenchmarks/fx_microbenchmarks.py b/benchmarks/dynamo/microbenchmarks/fx_microbenchmarks.py
@@ -24,7 +24,7 @@ def fn():
             pass
 
     t = min(timeit.repeat(fn, number=K, repeat=3))
-    print(f"iterating over {N*K} FX nodes took {t:.1f}s ({N*K/t:.0f} nodes/s)")
+    print(f"iterating over {N * K} FX nodes took {t:.1f}s ({N * K / t:.0f} nodes/s)")
 
 
 if __name__ == "__main__":

diff --git a/benchmarks/dynamo/microbenchmarks/overheads.py b/benchmarks/dynamo/microbenchmarks/overheads.py
@@ -19,7 +19,7 @@ def bench(name, fn, requires_grad):
     end = time.perf_counter()
 
     results = timeit.repeat(lambda: fn(x), number=1000, repeat=1000)
-    print(f"{name} {np.median(results)*1000:.1f}us (warmup={end-start:.1f}s)")
+    print(f"{name} {np.median(results) * 1000:.1f}us (warmup={end - start:.1f}s)")
 
 
 def main():

diff --git a/benchmarks/dynamo/runner.py b/benchmarks/dynamo/runner.py
@@ -548,7 +548,7 @@ def env_var(name):
         out_io.write(f"Number CUDA Devices: {torch.cuda.device_count()}\n")
         out_io.write(f"Device Name: {torch.cuda.get_device_name(0)}\n")
         out_io.write(
-            f"Device Memory [GB]: {torch.cuda.get_device_properties(0).total_memory/1e9}\n"
+            f"Device Memory [GB]: {torch.cuda.get_device_properties(0).total_memory / 1e9}\n"
         )
 
     title = "## Build Summary"

diff --git a/benchmarks/dynamo/training_loss.py b/benchmarks/dynamo/training_loss.py
@@ -193,9 +193,9 @@ def main():
     print(
         f"Train model on {args.epochs} epochs with backend {args.backend} and optimizer {args.optimizer}:"
     )
-    print(f"PyTorch spent {timedelta(seconds=native_elapsed/args.epochs)} per epoch")
+    print(f"PyTorch spent {timedelta(seconds=native_elapsed / args.epochs)} per epoch")
     print(
-        f"TorchDynamo spent {timedelta(seconds=dynamo_elapsed/args.epochs)} per epoch"
+        f"TorchDynamo spent {timedelta(seconds=dynamo_elapsed / args.epochs)} per epoch"
     )
 
 

diff --git a/benchmarks/instruction_counts/applications/ci.py b/benchmarks/instruction_counts/applications/ci.py
@@ -1,4 +1,5 @@
 """Collect instruction counts for continuous integration."""
+
 import argparse
 import hashlib
 import json

diff --git a/benchmarks/instruction_counts/core/api.py b/benchmarks/instruction_counts/core/api.py
@@ -1,4 +1,5 @@
 """Key enums and structs used to handle data flow within the benchmark."""
+
 import dataclasses
 import enum
 import itertools as it

diff --git a/benchmarks/instruction_counts/core/expand.py b/benchmarks/instruction_counts/core/expand.py
@@ -2,6 +2,7 @@
 
 This is mostly string manipulation, with just a bit of importlib magic.
 """
+
 import importlib.abc
 import importlib.util
 import itertools as it

diff --git a/benchmarks/instruction_counts/core/types.py b/benchmarks/instruction_counts/core/types.py
@@ -1,4 +1,5 @@
 """Type annotations for various benchmark objects."""
+
 from typing import Any, Dict, Optional, Tuple, Union
 
 from core.api import AutoLabels, GroupedBenchmark, TimerArgs

diff --git a/benchmarks/instruction_counts/execution/runner.py b/benchmarks/instruction_counts/execution/runner.py
@@ -1,4 +1,5 @@
 """Run benchmarks while handling parallelism, isolation, and fault tolerance."""
+
 import math
 import multiprocessing
 import subprocess

diff --git a/benchmarks/instruction_counts/execution/work.py b/benchmarks/instruction_counts/execution/work.py
@@ -1,4 +1,5 @@
 """Handle the details of subprocess calls and retries for a given benchmark run."""
+
 import dataclasses
 import json
 import os

diff --git a/benchmarks/instruction_counts/main.py b/benchmarks/instruction_counts/main.py
@@ -5,6 +5,7 @@
 components) in future iterations. However this allows us to excercise the
 underlying benchmark generation infrastructure in the mean time.
 """
+
 import argparse
 import sys
 from typing import List

diff --git a/benchmarks/instruction_counts/worker/main.py b/benchmarks/instruction_counts/worker/main.py
@@ -15,6 +15,7 @@
 Because this file only expects to run in a child context, error handling means
 plumbing failures up to the caller, not raising in this process.
 """
+
 import argparse
 import dataclasses
 import io

diff --git a/benchmarks/operator_benchmark/pt/diag_test.py b/benchmarks/operator_benchmark/pt/diag_test.py
@@ -24,9 +24,11 @@
 class DiagBenchmark(op_bench.TorchBenchmarkBase):
     def init(self, dim, M, N, diagonal, out, device):
         self.inputs = {
-            "input": torch.rand(M, N, device=device)
-            if dim == 2
-            else torch.rand(M, device=device),
+            "input": (
+                torch.rand(M, N, device=device)
+                if dim == 2
+                else torch.rand(M, device=device)
+            ),
             "diagonal": diagonal,
             "out": out,
             "out_tensor": torch.tensor(

diff --git a/benchmarks/operator_benchmark/pt/matmul_test.py b/benchmarks/operator_benchmark/pt/matmul_test.py
@@ -34,12 +34,16 @@
 class MatMulBenchmark(op_bench.TorchBenchmarkBase):
     def init(self, M, N, K, trans_a, trans_b, device):
         self.inputs = {
-            "input_one": torch.rand(M, N, device=device)
-            if trans_a
-            else torch.rand(N, M, device=device).t(),
-            "input_two": torch.rand(N, K, device=device)
-            if trans_b
-            else torch.rand(K, N, device=device).t(),
+            "input_one": (
+                torch.rand(M, N, device=device)
+                if trans_a
+                else torch.rand(N, M, device=device).t()
+            ),
+            "input_two": (
+                torch.rand(N, K, device=device)
+                if trans_b
+                else torch.rand(K, N, device=device).t()
+            ),
         }
         self.set_module_name("matmul")
 

diff --git a/benchmarks/sparse/dlmc/matmul_bench.py b/benchmarks/sparse/dlmc/matmul_bench.py
@@ -97,13 +97,15 @@ def filter_ops(operation):
                             f"{OPS_MAP[operation]}(x, y)",
                         ),
                         (
-                            test_name,
-                            device,
-                            "scipy:" + operation,
-                            "scipy_matmul(sx, sy)",
-                        )
-                        if device == "cpu"
-                        else None,
+                            (
+                                test_name,
+                                device,
+                                "scipy:" + operation,
+                                "scipy_matmul(sx, sy)",
+                            )
+                            if device == "cpu"
+                            else None
+                        ),
                     ],
                 )
             )
-Original file line number
+Diff line change
@@ Expand Up / @@ -4,6 +4,7 @@ @@
     That file can be fed into this script to generate the minimizes total, weighted matmul time as a function of allowed templates.
     """
     import json
     import click
@@ Expand Down @@