feat(ltx2): add attention specs to text connector and pass them

Author: Perseus14

src/maxdiffusion/models/ltx2/logical_sharding_ltx2.py

@@ -68,6 +68,12 @@ class TextConnectorShardingSpecs:
   net_0_bias: tuple = ("mlp",)
   net_2_kernel: tuple = ("mlp", None)
   net_2_bias: tuple = (None,)
+  # Attention specs (defaulting to trillium/safe defaults)
+  qkv_kernel: tuple = ("embed", "heads")
+  out_kernel: tuple = ("heads", "embed")
+  out_bias: tuple = ("embed",)
+  qkv_bias: tuple = ("heads",)
+  norm_scale: tuple = ("norm",)
 
 
 @dataclass
@@ -95,7 +101,11 @@ class VAEShardingSpecs:
             use_batched_text_encoder=True,
             text_encoder_kernel=(None, "embed"),
         ),
-        "text_connector": TextConnectorShardingSpecs(),
+        "text_connector": TextConnectorShardingSpecs(
+            qkv_kernel=(None, "heads"),
+            out_kernel=("heads", None),
+            out_bias=(None,),
+        ),
         "vae": VAEShardingSpecs(vae_conv_kernel=("batch", None, None, None)),
     },
     "trillium": {

src/maxdiffusion/models/ltx2/text_encoders/embeddings_connector_ltx2.py

@@ -49,6 +49,7 @@ def __init__(
         attention_kernel=attention_kernel,
         mesh=mesh,
         rngs=rngs,
+        sharding_specs=sharding_specs,
     )
     self.ff = NNXSimpleFeedForward(rngs=rngs, dim=dim, dim_out=dim, activation_fn="gelu_tanh", sharding_specs=sharding_specs)
     self.norm1 = nnx.RMSNorm(dim, epsilon=1e-6, dtype=jnp.float32, param_dtype=jnp.float32, use_scale=False, rngs=rngs)