moving sharding constraint to outside scan loop

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1247,13 +1247,6 @@ def step_fn(carry, t):
         video_embeds_sharded = video_embeds
         audio_embeds_sharded = audio_embeds
 
-        if not self.transformer.scan_layers:
-          activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
-          latents_jax_sharded = jax.lax.with_sharding_constraint(latents_jax, activation_axis_names)
-          audio_latents_jax_sharded = jax.lax.with_sharding_constraint(audio_latents_jax, activation_axis_names)
-          video_embeds_sharded = jax.lax.with_sharding_constraint(video_embeds, activation_axis_names)
-          audio_embeds_sharded = jax.lax.with_sharding_constraint(audio_embeds, activation_axis_names)
-
         noise_pred, noise_pred_audio = transformer_forward_pass(
             graphdef,
             state,
@@ -1301,6 +1294,13 @@ def step_fn(carry, t):
 
         return (new_latents_jax.astype(latents_jax.dtype), new_audio_latents_jax.astype(audio_latents_jax.dtype)), None
 
+      if not self.transformer.scan_layers:
+        activation_axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
+        latents_jax = jax.lax.with_sharding_constraint(latents_jax, activation_axis_names)
+        audio_latents_jax = jax.lax.with_sharding_constraint(audio_latents_jax, activation_axis_names)
+        video_embeds = jax.lax.with_sharding_constraint(video_embeds, activation_axis_names)
+        audio_embeds = jax.lax.with_sharding_constraint(audio_embeds, activation_axis_names)
+
       initial_carry = (latents_jax, audio_latents_jax)
       with jax.named_scope("denoising_loop"):
         (latents_jax, audio_latents_jax), _ = jax.lax.scan(step_fn, initial_carry, timesteps_jax)