[BE] Format .ci/ / .github/ / benchmarks/ / functorch/ / tools/ / torchgen/ with ruff format (pytorch#132577)

Pull Request resolved: pytorch#132577
Approved by: https://github.com/malfet

Author: XuehaiPan

.ci/pytorch/create_test_cert.py

@@ -45,8 +45,7 @@ def create_cert(path, C, ST, L, O, key):
         .not_valid_before(datetime.now(timezone.utc))
         .not_valid_after(
             # Our certificate will be valid for 10 days
-            datetime.now(timezone.utc)
-            + timedelta(days=10)
+            datetime.now(timezone.utc) + timedelta(days=10)
         )
         .add_extension(
             x509.BasicConstraints(ca=True, path_length=None),
@@ -91,8 +90,7 @@ def sign_certificate_request(path, csr_cert, ca_cert, private_ca_key):
         .not_valid_before(datetime.now(timezone.utc))
         .not_valid_after(
             # Our certificate will be valid for 10 days
-            datetime.now(timezone.utc)
-            + timedelta(days=10)
+            datetime.now(timezone.utc) + timedelta(days=10)
             # Sign our certificate with our private key
         )
         .sign(private_ca_key, hashes.SHA256())

.github/scripts/generate_binary_build_matrix.py

@@ -409,7 +409,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 ""
                         ),
@@ -457,7 +457,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" and gpu_arch_type != "xpu"
                             else ""
                         ),

.github/scripts/trymerge.py

@@ -1506,7 +1506,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]

benchmarks/distributed/ddp/diff.py

@@ -51,9 +51,7 @@ def main():
         print()
         print(f"{'':>10s}", end="")  # noqa: E999
         for _ in [75, 95]:
-            print(
-                f"{'sec/iter':>16s}{'ex/sec':>10s}{'diff':>10s}", end=""
-            )  # noqa: E999
+            print(f"{'sec/iter':>16s}{'ex/sec':>10s}{'diff':>10s}", end="")  # noqa: E999
         print()
 
         # Print measurements

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"]):

benchmarks/dynamo/common.py

@@ -1391,9 +1391,7 @@ def load(cls, model, example_inputs, device):
                 strict=False,
             ).module()
             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)
 
@@ -1559,12 +1557,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, npt.NDArray]:
         pt_inputs = self.format_pt_inputs(pt_inputs)
@@ -3134,9 +3130,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(
@@ -3290,9 +3286,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()
@@ -4324,7 +4320,14 @@ def run(runner, args, original_dir=None):
         runner.skip_models.clear()
 
     experiment = null_experiment
-    global current_name, current_device, current_batch_size, output_filename, disable_output, optimize_ctx, current_onnx_compiler
+    global \
+        current_name, \
+        current_device, \
+        current_batch_size, \
+        output_filename, \
+        disable_output, \
+        optimize_ctx, \
+        current_onnx_compiler
     optimize_ctx = contextlib.nullcontext()
 
     if args.disable_output:

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

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

benchmarks/fastrnns/bench.py

@@ -214,8 +214,7 @@ def bench(rnn_runners, group_name, print_json=False, sep=" ", **params):
             k: {"avg": v.avg_fwd, "std": v.std_fwd, "info": v.info_fwd}
             for k, v in results.items()
         },
-        group_name
-        + "-backward": {
+        f"{group_name}-backward": {
             k: {"avg": v.avg_bwd, "std": v.std_bwd, "info": v.info_bwd}
             for k, v in results.items()
         },

benchmarks/gpt_fast/mixtral_moe_quantize.py

@@ -184,7 +184,5 @@ def forward(self, x, expert_indices):
         ].to(x.dtype)
         expert_outs = torch.einsum(
             "tao, taio -> tai", (x1 * x3), w2_weights
-        ) * self.scales2[expert_indices].to(
-            x.dtype
-        )  # [T, A, D, D]
+        ) * self.scales2[expert_indices].to(x.dtype)  # [T, A, D, D]
         return expert_outs

benchmarks/instruction_counts/applications/ci.py

@@ -1,5 +1,7 @@
 """Collect instruction counts for continuous integration."""
+
 # mypy: ignore-errors
+
 import argparse
 import hashlib
 import json

benchmarks/instruction_counts/core/api.py

@@ -1,5 +1,7 @@
 """Key enums and structs used to handle data flow within the benchmark."""
+
 # mypy: ignore-errors
