feat(ltx2): use cleaner approach for connector attention sharding

Author: Perseus14

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -364,33 +364,19 @@ 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(), qkv_kernel)
+    qkv_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), specs.qkv_kernel)
     # Q, K, V biases: [out_features (heads)]
-    qkv_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), qkv_bias)
+    qkv_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), specs.qkv_bias)
 
     # Out kernel: [in_features (heads), out_features (embed)]
-    out_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), out_kernel)
+    out_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), specs.out_kernel)
     # Out bias: [out_features (embed)]
-    out_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), out_bias)
+    out_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), specs.out_bias)
 
     # Norm scales
-    norm_scale_init = nnx.with_partitioning(nnx.initializers.ones_init(), norm_scale)
+    norm_scale_init = nnx.with_partitioning(nnx.initializers.ones_init(), specs.norm_scale)
 
     # 2. Projections
     self.to_q = nnx.Linear(

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

@@ -49,7 +49,6 @@ 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)