dim changed in test

Author: prishajain1

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -68,21 +68,21 @@ def test_transformer_block_shapes(self):
         """
         print("\n=== Testing LTX2VideoTransformerBlock Shapes ===")
 
-        dim = 32
-        audio_dim = 16
-        cross_dim = 20 # context dim
+        dim = 1024
+        audio_dim = 1024
+        cross_dim = 64 # context dim
 
         # NNX sharding context
         with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
             block = LTX2VideoTransformerBlock(
                 rngs=self.rngs,
                 dim=dim,
                 num_attention_heads=8,
-                attention_head_dim=4,
+                attention_head_dim=128,
                 cross_attention_dim=cross_dim,
                 audio_dim=audio_dim,
                 audio_num_attention_heads=8,
-                audio_attention_head_dim=2,
+                audio_attention_head_dim=128,
                 audio_cross_attention_dim=cross_dim,
                 activation_fn="gelu",
                 qk_norm="rms_norm_across_heads",
@@ -181,14 +181,14 @@ def test_transformer_3d_model_instantiation_and_forward(self):
                 patch_size=self.patch_size,
                 patch_size_t=self.patch_size_t,
                 num_attention_heads=8,
-                attention_head_dim=4,
+                attention_head_dim=128,
                 num_layers=1, # 1 layer for speed
                 caption_channels=32, # small for test
                 cross_attention_dim=32,
                 audio_in_channels=self.audio_in_channels,
                 audio_out_channels= self.audio_in_channels,
                 audio_num_attention_heads=8,
-                audio_attention_head_dim=2,
+                audio_attention_head_dim=128,
                 audio_cross_attention_dim=32,
                 mesh=self.mesh,
             )