+
 import dataclasses
 import enum
 import itertools as it

benchmarks/instruction_counts/core/expand.py

@@ -2,7 +2,9 @@
 
 This is mostly string manipulation, with just a bit of importlib magic.
 """
+
 # mypy: ignore-errors
+
 import importlib.abc
 import importlib.util
 import itertools as it

benchmarks/instruction_counts/core/types.py

@@ -1,5 +1,7 @@
 """Type annotations for various benchmark objects."""
+
 # mypy: ignore-errors
+
 from typing import Any, Dict, Optional, Tuple, Union
 
 from core.api import AutoLabels, GroupedBenchmark, TimerArgs

benchmarks/instruction_counts/definitions/setup.py

@@ -1,5 +1,7 @@
 """Define some common setup blocks which benchmarks can reuse."""
+
 # mypy: ignore-errors
+
 import enum
 
 from core.api import GroupedSetup

benchmarks/instruction_counts/execution/runner.py

@@ -1,5 +1,7 @@
 """Run benchmarks while handling parallelism, isolation, and fault tolerance."""
+
 # mypy: ignore-errors
+
 import math
 import multiprocessing
 import subprocess

benchmarks/instruction_counts/execution/work.py

@@ -1,5 +1,7 @@
 """Handle the details of subprocess calls and retries for a given benchmark run."""
+
 # mypy: ignore-errors
+
 import dataclasses
 import json
 import os

benchmarks/instruction_counts/main.py

@@ -5,7 +5,9 @@
 components) in future iterations. However this allows us to excercise the
 underlying benchmark generation infrastructure in the mean time.
 """
+
 # mypy: ignore-errors
+
 import argparse
 import sys
 from typing import List

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

benchmarks/operator_benchmark/pt/qrnn_test.py

@@ -48,14 +48,20 @@ def init(self, I, H, NL, B, D, dtype):
         )[0]
 
         x = torch.randn(
-            sequence_len, batch_size, I  # sequence length  # batch size
-        )  # Number of features in X
+            sequence_len,  # sequence length
+            batch_size,  # batch size
+            I,  # Number of features in X
+        )
         h = torch.randn(
-            NL * (D + 1), batch_size, H  # layer_num * dir_num  # batch size
-        )  # hidden size
+            NL * (D + 1),  # layer_num * dir_num
+            batch_size,  # batch size
+            H,  # hidden size
+        )
         c = torch.randn(
-            NL * (D + 1), batch_size, H  # layer_num * dir_num  # batch size
-        )  # hidden size
+            NL * (D + 1),  # layer_num * dir_num
+            batch_size,  # batch size
+            H,  # hidden size
+        )
 
         self.inputs = {"x": x, "h": h, "c": c}
         self.set_module_name("QLSTM")

benchmarks/transformer/better_transformer_vs_mha_functional.py

