fix

Author: prishajain1

src/maxdiffusion/tests/ltx_2_transformer_test.py

@@ -92,8 +92,8 @@ def test_transformer_block_shapes(self):
             # Create dummy inputs
             hidden_states = jnp.zeros((self.batch_size, self.seq_len, dim))
             audio_hidden_states = jnp.zeros((self.batch_size, 128, audio_dim)) # 128 audio frames for TPFA
-            encoder_hidden_states = jnp.zeros((self.batch_size, 5, cross_dim))
-            audio_encoder_hidden_states = jnp.zeros((self.batch_size, 5, cross_dim)) # reusing cross_dim for audio context 
+            encoder_hidden_states = jnp.zeros((self.batch_size, 128, cross_dim)) # 128 for TPFA
+            audio_encoder_hidden_states = jnp.zeros((self.batch_size, 128, cross_dim)) # reusing cross_dim for audio context 
 
             # Dummy scale/shift/gate modulations
             # These match the shapes expected by the block internal calculation logic
@@ -225,10 +225,10 @@ def test_transformer_3d_model_instantiation_and_forward(self):
 
         timestep = jnp.array([1.0]) # (B,)
 
-        encoder_hidden_states = jnp.zeros((self.batch_size, 5, 32)) # (B, Lc, Dc)
-        audio_encoder_hidden_states = jnp.zeros((self.batch_size, 5, 32))
-        encoder_attention_mask = jnp.ones((self.batch_size, 5), dtype=jnp.float32)
-        audio_encoder_attention_mask = jnp.ones((self.batch_size, 5), dtype=jnp.float32)
+        encoder_hidden_states = jnp.zeros((self.batch_size, 128, 32)) # (B, Lc, Dc) # 128 for TPFA
+        audio_encoder_hidden_states = jnp.zeros((self.batch_size, 128, 32))
+        encoder_attention_mask = jnp.ones((self.batch_size, 128), dtype=jnp.float32)
+        audio_encoder_attention_mask = jnp.ones((self.batch_size, 128), dtype=jnp.float32)
 
         # Forward
         with self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):
@@ -303,8 +303,8 @@ def test_scan_remat_parity(self):
         hidden_states = jnp.ones((self.batch_size, self.seq_len, self.in_channels)) * 0.5
         audio_hidden_states = jnp.ones((self.batch_size, 128, self.audio_in_channels)) * 0.5
         timestep = jnp.array([1.0])
-        encoder_hidden_states = jnp.ones((self.batch_size, 5, 32)) * 0.1
-        audio_encoder_hidden_states = jnp.ones((self.batch_size, 5, 32)) * 0.1
+        encoder_hidden_states = jnp.ones((self.batch_size, 128, 32)) * 0.1
+        audio_encoder_hidden_states = jnp.ones((self.batch_size, 128, 32)) * 0.1
 
         inp_args = dict(
             hidden_states=hidden_states,