Run pyink on files that landed unformatted (PRs #436, #438)

src/maxdiffusion/generate_wan.py

@@ -302,10 +302,7 @@ def run(config, pipeline=None, filename_prefix="", commit_hash=None):
   videos = call_pipeline(config, pipeline, prompt, negative_prompt, num_inference_steps=warmup_steps)
   if isinstance(videos, tuple):
     videos, warmup_trace = videos
-    max_logging.log(
-        "Warmup breakdown: "
-        + ", ".join(f"{stage}={seconds:.1f}s" for stage, seconds in warmup_trace.items())
-    )
+    max_logging.log("Warmup breakdown: " + ", ".join(f"{stage}={seconds:.1f}s" for stage, seconds in warmup_trace.items()))
 
   max_logging.log("===================== Model details =======================")
   max_logging.log(f"model name: {config.model_name}")

src/maxdiffusion/models/attention_flax.py

@@ -1113,9 +1113,7 @@ def wrap_ulysses_ring_attention(query, key, value):
           use_fixed_m=use_fixed_m,
       )
       if use_fixed_m:
-        attention_output = jnp.swapaxes(
-            jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr), 2, 3
-        )
+        attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0, None))(query, key, value, mk_arr), 2, 3)
       else:
         attention_output = jnp.swapaxes(jax.vmap(splash_kernel, in_axes=(0, 0, 0))(query, key, value), 2, 3)
     else:

src/maxdiffusion/models/wan/wan_utils.py

@@ -383,9 +383,7 @@ def convert_chunk(ckpt_shard_path, chunk_keys):
 
         # rename_key_and_reshape_tensor only reindexes/transposes views; the
         # single real copy happens on assignment into the target buffer below.
-        flax_key, flax_tensor = rename_key_and_reshape_tensor(
-            pt_tuple_key, tensor, random_flax_state_dict, scan_layers
-        )
+        flax_key, flax_tensor = rename_key_and_reshape_tensor(pt_tuple_key, tensor, random_flax_state_dict, scan_layers)
         flax_key = rename_for_nnx(flax_key)
         flax_key = _tuple_str_to_int(flax_key)
 
@@ -428,9 +426,7 @@ def convert_chunk(ckpt_shard_path, chunk_keys):
 
   validate_flax_state_dict(eval_shapes, flax_state_dict)
   flax_state_dict = unflatten_dict(flax_state_dict)
-  max_logging.log(
-      f"Converted {subfolder or 'transformer'} weights to host arrays in {time.perf_counter() - t_start:.1f}s"
-  )
+  max_logging.log(f"Converted {subfolder or 'transformer'} weights to host arrays in {time.perf_counter() - t_start:.1f}s")
   return flax_state_dict
 
 

src/maxdiffusion/pipelines/wan/wan_pipeline.py

@@ -279,15 +279,11 @@ def stages_via_ici(val, sharding) -> bool:
 
     put_arrays = [None] * len(put_specs)
     if staged_ids:
-      staged = jax.device_put(
-          [put_specs[i][1] for i in staged_ids], [dim0_sharding] * len(staged_ids)
-      )
+      staged = jax.device_put([put_specs[i][1] for i in staged_ids], [dim0_sharding] * len(staged_ids))
       # out_shardings must be the exact target sharding objects (not an
       # equivalent P()): downstream jit cache keys include arg shardings, so
       # a different-but-equivalent spec would force a full recompile.
-      replicate_fn = jax.jit(
-          lambda xs: xs, out_shardings=[put_specs[i][2] for i in staged_ids]
-      )
+      replicate_fn = jax.jit(lambda xs: xs, out_shardings=[put_specs[i][2] for i in staged_ids])
       for i, replicated in zip(staged_ids, replicate_fn(staged)):
         put_arrays[i] = replicated
     if direct_ids:

src/maxdiffusion/tests/custom_splash_fixed_m_test.py

@@ -47,13 +47,9 @@ class CustomSplashFixedMTest(unittest.TestCase):
   def setUp(self):
     super().setUp()
     self.scale = 1.0 / math.sqrt(self.head_dim)
-    self.block_sizes = custom_splash._BlockSizes(
-        block_q=2048, block_kv=1024, block_kv_compute=512
-    )
+    self.block_sizes = custom_splash._BlockSizes(block_q=2048, block_kv=1024, block_kv_compute=512)
 
-  def _random_qkv(
-      self, q_gain: float = 1.0, k_gain: float = 1.0
-  ) -> tuple[jax.Array, jax.Array, jax.Array]:
+  def _random_qkv(self, q_gain: float = 1.0, k_gain: float = 1.0) -> tuple[jax.Array, jax.Array, jax.Array]:
     """Returns bf16 (q, k, v), optionally amplifying head 0 of q and k."""
     shape = (self.num_heads, self.seq_len, self.head_dim)
     q = jax.random.normal(jax.random.PRNGKey(0), shape, jnp.bfloat16)
@@ -63,18 +59,14 @@ def _random_qkv(
     k = k.at[0].multiply(k_gain)
     return q, k, v
 
-  def _reference(
-      self, q: jax.Array, k: jax.Array, v: jax.Array
-  ) -> jax.Array:
+  def _reference(self, q: jax.Array, k: jax.Array, v: jax.Array) -> jax.Array:
     """Per-head f32 softmax attention reference."""
     qf, kf, vf = (x.astype(jnp.float32) for x in (q, k, v))
     logits = jnp.einsum("hsd,htd->hst", qf, kf) * self.scale
     probs = jax.nn.softmax(logits, axis=-1)
     return jnp.einsum("hst,htd->hsd", probs, vf)
 
-  def _run_kernel(
-      self, q: jax.Array, k: jax.Array, v: jax.Array, use_fixed_m: bool
-  ) -> tuple[jax.Array, jax.Array | None]:
+  def _run_kernel(self, q: jax.Array, k: jax.Array, v: jax.Array, use_fixed_m: bool) -> tuple[jax.Array, jax.Array | None]:
     """Runs the custom kernel using the production scaling convention.
 
     Args:
@@ -95,9 +87,7 @@ def _run_kernel(
       k_in = k_in - jnp.mean(k_in, axis=1, keepdims=True)
       qn = jnp.sqrt((q_in.astype(jnp.float32) ** 2).sum(-1)).max(axis=1)
       mk_h = jnp.sqrt((k_in.astype(jnp.float32) ** 2).sum(-1)).max(axis=1)
-      eligible = (qn * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(
-          jnp.float32
-      )
+      eligible = (qn * mk_h <= custom_splash._FIXED_M_SAFE_BOUND).astype(jnp.float32)
       mk = jnp.stack([mk_h, eligible])
     kernel = custom_splash.make_splash_mha(
         block_sizes=self.block_sizes,