debug added

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -18,6 +18,7 @@
 from functools import partial
 
 import numpy as np
+import functools
 import torch
 import jax
 import jax.numpy as jnp
@@ -1231,6 +1232,9 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
       )
 
       for i, t in enumerate(timesteps):
+        print(f"\n--- DEBUG SHARDING Step {i} Time {t} ---")
+        print_shardings(state, prefix="transformer_state.")
+
         noise_pred, noise_pred_audio = transformer_forward_pass(
             graphdef,
             state,
@@ -1367,6 +1371,27 @@ def run_connectors(graphdef, state, hidden_states, attention_mask):
         "fps",
     ),
 )
+
+def print_shardings(pytree, prefix=""):
+    flat_tree, _ = jax.tree_util.tree_flatten(pytree)
+    for i, leaf in enumerate(flat_tree):
+        if hasattr(leaf, 'sharding'):
+            print(f"{prefix}leaf_{i} sharding: {leaf.sharding}")
+        else:
+            print(f"{prefix}leaf_{i} has no sharding attribute")
+
+@partial(
+    jax.jit,
+    static_argnames=(
+        "do_classifier_free_guidance",
+        "guidance_scale",
+        "latent_num_frames",
+        "latent_height",
+        "latent_width",
+        "audio_num_frames",
+        "fps",
+    ),  
+)
 def transformer_forward_pass(
     graphdef,
     state,