dimensions changed for tests

Author: prishajain1

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -77,12 +77,12 @@ def test_transformer_block_shapes(self):
             block = LTX2VideoTransformerBlock(
                 rngs=self.rngs,
                 dim=dim,
-                num_attention_heads=4,
-                attention_head_dim=8,
+                num_attention_heads=8,
+                attention_head_dim=4,
                 cross_attention_dim=cross_dim,
                 audio_dim=audio_dim,
-                audio_num_attention_heads=4,
-                audio_attention_head_dim=4,
+                audio_num_attention_heads=8,
+                audio_attention_head_dim=2,
                 audio_cross_attention_dim=cross_dim,
                 activation_fn="gelu",
                 qk_norm="rms_norm_across_heads",
@@ -180,15 +180,15 @@ def test_transformer_3d_model_instantiation_and_forward(self):
                 out_channels=self.out_channels,
                 patch_size=self.patch_size,
                 patch_size_t=self.patch_size_t,
-                num_attention_heads=2,
-                attention_head_dim=16,
+                num_attention_heads=8,
+                attention_head_dim=4,
                 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=2,
-                audio_attention_head_dim=16,
+                audio_num_attention_heads=8,
+                audio_attention_head_dim=2,
                 audio_cross_attention_dim=32,
                 mesh=self.mesh,
             )
@@ -271,15 +271,15 @@ def test_scan_remat_parity(self):
             out_channels=self.out_channels,
             patch_size=self.patch_size,
             patch_size_t=self.patch_size_t,
-            num_attention_heads=2,
-            attention_head_dim=16,
+            num_attention_heads=8,
+            attention_head_dim=4,
             num_layers=2, # Need >1 layer to test scan effectively
             caption_channels=32,
             cross_attention_dim=32,
             audio_in_channels=self.audio_in_channels,
             audio_out_channels= self.audio_in_channels,
-            audio_num_attention_heads=2,
-            audio_attention_head_dim=16,
+            audio_num_attention_heads=8,
+            audio_attention_head_dim=4,
             audio_cross_attention_dim=32
         )