@@ -152,8 +152,8 @@ def run(
     result_entry["sequence_length"] = sequence_length
     result_entry["n_heads"] = num_heads
     result_entry["embed_dim"] = embed_dim
-    result_entry["time_native_mha_slow(\u00B5s)"] = f"{time_native_mha_slow:.3f}"
-    result_entry["time_native_mha_fast (\u00B5s)"] = f"{time_native_mha_fast:.3f}"
+    result_entry["time_native_mha_slow(\u00b5s)"] = f"{time_native_mha_slow:.3f}"
+    result_entry["time_native_mha_fast (\u00b5s)"] = f"{time_native_mha_fast:.3f}"
     result_entry["speedup flash_mha v native_mha"] = f"{speedup_fast_internal:.3f}"
     result_entry["padding"] = f"{padding:.3f}"
     return result_entry

benchmarks/transformer/sdp.py

@@ -82,10 +82,10 @@ def get_entries(self) -> List:
     @classmethod
     def get_entry_names(cls) -> List[str]:
         return [
-            "nn_mha_time (\u00B5s)",
-            "compiled_nn_mha_time (\u00B5s)",
-            "composite_mha_time (\u00B5s)",
-            "compiled_composite_mha_time (\u00B5s)",
+            "nn_mha_time (\u00b5s)",
+            "compiled_nn_mha_time (\u00b5s)",
+            "composite_mha_time (\u00b5s)",
+            "compiled_composite_mha_time (\u00b5s)",
         ]
 
 

functorch/dim/dim.py

@@ -32,8 +32,7 @@ def __del__(self):
         if self._vmap_level is not None:
             _vmap_active_levels[self._vmap_stack].alive = False  # noqa: F821
             while (
-                not _vmap_levels[-1].alive
-                and current_level() == _vmap_levels[-1].level  # noqa: F821
+                not _vmap_levels[-1].alive and current_level() == _vmap_levels[-1].level  # noqa: F821
             ):
                 _vmap_decrement_nesting()  # noqa: F821
                 _vmap_levels.pop()

functorch/einops/_parsing.py

@@ -22,6 +22,7 @@
 OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 SOFTWARE.
 """
+
 from __future__ import annotations
 
 import keyword
@@ -283,16 +284,16 @@ def comma_separate(collection: Collection[Union[str, Collection[str]]]) -> str:
         str: the comma-separated string
 
     Examples:
-        >>> comma_separate(('d0',))
+        >>> comma_separate(("d0",))
         'd0'
 
-        >>> comma_separate(('d0', 'd1', 'd2', 'd3'))
+        >>> comma_separate(("d0", "d1", "d2", "d3"))
         'd0, d1, d2, d3'
 
-        >>> comma_separate([('d1', 'd4')])
+        >>> comma_separate([("d1", "d4")])
         '(d1, d4)'
 
-        >>> comma_separate([('d0',), (), ('d1',), ('d2',), ('d3', 'd4')])
+        >>> comma_separate([("d0",), (), ("d1",), ("d2",), ("d3", "d4")])
         '(d0,), (), (d1,), (d2,), (d3, d4)'
     """
     return ", ".join(

functorch/einops/rearrange.py

@@ -95,7 +95,7 @@ def _create_rearrange_callable(
             raise ValueError(f"Unexpected dimension: {dimension}")
 
     def composition_to_dims(
-        composition: Sequence[Union[List[Union[str, AnonymousAxis]], str]]
+        composition: Sequence[Union[List[Union[str, AnonymousAxis]], str]],
     ) -> List[Union[str, Tuple[str, ...]]]:
         """Convert a `ParsedExpression.composition` into a `Tensor.__getitem__` index of strings representing first
         class dims."""
@@ -171,31 +171,31 @@ def rearrange(
         >>> images = torch.randn((32, 30, 40, 3))
 
         >>> # stack along first (batch) axis, output is a single array
-        >>> rearrange(images, 'b h w c -> b h w c').shape
+        >>> rearrange(images, "b h w c -> b h w c").shape
         torch.Size([32, 30, 40, 3])
 
         >>> # concatenate images along height (vertical axis), 960 = 32 * 30
-        >>> rearrange(images, 'b h w c -> (b h) w c').shape
+        >>> rearrange(images, "b h w c -> (b h) w c").shape
         torch.Size([960, 40, 3])
 
         >>> # concatenated images along horizontal axis, 1280 = 32 * 40
-        >>> rearrange(images, 'b h w c -> h (b w) c').shape
+        >>> rearrange(images, "b h w c -> h (b w) c").shape
         torch.Size([30, 1280, 3])
 
         >>> # reordered axes to "b c h w" format for deep learning
-        >>> rearrange(images, 'b h w c -> b c h w').shape
+        >>> rearrange(images, "b h w c -> b c h w").shape
         torch.Size([32, 3, 30, 40])
 
         >>> # flattened each image into a vector, 3600 = 30 * 40 * 3
-        >>> rearrange(images, 'b h w c -> b (c h w)').shape
+        >>> rearrange(images, "b h w c -> b (c h w)").shape
         torch.Size([32, 3600])
 
         >>> # split each image into 4 smaller (top-left, top-right, bottom-left, bottom-right), 128 = 32 * 2 * 2
-        >>> rearrange(images, 'b (h1 h) (w1 w) c -> (b h1 w1) h w c', h1=2, w1=2).shape
+        >>> rearrange(images, "b (h1 h) (w1 w) c -> (b h1 w1) h w c", h1=2, w1=2).shape
         torch.Size([128, 15, 20, 3])
 
         >>> # space-to-depth operation
-        >>> rearrange(images, 'b (h h1) (w w1) c -> b h w (c h1 w1)', h1=2, w1=2).shape
+        >>> rearrange(images, "b (h h1) (w w1) c -> b h w (c h1 w1)", h1=2, w1=2).shape
         torch.Size([32, 15, 20, 12])
     """
     if not isinstance(tensor, torch.Tensor):