[pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

Author: pre-commit-ci[bot]

thunder/clang/__init__.py

@@ -353,7 +353,9 @@ def tensor_from_sequence(
 def diagonal(a: TensorLike, offset: int = 0, dim1: int = 0, dim2: int = 1) -> TensorLike:
     utils.check(
         a.ndim >= 2,
-        lambda: f"diagonal() expected a tensor with at least two dimensions, but got a tensor with {a.ndims} dimensions",
+        lambda: (
+            f"diagonal() expected a tensor with at least two dimensions, but got a tensor with {a.ndims} dimensions"
+        ),
     )
 
     diag_length = max(0, min(a.shape[dim1] + min(offset, 0), a.shape[dim2] - max(offset, 0)))

thunder/core/dtypes.py

@@ -334,15 +334,19 @@ def resolve_dtypes(args: Iterable) -> set[dtype]:
             for a in arg:
                 baseutils.check(
                     isinstance(a, dtype),
-                    lambda: f"Iterables passed to resolve_dtypes must only contain dtypes, but found an Iterable with {a}",
+                    lambda: (
+                        f"Iterables passed to resolve_dtypes must only contain dtypes, but found an Iterable with {a}"
+                    ),
                     exception_type=NotImplementedError,
                 )
                 if not a.is_weak:
                     dtypes.add(a)
 
         baseutils.check(
             arg in (dtype, exact, signedinteger, unsignedinteger, bool_, inexact, floating, complexfloating),
-            lambda: f"Excepted arguments to resolve_dtypes to be dtypes, sets of dtypes, or a dtype (sub)class, but got {arg}",
+            lambda: (
+                f"Excepted arguments to resolve_dtypes to be dtypes, sets of dtypes, or a dtype (sub)class, but got {arg}"
+            ),
             exception_type=AssertionError,
         )
 

thunder/core/prims.py

@@ -917,8 +917,10 @@ def _index_copy_grad(a: TensorProxy, /, index: TensorProxy, src: TensorProxy, di
 def _scatter_add_prim_grad(a: TensorProxy, /, index: TensorProxy, value: TensorProxy, dim: int) -> TensorProxy:
     utils.check(
         not value._requires_grad or value.shape == index.shape,
-        lambda: "The gradient for the value Tensor is implemented only when value.shape == index.shape. "
-        "value shape is {value.shape} while index shape is {index.shape}",
+        lambda: (
+            "The gradient for the value Tensor is implemented only when value.shape == index.shape. "
+            "value shape is {value.shape} while index shape is {index.shape}"
+        ),
     )
 
     fwd = prims.scatter_add(a, index, value, dim)
@@ -2985,8 +2987,10 @@ def put_grad(v: Variable, val: Any) -> None:
         if len(symbol.args) != (orig_res_len := len(result)):
             check(
                 orig_res_len <= len(symbol.args),
-                lambda: f"Backward for {symbol.sym.id} returned {orig_res_len} values, "
-                + f"but expected at most {len(symbol.args)}",
+                lambda: (
+                    f"Backward for {symbol.sym.id} returned {orig_res_len} values, "
+                    + f"but expected at most {len(symbol.args)}"
+                ),
             )
             # Assuming that the non-differentiable arguments were dropped from
             # the backward function, we are going to append None to the result
@@ -2998,8 +3002,10 @@ def put_grad(v: Variable, val: Any) -> None:
             n_differentiable_args = sum(bool(_is_differentiable(arg)) for arg in symbol.args)
             check(
                 n_differentiable_args <= orig_res_len,
-                lambda: f"Backward for {symbol.sym.id} returned {orig_res_len} value(s), "
-                + f"but expected {n_differentiable_args}",
+                lambda: (
+                    f"Backward for {symbol.sym.id} returned {orig_res_len} value(s), "
+                    + f"but expected {n_differentiable_args}"
+                ),
             )
 
             result = tuple(next(iter_result) if _is_differentiable(arg) else None for arg in symbol.args)
@@ -3033,7 +3039,9 @@ def vjp_call(primals, cotangents, trace: Trace, **kwargs):
     result, env = augmented_forward_pass(*primals, trace=trace, **kwargs)
     check(
         len(result) == len(cotangents) if isinstance(result, Sequence) else True,
-        lambda: f"Expected cotangents to be a sequence of length {len(result)}, got a sequence of length {len(cotangents)}",
+        lambda: (
+            f"Expected cotangents to be a sequence of length {len(result)}, got a sequence of length {len(cotangents)}"
+        ),
     )
     return result, backward_pass(env, trace, cotangents)
 

thunder/core/transforms.py

@@ -141,11 +141,13 @@ def find_backward_output(forward_input):
     # same as the number of primal outputs of the augmented forward trace
     utils.check(
         len(utils.sequencify(bsym.output)) == len(utils.sequencify(augmented_forward_trace.output[0])),
-        lambda: f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
-        "The number of outputs of the original forward function must be the same as the number of primal outputs of the augmented forward trace.\n"
-        f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
-        f"Number of primal outputs of the augmented forward trace: {len(utils.sequencify(augmented_forward_trace.output[0]))}\n"
-        "Please check the forward function and the augmented forward trace to ensure that they have the same number of outputs.",
+        lambda: (
+            f"While generating forward and backward functions for {bsym.sym.name}, encountered an error.\n"
+            "The number of outputs of the original forward function must be the same as the number of primal outputs of the augmented forward trace.\n"
+            f"Number of outputs of the original forward function: {len(utils.sequencify(bsym.output))}\n"
+            f"Number of primal outputs of the augmented forward trace: {len(utils.sequencify(augmented_forward_trace.output[0]))}\n"
+            "Please check the forward function and the augmented forward trace to ensure that they have the same number of outputs."
+        ),
     )
 
     # Check if any of the bound symbols in the backward trace are also in the

thunder/core/vjp_utils.py

@@ -67,7 +67,9 @@ def remove_grad_sync(backward_trace_with_grad_sync: TraceCtx) -> TraceCtx:
         bsym_of_preaveraging: BoundSymbol = producers[bsym_of_allreduce.flat_proxy_args[0]]
         utils.check(
             bsym_of_preaveraging.sym.id in {PrimIDs.DIV, "torch.true_divide"},
-            lambda: f"expected to be either of {(PrimIDs.DIV, 'torch.true_divide')} but {bsym_of_preaveraging.sym.id=} for {synced_grad=}",
+            lambda: (
+                f"expected to be either of {(PrimIDs.DIV, 'torch.true_divide')} but {bsym_of_preaveraging.sym.id=} for {synced_grad=}"
+            ),
         )
         bsym_to_remove.extend([bsym_of_allreduce, bsym_of_wait, bsym_of_preaveraging])
         synced_to_unsynced[variableify(synced_grad)] = bsym_of_preaveraging.flat_proxy_args[0]
@@ -172,7 +174,9 @@ def __call__(self, bsym: BoundSymbol) -> None:
         if grads_of_bsym:
             utils.check(
                 bsym.sym.id in {PrimIDs.DIV, "torch.true_divide"},
-                lambda: f"This bsym's sym.id is expected to be {PrimIDs.DIV=} or 'torch.true_divide' but {bsym.sym.id=}",
+                lambda: (
+                    f"This bsym's sym.id is expected to be {PrimIDs.DIV=} or 'torch.true_divide' but {bsym.sym.id=}"
+                ),
             )
             utils.check(len(grads_of_bsym) == 1, lambda: f"{len(grads_of_bsym)=} is expected to be 1")
             self.gradient_buckets.tell(grads_of_bsym[0], self.process_group)

thunder/distributed/transforms/ddp.py

@@ -149,7 +149,9 @@ def stash_unsharded_grads_and_return_none_as_grads(
             prod_of_unsharded_grad = producers[preaverage_bsym.flat_proxy_args[0]]
             utils.check(
                 preaverage_bsym.sym.id in {PrimIDs.DIV, "torch.true_divide"},
-                lambda: f"expected to be either of {(PrimIDs.DIV, 'torch.true_divide')} but {preaverage_bsym.sym.id=} for {output_proxy=}",
+                lambda: (
+                    f"expected to be either of {(PrimIDs.DIV, 'torch.true_divide')} but {preaverage_bsym.sym.id=} for {output_proxy=}"
+                ),
             )
             wait_bsym = consumers[reduce_scatter_bsym.flat_proxy_outs[0]][0]
 
@@ -221,7 +223,9 @@ def __post_init__(self) -> None:
             case _:
                 utils.check(
                     False,
-                    lambda: f"Invalid {self.comm_to_bucket}, {(dist_prims.PrimIDs.ALL_GATHER, dist_prims.PrimIDs.REDUCE_SCATTER)} are supported",
+                    lambda: (
+                        f"Invalid {self.comm_to_bucket}, {(dist_prims.PrimIDs.ALL_GATHER, dist_prims.PrimIDs.REDUCE_SCATTER)} are supported"
+                    ),
                 )
 
     @property
@@ -265,7 +269,9 @@ def maybe_swap_proxies_of_bsym_and_update_swap_map(bsym: BoundSymbol) -> bool:
             param = bsym.flat_proxy_args[0]
             utils.check(
                 param in self.params,
-                lambda: f"{variableify(param)} not found in param set: {(variableify(p) for p in self.original_params)}",
+                lambda: (
+                    f"{variableify(param)} not found in param set: {(variableify(p) for p in self.original_params)}"
+                ),
             )
             if param not in self.param_to_bucket:
                 # This path is highly likely to be backward reduce-scatter bucketing:
@@ -458,7 +464,9 @@ def __init__(
             hasattr(compile_data.fn, "process_group_for_ddp")
             and hasattr(compile_data.fn, "bucketing_strategy")
             and hasattr(compile_data.fn, "sharding_strategy"),
-            lambda: f"Given module does not seem to have all the attributes of `process_group_for_ddp`, `bucketing_strategy`, and `sharding_strategy`, {hasattr(compile_data.fn, 'bucketing_strategy')=}, {hasattr(compile_data.fn, 'sharding_strategy')=}",
+            lambda: (
+                f"Given module does not seem to have all the attributes of `process_group_for_ddp`, `bucketing_strategy`, and `sharding_strategy`, {hasattr(compile_data.fn, 'bucketing_strategy')=}, {hasattr(compile_data.fn, 'sharding_strategy')=}"
+            ),
         )
         self.bucketing_strategy: FSDPBucketingStrategy = compile_data.fn.bucketing_strategy
         self.apply_bucketing = self.bucketing_strategy != FSDPBucketingStrategy.NONE
@@ -539,8 +547,10 @@ def apply_bucketing_to_forward_trace(self, fwd_trace: TraceCtx) -> TraceCtx:
 
         collective_comm_bsyms: tuple[BoundSymbol, ...] = tuple(
             filter(
-                lambda bsym: bsym.sym.id == dist_prims.PrimIDs.ALL_GATHER
-                and any(arg in arg_to_index_in_flat_args for arg in bsym.flat_proxy_args),
+                lambda bsym: (
+                    bsym.sym.id == dist_prims.PrimIDs.ALL_GATHER
+                    and any(arg in arg_to_index_in_flat_args for arg in bsym.flat_proxy_args)
+                ),
                 fsdp_fwd_trace.bound_symbols,
             )
         )

thunder/distributed/transforms/fsdp.py

@@ -224,7 +224,9 @@ def limit_in_flight_allgathers(
                 pack_consumer = consumers.get(bsym.flat_proxy_outs[0], None)
                 check(
                     pack_consumer is not None and len(pack_consumer) in (1, 2),
-                    lambda: f"Pack's operand {bsym.flat_proxy_outs[0]} expected to be consumed by all-gather and del: {pack_consumer}",
+                    lambda: (
+                        f"Pack's operand {bsym.flat_proxy_outs[0]} expected to be consumed by all-gather and del: {pack_consumer}"
+                    ),
                 )
                 # skip the pack operator corresponds to allgather
                 if pack_consumer[0].sym.id != all_gather_prim_impl.id:
@@ -247,7 +249,9 @@ def limit_in_flight_allgathers(
                             wait_consumer is not None
                             and len(wait_consumer) in (1, 2)
                             and wait_consumer[0].sym.id == unpack_for_fsdp_prim_impl.id,
-                            lambda: f"wait of {bsym.flat_proxy_outs[0]} expected to be consumed unpack and del: {wait_consumer}",
+                            lambda: (
+                                f"wait of {bsym.flat_proxy_outs[0]} expected to be consumed unpack and del: {wait_consumer}"
+                            ),
                         )
                         unpack_bsyms.append(wait_consumer[0])
                 bound_symbols.append(bsym)

thunder/distributed/utils.py

@@ -528,7 +528,9 @@ def try_and_log_benchmark(compile_fn, filename):
     for m1, m2, name in zip(measure1, measure2, ("forward", "backward")):
         check(
             (m1 is None) == (m2 is None),
-            lambda: f"{name} measurement for the two compilation methods should either both be None or both not None, but got {m1} and {m2}",
+            lambda: (
+                f"{name} measurement for the two compilation methods should either both be None or both not None, but got {m1} and {m2}"
+            ),
         )
         if m1 is None:
             continue

thunder/dynamo/benchmark_utils.py

@@ -78,7 +78,9 @@ def test_func(benchmark):
         check(isinstance(executors, dict) and executors, lambda: "'executors' must be a non-empty dictionary.")
         check(
             not any("-" in k for k in executors.keys()),
-            lambda: "Executor names cannot contain '-' as it conflicts with the 'benchmark-group-by' function. Please rename it using a different character.",
+            lambda: (
+                "Executor names cannot contain '-' as it conflicts with the 'benchmark-group-by' function. Please rename it using a different character."
+            ),
         )
         self.executors = executors
         self._get_debug_options(**debug_options)