annotation for denoising loop

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1302,7 +1302,8 @@ def step_fn(carry, t):
         return (new_latents_jax.astype(latents_jax.dtype), new_audio_latents_jax.astype(audio_latents_jax.dtype)), None
 
       initial_carry = (latents_jax, audio_latents_jax)
-      (latents_jax, audio_latents_jax), _ = jax.lax.scan(step_fn, initial_carry, timesteps_jax)
+      with jax.named_scope("denoising_loop"):
+        (latents_jax, audio_latents_jax), _ = jax.lax.scan(step_fn, initial_carry, timesteps_jax)