fixes missing mesh issue and rope reshape error

Author: prishajain1

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -230,6 +230,7 @@ def __call__(self, coords: Array) -> Tuple[Array, Array]:
         # We want [B, 3, N] / [3] (broadcasting over B, N)
         # JAX broadcasting: last dims match? No.
         # reshape max_positions to [1, 3, 1]
+        max_positions = max_positions[:num_pos_dims]
         max_positions = max_positions.reshape(1, num_pos_dims, 1)
         grid = coords / max_positions
 

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -108,7 +108,8 @@ def __init__(
             bias=attention_bias,
             out_bias=attention_out_bias,
             eps=norm_eps,
-            dtype=dtype
+            dtype=dtype,
+            mesh=mesh
         )
 
         self.audio_norm1 = nnx.RMSNorm(audio_dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
@@ -121,7 +122,8 @@ def __init__(
             bias=attention_bias,
             out_bias=attention_out_bias,
             eps=norm_eps,
-            dtype=dtype
+            dtype=dtype,
+            mesh=mesh
         )
 
         # 2. Prompt Cross-Attention
@@ -136,7 +138,8 @@ def __init__(
             bias=attention_bias,
             out_bias=attention_out_bias,
             eps=norm_eps,
-            dtype=dtype
+            dtype=dtype,
+            mesh=mesh
         )
 
         self.audio_norm2 = nnx.RMSNorm(audio_dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
@@ -150,7 +153,8 @@ def __init__(
             bias=attention_bias,
             out_bias=attention_out_bias,
             eps=norm_eps,
-            dtype=dtype
+            dtype=dtype,
+            mesh=mesh
         )
 
         # 3. Audio-to-Video (a2v) and Video-to-Audio (v2a) Cross-Attention
@@ -165,7 +169,8 @@ def __init__(
             bias=attention_bias,
             out_bias=attention_out_bias,
             eps=norm_eps,
-            dtype=dtype
+            dtype=dtype,
+            mesh=mesh
         )
 
         self.video_to_audio_norm = nnx.RMSNorm(audio_dim, epsilon=self.norm_eps, use_scale=self.norm_elementwise_affine, rngs=rngs, dtype=dtype, param_dtype=weights_dtype)
@@ -179,7 +184,8 @@ def __init__(
             bias=attention_bias,
             out_bias=attention_out_bias,
             eps=norm_eps,
-            dtype=dtype
+            dtype=dtype,
+            mesh=mesh
         )
 
         # 4. Feed Forward