fix(ltx2): add robust fallback for missing sharding specs in attention

Author: Perseus14

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -364,19 +364,33 @@ def __init__(
     else:
       specs = sharding_specs
 
+    # Use getattr with fallback to default specs if attribute is missing
+    def get_spec(attr_name):
+      val = getattr(specs, attr_name, None)
+      if val is None:
+        default_specs = get_sharding_specs("default", "ltx2_dit")
+        return getattr(default_specs, attr_name)
+      return val
+
+    qkv_kernel = get_spec("qkv_kernel")
+    qkv_bias = get_spec("qkv_bias")
+    out_kernel = get_spec("out_kernel")
+    out_bias = get_spec("out_bias")
+    norm_scale = get_spec("norm_scale")
+
     # 1. Define Partitioned Initializers (Logical Axes)
     # Q, K, V kernels: [in_features (embed), out_features (heads)]
-    qkv_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), specs.qkv_kernel)
+    qkv_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), qkv_kernel)
     # Q, K, V biases: [out_features (heads)]
-    qkv_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), specs.qkv_bias)
+    qkv_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), qkv_bias)
 
     # Out kernel: [in_features (heads), out_features (embed)]
-    out_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), specs.out_kernel)
+    out_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), out_kernel)
     # Out bias: [out_features (embed)]
-    out_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), specs.out_bias)
+    out_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), out_bias)
 
     # Norm scales
-    norm_scale_init = nnx.with_partitioning(nnx.initializers.ones_init(), specs.norm_scale)
+    norm_scale_init = nnx.with_partitioning(nnx.initializers.ones_init(), norm_scale)
 
     # 2. Projections
     self.to_q = nnx.Linear(