test: dynamically scale LTX2 sequence length based on device count

Author: Perseus14

src/maxdiffusion/tests/ltx2/test_transformer_ltx2.py

@@ -173,14 +173,15 @@ def test_kv_cache_rejects_timestep_modulated_prompt_embeddings(self):
   def test_kv_cache_matches_uncached_output_with_and_without_layer_scan(self):
     """Layer-scanned KV inputs must be equivalent to normal cross-attention projections."""
     batch_size = 1
-    num_frames, height, width = 2, 2, 2
-    audio_num_frames = 4
+    seq_len = jax.device_count()
+    num_frames, height, width = seq_len, 2, 2
+    audio_num_frames = seq_len
     video_dim, audio_dim, caption_dim = 64, 64, 32
     hidden_states = jnp.ones((batch_size, num_frames * height * width, 8), dtype=jnp.float32)
     audio_hidden_states = jnp.ones((batch_size, audio_num_frames, 4), dtype=jnp.float32)
-    encoder_hidden_states = jnp.ones((batch_size, 4, caption_dim), dtype=jnp.float32)
-    audio_encoder_hidden_states = jnp.ones((batch_size, 4, caption_dim), dtype=jnp.float32)
-    attention_mask = jnp.ones((batch_size, 4), dtype=jnp.float32)
+    encoder_hidden_states = jnp.ones((batch_size, seq_len, caption_dim), dtype=jnp.float32)
+    audio_encoder_hidden_states = jnp.ones((batch_size, seq_len, caption_dim), dtype=jnp.float32)
+    attention_mask = jnp.ones((batch_size, seq_len), dtype=jnp.float32)
 
     for scan_layers in (False, True):
       with self.subTest(scan_layers=scan_layers), self.mesh, nn_partitioning.axis_rules(self.config.logical_axis_rules):