fix for tpu flash attention

Author: prishajain1

src/maxdiffusion/models/ltx2/transformer_ltx2.py

@@ -821,14 +821,16 @@ def __call__(
         if audio_encoder_attention_mask is not None: print_shape("audio_encoder_attention_mask input", audio_encoder_attention_mask)
 
 
-        if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
-            encoder_attention_mask = (1 - encoder_attention_mask.astype(self.dtype)) * -10000.0
-            encoder_attention_mask = jnp.expand_dims(encoder_attention_mask, axis=1)
+        if self.attention_kernel == "dot_product":
+            if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
+                encoder_attention_mask = (1 - encoder_attention_mask.astype(self.dtype)) * -10000.0
+                encoder_attention_mask = jnp.expand_dims(encoder_attention_mask, axis=1)
+            
+            if audio_encoder_attention_mask is not None and audio_encoder_attention_mask.ndim == 2:
+                audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.astype(self.dtype)) * -10000.0
+                audio_encoder_attention_mask = jnp.expand_dims(audio_encoder_attention_mask, axis=1)
+        
         if encoder_attention_mask is not None: print_shape("encoder_attention_mask bias", encoder_attention_mask)
-
-        if audio_encoder_attention_mask is not None and audio_encoder_attention_mask.ndim == 2:
-            audio_encoder_attention_mask = (1 - audio_encoder_attention_mask.astype(self.dtype)) * -10000.0
-            audio_encoder_attention_mask = jnp.expand_dims(audio_encoder_attention_mask, axis=1)
         if audio_encoder_attention_mask is not None: print_shape("audio_encoder_attention_mask bias", audio_encoder_attention_mask)
 
         batch_size = hidden_states.shape[0]