Add attention block

Author: Perseus14

src/maxdiffusion/models/ltx2/__init__.py

@@ -23,7 +23,7 @@
 import pandas as pd
 from jax.sharding import Mesh
 from jax.experimental.pallas.ops.tpu.splash_attention import splash_attention_kernel
-
+from ...models.ltx2.attention_ltx2 import LTX2Attention, LTX2RotaryPosEmbed
 # Set JAX to use float32 for higher precision checks
 jax.config.update("jax_default_matmul_precision", "float32")
 
@@ -175,9 +175,8 @@ def forward(self, x, context=None, q_rope=None, k_rope=None, mask=None):
 
 
 # ==========================================
-# 2. JAX Imports & Test Suite
+# 2. JAX Test Suite
 # ==========================================
-from ...models.ltx2.attention_ltx2 import LTX2Attention, LTX2RotaryPosEmbed
 
 
 class LTX2AttentionTest(unittest.TestCase):