pyink checks

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -384,7 +384,9 @@ def ring_scan_body(carry, _):
           return (m, l, o, k_next, v_next), None
 
         initial_carry = (m, l, o, k1, v1)
-        (m_final, l_final, o_final, _, _), _ = jax.lax.scan(ring_scan_body, initial_carry, None, length=num_context_shards - 1)
+        (m_final, l_final, o_final, _, _), _ = jax.lax.scan(
+            ring_scan_body, initial_carry, None, length=num_context_shards - 1
+        )
 
         attention_output = o_final / l_final[..., None]
       else:
@@ -749,6 +751,7 @@ def __init__(
     self.dpa_layer = None
     if attention_kernel == "cudnn_flash_te":
       from transformer_engine.jax.flax.transformer import DotProductAttention  # pytype: disable=import-error
+
       jax.config.update("jax_use_shardy_partitioner", False)
 
       dpa_layer = DotProductAttention(
@@ -829,6 +832,7 @@ def setup(self):
     self.dpa_layer = None
     if self.attention_kernel == "cudnn_flash_te":
       from transformer_engine.jax.flax.transformer import DotProductAttention  # pytype: disable=import-error
+
       jax.config.update("jax_use_shardy_partitioner", False)
 
       dpa_layer = DotProductAttention(
@@ -848,7 +852,6 @@ def setup(self):
       variables = {}
       self.dpa_layer = functools.partial(dpa_layer.apply, variables)
 
-
   def apply_attention(self, query: Array, key: Array, value: Array, attention_mask: Array = None):
     return _apply_attention(
         query=query,

src/maxdiffusion/models/wan/autoencoder_kl_wan.py

@@ -29,10 +29,11 @@
 
 CACHE_T = 2
 try:
-  flax.config.update('flax_always_shard_variable', False)
+  flax.config.update("flax_always_shard_variable", False)
 except LookupError:
   pass
 
+
 # Helper to ensure kernel_size, stride, padding are tuples of 3 integers
 def _canonicalize_tuple(x: Union[int, Sequence[int]], rank: int, name: str) -> Tuple[int, ...]:
   """Canonicalizes a value to a tuple of integers."""

src/maxdiffusion/pipelines/wan/wan_pipeline_2_1.py

@@ -118,7 +118,8 @@ def __call__(
     )
     # Set the TE shard_guard context_manager if using TE cudnn_flash attention
     if self.config.attention == "cudnn_flash_te":
-      from transformer_engine.jax.sharding import global_shard_guard, MeshResource # pytype: disable=import-error
+      from transformer_engine.jax.sharding import global_shard_guard, MeshResource  # pytype: disable=import-error
+
       shard_guard = global_shard_guard(MeshResource(cp_resource="context"))
     else:
       shard_guard = nullcontext()

src/maxdiffusion/pipelines/wan/wan_pipeline_2_2.py

@@ -140,7 +140,8 @@ def __call__(
     )
     # Set the TE shard_guard context_manager if using TE cudnn_flash attention
     if self.config.attention == "cudnn_flash_te":
-      from transformer_engine.jax.sharding import global_shard_guard, MeshResource # pytype: disable=import-error
+      from transformer_engine.jax.sharding import global_shard_guard, MeshResource  # pytype: disable=import-error
+
       shard_guard = global_shard_guard(MeshResource(cp_resource="context"))
     else:
       shard_guard = nullcontext()

src/maxdiffusion/pyconfig.py

@@ -201,7 +201,9 @@ def user_init(raw_keys):
     raw_keys["logical_axis_rules"] = _lists_to_tuples(raw_keys["logical_axis_rules"])
     # Verify qkv is sharded across sequence.
     if raw_keys["attention"] == "ring" or raw_keys["attention_sharding_uniform"]:
-      max_logging.log(f"Adding sequence sharding to q and kv if not already present because {raw_keys['attention']}=='ring' or {raw_keys['attention_sharding_uniform']} is set.")
+      max_logging.log(
+          f"Adding sequence sharding to q and kv if not already present because {raw_keys['attention']}=='ring' or {raw_keys['attention_sharding_uniform']} is set."
+      )
       logical_axis_rules = list(raw_keys["logical_axis_rules"])
       max_logging.log(f"Initial logical axis rules: {logical_axis_rules}")
       new_rules = []

src/maxdiffusion/tests/wan_transformer_test.py

@@ -125,9 +125,7 @@ def test_wan_time_text_embedding(self):
 
       encoder_hidden_states_shape = (batch_size, time_freq_dim * 2, text_embed_dim)
       dummy_encoder_hidden_states = jnp.ones(encoder_hidden_states_shape)
-      temb, timestep_proj, encoder_hidden_states, _, _ = layer(
-          dummy_timestep, dummy_encoder_hidden_states
-      )
+      temb, timestep_proj, encoder_hidden_states, _, _ = layer(dummy_timestep, dummy_encoder_hidden_states)
       assert temb.shape == (batch_size, dim)
       assert timestep_proj.shape == (batch_size, time_proj_dim)
       assert encoder_hidden_states.shape == (batch_size, time_freq_dim * 2, dim)

src/maxdiffusion/trainers/wan_trainer.py

@@ -309,7 +309,7 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
         pretty_string = pprint.pformat(state_spec.opt_state, indent=4, width=60)
         max_logging.log(pretty_string)
         max_logging.log("------------------------------------------------")
-    if self.config.hardware != 'gpu':
+    if self.config.hardware != "gpu":
       max_utils.delete_pytree(params)
     data_shardings = self.get_data_shardings(mesh)
     eval_data_shardings = self.get_eval_data_shardings(mesh)
@@ -368,14 +368,16 @@ def training_loop(self, pipeline, optimizer, learning_rate_scheduler, train_data
 
         # Designate the context parallel axis for sharding
         if self.config.attention == "cudnn_flash_te":
-          from transformer_engine.jax.sharding import global_shard_guard, MeshResource # pytype: disable=import-error
+          from transformer_engine.jax.sharding import global_shard_guard, MeshResource  # pytype: disable=import-error
+
           shard_guard = global_shard_guard(MeshResource(cp_resource="context"))
         else:
           shard_guard = nullcontext()
 
         next_batch_future = executor.submit(load_next_batch, train_data_iterator, example_batch, self.config)
-        with jax.profiler.StepTraceAnnotation("train", step_num=step), pipeline.mesh, \
-        shard_guard, nn_partitioning.axis_rules(self.config.logical_axis_rules):
+        with jax.profiler.StepTraceAnnotation(
+            "train", step_num=step
+        ), pipeline.mesh, shard_guard, nn_partitioning.axis_rules(self.config.logical_axis_rules):
           state, scheduler_state, train_metric, rng = p_train_step(state, example_batch, rng, scheduler_state)
           train_metric["scalar"]["learning/loss"].block_until_ready()
         last_step_completion = datetime.datetime.now()