passed mesh param to test for transformer

Author: prishajain1

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -86,6 +86,7 @@ def test_transformer_block_shapes(self):
                 audio_cross_attention_dim=cross_dim,
                 activation_fn="gelu",
                 qk_norm="rms_norm_across_heads",
+                mesh=self.mesh,
             )
 
             # Create dummy inputs
@@ -188,7 +189,8 @@ def test_transformer_3d_model_instantiation_and_forward(self):
                 audio_out_channels= self.audio_in_channels,
                 audio_num_attention_heads=2,
                 audio_attention_head_dim=16,
-                audio_cross_attention_dim=32
+                audio_cross_attention_dim=32,
+                mesh=self.mesh,
             )
 
         # Inputs
@@ -283,9 +285,9 @@ def test_scan_remat_parity(self):
 
         # 1. Initialize models
         with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
-            model_scan = LTX2VideoTransformer3DModel(**args, scan_layers=True)
-            model_loop = LTX2VideoTransformer3DModel(**args, scan_layers=False)
-            model_remat = LTX2VideoTransformer3DModel(**args, scan_layers=True, remat_policy="full")
+            model_scan = LTX2VideoTransformer3DModel(**args, scan_layers=True, mesh=self.mesh)
+            model_loop = LTX2VideoTransformer3DModel(**args, scan_layers=False, mesh=self.mesh)
+            model_remat = LTX2VideoTransformer3DModel(**args, scan_layers=True, remat_policy="full", mesh=self.mesh)
 
         # 2. Sync weights (crucial for parity)
         # We can just copy params from scan to loop/remat