feat(ltx2): add defaults to LTX2DiTShardingSpecs matching Trillium

Author: Perseus14

src/maxdiffusion/models/ltx2/logical_sharding_ltx2.py

@@ -25,9 +25,9 @@ class LTX2DiTShardingSpecs:
   """Sharding specs for the LTX2 Diffusion Transformer."""
 
   # --- Attention Layers (LTX2Attention) ---
-  qkv_kernel: tuple
-  out_kernel: tuple
-  out_bias: tuple
+  qkv_kernel: tuple = ("embed", "heads")
+  out_kernel: tuple = ("heads", "embed")
+  out_bias: tuple = ("embed",)
   qkv_bias: tuple = ("heads",)
 
   # --- Feed Forward Network (NNXSimpleFeedForward) ---
@@ -113,11 +113,7 @@ class VAEShardingSpecs:
         "vae": VAEShardingSpecs(),
     },
     "trillium": {
-        "ltx2_dit": LTX2DiTShardingSpecs(
-            qkv_kernel=("embed", "heads"),
-            out_kernel=("heads", "embed"),
-            out_bias=("embed",),
-        ),
+        "ltx2_dit": LTX2DiTShardingSpecs(),
         "text_encoder": TextEncoderShardingSpecs(
             use_batched_text_encoder=False,
             text_encoder_kernel=(None, "embed"),