dot attn for audio self attn+ replication of audio stream over context axis

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -159,7 +159,7 @@ def __init__(
         eps=norm_eps,
         dtype=dtype,
         mesh=mesh,
-        attention_kernel=self.attention_kernel,
+        attention_kernel="dot_product",
         rope_type=rope_type,
         flash_block_sizes=flash_block_sizes,
     )
@@ -212,7 +212,7 @@ def __init__(
         eps=norm_eps,
         dtype=dtype,
         mesh=mesh,
-        attention_kernel=self.attention_kernel,
+        attention_kernel="dot_product",
         rope_type=rope_type,
         flash_block_sizes=flash_block_sizes,
     )
@@ -350,7 +350,8 @@ def __call__(
 
     axis_names = nn.logical_to_mesh_axes(("activation_batch", "activation_length", "activation_embed"))
     hidden_states = jax.lax.with_sharding_constraint(hidden_states, axis_names)
-    audio_hidden_states = jax.lax.with_sharding_constraint(audio_hidden_states, axis_names)
+    axis_names_audio = nn.logical_to_mesh_axes(("activation_batch", None, "activation_embed"))
+    audio_hidden_states = jax.lax.with_sharding_constraint(audio_hidden_states, axis_names_audio)
 
     if encoder_hidden_states is not None:
       encoder_hidden_states = jax.lax.with_sharding_constraint(encoder_hidden_states, axis_names)