fix list-type prompt embeds shape error

Author: prishajain1

src/maxdiffusion/pipelines/ltx2/ltx2_pipeline.py

@@ -1399,13 +1399,25 @@ def __call__(
       latents_jax = jnp.concatenate([latents_jax] * 2, axis=0)
       audio_latents_jax = jnp.concatenate([audio_latents_jax] * 2, axis=0)
     def _print_stats_gemma(name, tensor):
-      print(
-          f"DEBUG {name} shape: {tensor.shape}, mean: {jnp.round(jnp.mean(tensor), 6)}, min: {jnp.round(jnp.min(tensor), 4)}, max: {jnp.round(jnp.max(tensor), 4)}, std: {jnp.round(jnp.std(tensor), 4)}"
-      )
+      if isinstance(tensor, list):
+        for idx, t in enumerate(tensor):
+          print(
+              f"DEBUG {name}_{idx} shape: {t.shape}, mean: {jnp.round(jnp.mean(t), 6)}, min: {jnp.round(jnp.min(t), 4)}, max: {jnp.round(jnp.max(t), 4)}, std: {jnp.round(jnp.std(t), 4)}"
+          )
+      else:
+        print(
+            f"DEBUG {name} shape: {tensor.shape}, mean: {jnp.round(jnp.mean(tensor), 6)}, min: {jnp.round(jnp.min(tensor), 4)}, max: {jnp.round(jnp.max(tensor), 4)}, std: {jnp.round(jnp.std(tensor), 4)}"
+        )
+
     if do_cfg and do_stg:
-      _print_stats_gemma("text_encoder_output_flattened", prompt_embeds_jax[:2])
+      if isinstance(prompt_embeds_jax, list):
+        prompt_embeds_to_print = [x[:2] for x in prompt_embeds_jax]
+      else:
+        prompt_embeds_to_print = prompt_embeds_jax[:2]
     else:
-      _print_stats_gemma("text_encoder_output_flattened", prompt_embeds_jax)
+      prompt_embeds_to_print = prompt_embeds_jax
+
+    _print_stats_gemma("text_encoder_output_flattened", prompt_embeds_to_print)
 
     if hasattr(self, "mesh") and self.mesh is not None:
       data_sharding_3d = NamedSharding(self.mesh, P())