[BE][CI] bump ruff to 0.9.2: multiline assert statements (pytorch#144546)

Reference: https://docs.astral.sh/ruff/formatter/black/#assert-statements

> Unlike Black, Ruff prefers breaking the message over breaking the assertion, similar to how both Ruff and Black prefer breaking the assignment value over breaking the assignment target:
>
> ```python
> # Input
> assert (
>     len(policy_types) >= priority + num_duplicates
> ), f"This tests needs at least {priority+num_duplicates} many types."
>
>
> # Black
> assert (
>     len(policy_types) >= priority + num_duplicates
> ), f"This tests needs at least {priority+num_duplicates} many types."
>
> # Ruff
> assert len(policy_types) >= priority + num_duplicates, (
>     f"This tests needs at least {priority + num_duplicates} many types."
> )
> ```

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

Author: XuehaiPan

.github/scripts/label_utils.py

@@ -63,9 +63,9 @@ def gh_get_labels(org: str, repo: str) -> list[str]:
     update_labels(labels, info)
 
     last_page = get_last_page_num_from_header(header)
-    assert (
-        last_page > 0
-    ), "Error reading header info to determine total number of pages of labels"
+    assert last_page > 0, (
+        "Error reading header info to determine total number of pages of labels"
+    )
     for page_number in range(2, last_page + 1):  # skip page 1
         _, info = request_for_labels(prefix + f"&page={page_number}")
         update_labels(labels, info)

.lintrunner.toml

@@ -1476,7 +1476,7 @@ init_command = [
     'black==23.12.1',
     'usort==1.0.8.post1',
     'isort==5.13.2',
-    'ruff==0.8.4',  # sync with RUFF
+    'ruff==0.9.2',  # sync with RUFF
 ]
 is_formatter = true
 
@@ -1561,7 +1561,7 @@ init_command = [
     'python3',
     'tools/linter/adapters/pip_init.py',
     '--dry-run={{DRYRUN}}',
-    'ruff==0.8.4',  # sync with PYFMT
+    'ruff==0.9.2',  # sync with PYFMT
 ]
 is_formatter = true
 

benchmarks/dynamo/cachebench.py

@@ -101,9 +101,9 @@ def _run_torchbench_model(
     torchbench_file = os.path.join(
         os.path.dirname(cur_file), BENCHMARK_FILE[cmd_args.benchmark]
     )
-    assert os.path.exists(
-        torchbench_file
-    ), f"Torchbench does not exist at {torchbench_file}"
+    assert os.path.exists(torchbench_file), (
+        f"Torchbench does not exist at {torchbench_file}"
+    )
 
     dynamic = cmd_args.dynamic
     dynamic_args = ["--dynamic-shapes", "--dynamic-batch-only"] if dynamic else []

benchmarks/dynamo/common.py

@@ -1006,9 +1006,9 @@ def latency_experiment_summary(suite_name, args, model, timings, **kwargs):
         row,
     )
     c_headers, c_data = torch._dynamo.utils.compile_times(repr="csv", aggregate=True)
-    assert (
-        output_filename.find(".csv") > 0
-    ), f"expected output_filename to be a .csv, but got {output_filename}"
+    assert output_filename.find(".csv") > 0, (
+        f"expected output_filename to be a .csv, but got {output_filename}"
+    )
     write_outputs(
         output_filename[:-4] + "_compilation_metrics.csv",
         first_headers + c_headers,
@@ -1182,9 +1182,9 @@ def maybe_mark_profile(*args, **kwargs):
         row,
     )
     c_headers, c_data = torch._dynamo.utils.compile_times(repr="csv", aggregate=True)
-    assert (
-        output_filename.find(".csv") > 0
-    ), f"expected output_filename to be a .csv, but got {output_filename}"
+    assert output_filename.find(".csv") > 0, (
+        f"expected output_filename to be a .csv, but got {output_filename}"
+    )
     write_outputs(
         output_filename[:-4] + "_compilation_metrics.csv",
         first_headers + c_headers,
@@ -1997,16 +1997,16 @@ def get_fsdp_auto_wrap_policy(self, model_name: str):
     def deepcopy_and_maybe_parallelize(self, model):
         model = self.deepcopy_model(model)
         if self.args.ddp:
-            assert (
-                torch.distributed.is_available()
-            ), "Can't use DDP without a distributed enabled build"
+            assert torch.distributed.is_available(), (
+                "Can't use DDP without a distributed enabled build"
+            )
             from torch.nn.parallel import DistributedDataParallel as DDP
 
             model = DDP(model, find_unused_parameters=True)
         elif self.args.fsdp:
-            assert (
-                torch.distributed.is_available()
-            ), "Can't use FSDP without a distributed enabled build"
+            assert torch.distributed.is_available(), (
+                "Can't use FSDP without a distributed enabled build"
+            )
             from torch.distributed.fsdp import (
                 FullyShardedDataParallel as FSDP,
                 MixedPrecision,
@@ -2375,9 +2375,9 @@ def run_performance_test_non_alternate(
         self, name, model, example_inputs, optimize_ctx, experiment, tag=None
     ):
         "Run performance test in non-alternately."
-        assert (
-            experiment.func is latency_experiment
-        ), "Must run with latency_experiment."
+        assert experiment.func is latency_experiment, (
+            "Must run with latency_experiment."
+        )
 
         def warmup(fn, model, example_inputs, mode, niters=10):
             peak_mem = 0

benchmarks/dynamo/microbenchmarks/bench_mm_fusion.py

@@ -81,9 +81,9 @@ def fn():
         torch._dynamo.reset()
         torch._inductor.metrics.reset()
         triton_mm_ms, _, _ = benchmarker.benchmark_gpu(fn)
-        assert (
-            torch._inductor.metrics.generated_kernel_count == 1
-        ), "codegen #kernel != 1"
+        assert torch._inductor.metrics.generated_kernel_count == 1, (
+            "codegen #kernel != 1"
+        )
         row.extend([tflops(torch_mm_ms), tflops(triton_mm_ms)])
 
     p.add_row(row)

benchmarks/dynamo/microbenchmarks/operator_inp_utils.py

@@ -265,9 +265,9 @@ def __init__(self, json_file_path):
     def get_inputs_for_operator(
         self, operator, dtype=None, device="cuda"
     ) -> Generator[tuple[Iterable[Any], dict[str, Any]], None, None]:
-        assert (
-            str(operator) in self.operator_db
-        ), f"Could not find {operator}, must provide overload"
+        assert str(operator) in self.operator_db, (
+            f"Could not find {operator}, must provide overload"
+        )
 
         if "embedding" in str(operator):
             log.warning("Embedding inputs NYI, input data cannot be randomized")
@@ -302,9 +302,9 @@ def get_all_ops(self):
             yield op
 
     def get_call_frequency(self, op):
-        assert (
-            str(op) in self.operator_db
-        ), f"Could not find {op}, must provide overload"
+        assert str(op) in self.operator_db, (
+            f"Could not find {op}, must provide overload"
+        )
 
         count = 0
         for counter in self.operator_db[str(op)].values():

benchmarks/dynamo/runner.py

@@ -711,9 +711,9 @@ def clean_batch_sizes(self, frames):
             for idx, (batch_a, batch_b) in enumerate(
                 zip(batch_sizes, frame_batch_sizes)
             ):
-                assert (
-                    batch_a == batch_b or batch_a == 0 or batch_b == 0
-                ), f"a={batch_a}, b={batch_b}"
+                assert batch_a == batch_b or batch_a == 0 or batch_b == 0, (
+                    f"a={batch_a}, b={batch_b}"
+                )
                 batch_sizes[idx] = max(batch_a, batch_b)
         for frame in frames:
             frame["batch_size"] = batch_sizes

benchmarks/functional_autograd_benchmark/torchaudio_models.py

@@ -538,21 +538,21 @@ def forward(
             query.size(-1),
         )
         q, k, v = self.in_proj_container(query, key, value)
-        assert (
-            q.size(-1) % self.nhead == 0
-        ), "query's embed_dim must be divisible by the number of heads"
+        assert q.size(-1) % self.nhead == 0, (
+            "query's embed_dim must be divisible by the number of heads"
+        )
         head_dim = q.size(-1) // self.nhead
         q = q.reshape(tgt_len, bsz * self.nhead, head_dim)
 
-        assert (
-            k.size(-1) % self.nhead == 0
-        ), "key's embed_dim must be divisible by the number of heads"
+        assert k.size(-1) % self.nhead == 0, (
+            "key's embed_dim must be divisible by the number of heads"
+        )
         head_dim = k.size(-1) // self.nhead
         k = k.reshape(src_len, bsz * self.nhead, head_dim)
 
-        assert (
-            v.size(-1) % self.nhead == 0
-        ), "value's embed_dim must be divisible by the number of heads"
+        assert v.size(-1) % self.nhead == 0, (
+            "value's embed_dim must be divisible by the number of heads"
+        )
         head_dim = v.size(-1) // self.nhead
         v = v.reshape(src_len, bsz * self.nhead, head_dim)
 
@@ -629,9 +629,9 @@ def forward(
                 attn_mask = torch.nn.functional.pad(_attn_mask, [0, 1])
 
         tgt_len, head_dim = query.size(-3), query.size(-1)
-        assert (
-            query.size(-1) == key.size(-1) == value.size(-1)
-        ), "The feature dim of query, key, value must be equal."
+        assert query.size(-1) == key.size(-1) == value.size(-1), (
+            "The feature dim of query, key, value must be equal."
+        )
         assert key.size() == value.size(), "Shape of key, value must match"
         src_len = key.size(-3)
         batch_heads = max(query.size(-2), key.size(-2))

benchmarks/functional_autograd_benchmark/torchvision_models.py

@@ -884,9 +884,9 @@ def __init__(
         self.cost_class = cost_class
         self.cost_bbox = cost_bbox
         self.cost_giou = cost_giou
-        assert (
-            cost_class != 0 or cost_bbox != 0 or cost_giou != 0
-        ), "all costs cant be 0"
+        assert cost_class != 0 or cost_bbox != 0 or cost_giou != 0, (
+            "all costs cant be 0"
+        )
 
     @torch.no_grad()
     def forward(self, outputs, targets):

benchmarks/gpt_fast/model.py

@@ -51,9 +51,9 @@ def from_name(cls, name: str):
         # take longer name (as it have more symbols matched)
         if len(config) > 1:
             config.sort(key=len, reverse=True)
-            assert len(config[0]) != len(
-                config[1]
-            ), name  # make sure only one 'best' match
+            assert len(config[0]) != len(config[1]), (
+                name
+            )  # make sure only one 'best' match
 
         return cls(**transformer_configs[config[0]])
 

benchmarks/instruction_counts/core/expand.py

@@ -80,9 +80,9 @@ def _generate_torchscript_file(model_src: str, name: str) -> Optional[str]:
 
     # And again, the type checker has no way of knowing that this line is valid.
     jit_model = module.jit_model  # type: ignore[attr-defined]
-    assert isinstance(
-        jit_model, (torch.jit.ScriptFunction, torch.jit.ScriptModule)
-    ), f"Expected ScriptFunction or ScriptModule, got: {type(jit_model)}"
+    assert isinstance(jit_model, (torch.jit.ScriptFunction, torch.jit.ScriptModule)), (
+        f"Expected ScriptFunction or ScriptModule, got: {type(jit_model)}"
+    )
     jit_model.save(artifact_path)  # type: ignore[call-arg]
 
     # Cleanup now that we have the actual serialized model.

benchmarks/operator_benchmark/benchmark_core.py

@@ -276,9 +276,9 @@ def split(s):
                 if c in open_to_close.keys():
                     curr_brackets.append(c)
                 elif c in open_to_close.values():
-                    assert (
-                        curr_brackets and open_to_close[curr_brackets[-1]] == c
-                    ), "ERROR: not able to parse the string!"
+                    assert curr_brackets and open_to_close[curr_brackets[-1]] == c, (
+                        "ERROR: not able to parse the string!"
+                    )
                     curr_brackets.pop()
                 elif c == "," and (not curr_brackets):
                     break_idxs.append(i)

benchmarks/sparse/triton_ops.py

@@ -3,9 +3,9 @@
 
 
 def create_blocked_tensor(B, M, N, blocksize, sparsity, dtype, device):
-    assert (
-        sparsity <= 1.0 and sparsity >= 0.0
-    ), "sparsity should be a value between 0 and 1"
+    assert sparsity <= 1.0 and sparsity >= 0.0, (
+        "sparsity should be a value between 0 and 1"
+    )
     assert M % blocksize[0] == 0
     assert N % blocksize[1] == 0
     shape = (B, M // blocksize[0], N // blocksize[1])[int(B == 0) :]

benchmarks/transformer/attention_bias_benchmarks.py

@@ -84,9 +84,9 @@ def __init__(self, num_heads, embed_dim, device=None, dtype=None):
 
         self.head_dim = embed_dim // num_heads
         self.embed_dim = embed_dim
-        assert (
-            self.head_dim * num_heads == self.embed_dim
-        ), "embed_dim must be divisible by num_heads"
+        assert self.head_dim * num_heads == self.embed_dim, (
+            "embed_dim must be divisible by num_heads"
+        )
 
         self.q_proj_weight = Parameter(
             torch.empty((embed_dim, embed_dim), **factory_kwargs)

benchmarks/transformer/score_mod.py

@@ -49,9 +49,9 @@ class ExperimentConfig:
     backends: list[str]
 
     def __post_init__(self):
-        assert (
-            len(self.shape) == 6
-        ), "Shape must be of length 6"  # [B, Hq, M, Hkv, N, D]
+        assert len(self.shape) == 6, (
+            "Shape must be of length 6"
+        )  # [B, Hq, M, Hkv, N, D]
 
     def asdict(self):
         # Convert the dataclass instance to a dictionary

functorch/dim/reference.py

@@ -625,9 +625,9 @@ def split(self, split_size_or_sections, dim=0):
             unbound.append(i)
 
     if unbound:
-        assert (
-            total_bound_size <= size
-        ), f"result dimensions are larger than original: {total_bound_size} vs {size} ({split_size_or_sections})"
+        assert total_bound_size <= size, (
+            f"result dimensions are larger than original: {total_bound_size} vs {size} ({split_size_or_sections})"
+        )
         remaining_size = size - total_bound_size
         chunk_size = -(-remaining_size // len(unbound))
         for u in unbound:
@@ -636,9 +636,9 @@ def split(self, split_size_or_sections, dim=0):
             sizes[u] = sz
             remaining_size -= sz
     else:
-        assert (
-            total_bound_size == size
-        ), f"result dimensions do not match original: {total_bound_size} vs {size} ({split_size_or_sections})"
+        assert total_bound_size == size, (
+            f"result dimensions do not match original: {total_bound_size} vs {size} ({split_size_or_sections})"
+        )
     return tuple(
         t.index(dim, d)
         for d, t in zip(split_size_or_sections, _orig_split(self, sizes, dim=dim))

scripts/compile_tests/download_reports.py

@@ -62,9 +62,9 @@ def subdir_path(config):
     for config in configs:
         required_jobs.extend(list(CONFIGS[config]))
     for job in required_jobs:
-        assert (
-            job in workflow_jobs
-        ), f"{job} not found, is the commit_sha correct? has the job finished running? The GitHub API may take a couple minutes to update."
+        assert job in workflow_jobs, (
+            f"{job} not found, is the commit_sha correct? has the job finished running? The GitHub API may take a couple minutes to update."
+        )
 
     # This page lists all artifacts.
     listings = requests.get(

scripts/export/update_schema.py

@@ -23,9 +23,9 @@
     )
     args = parser.parse_args()
 
-    assert os.path.exists(
-        args.prefix
-    ), f"Assuming path {args.prefix} is the root of pytorch directory, but it doesn't exist."
+    assert os.path.exists(args.prefix), (
+        f"Assuming path {args.prefix} is the root of pytorch directory, but it doesn't exist."
+    )
 
     commit = schema_check.update_schema()
 
@@ -40,7 +40,9 @@
                 f"Treespec version downgraded from {commit.base['TREESPEC_VERSION']} to {commit.result['TREESPEC_VERSION']}."
             )
     else:
-        assert args.force_unsafe, "Existing schema yaml file not found, please use --force-unsafe to try again."
+        assert args.force_unsafe, (
+            "Existing schema yaml file not found, please use --force-unsafe to try again."
+        )
 
     next_version, reason = schema_check.check(commit, args.force_unsafe)
 

test/onnx/dynamo/test_dynamo_with_onnxruntime_backend.py

@@ -182,9 +182,9 @@ def _test_model_numerically(
                             baseline_param.grad, param.grad, atol=atol, rtol=rtol
                         )
             else:
-                assert (
-                    test_backward is False
-                ), "Calculating backward with multiple outputs is not supported yet."
+                assert test_backward is False, (
+                    "Calculating backward with multiple outputs is not supported yet."
+                )
                 for baseline_elem, result_elem in zip(baseline_result, result):
                     torch.testing.assert_close(
                         baseline_elem, result_elem, atol=atol, rtol=rtol

test/onnx/exporter/test_hf_models_e2e.py

@@ -98,15 +98,13 @@ def test_onnx_export_with_custom_axis_names_in_dynamic_shapes(self):
             self.assertEqual(dim.value, custom_name)
 
 
-def _prepare_llm_model_gptj_to_test() -> (
-    tuple[
-        torch.nn.Module,
-        dict[str, Any],
-        dict[str, dict[int, str]],
-        list[str],
-        list[str],
-    ]
-):
+def _prepare_llm_model_gptj_to_test() -> tuple[
+    torch.nn.Module,
+    dict[str, Any],
+    dict[str, dict[int, str]],
+    list[str],
+    list[str],
+]:
     model = transformers.GPTJForCausalLM.from_pretrained(
         "hf-internal-testing/tiny-random-gptj"
     )