testing with sharding changes

Author: prishajain1

src/maxdiffusion/models/ltx2/attention_ltx2.py

@@ -359,13 +359,13 @@ def __init__(
     # 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(), ("embed", "heads"))
-    # Q, K, V biases: [out_features (embed)]
-    qkv_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), ("embed",))
+    # Q, K, V biases: [out_features (heads)]
+    qkv_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), ("heads",))
 
     # Out kernel: [in_features (heads), out_features (embed)]
     out_kernel_init = nnx.with_partitioning(nnx.initializers.lecun_normal(), ("heads", "embed"))
-    # Out bias: [out_features (heads)]
-    out_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), ("heads",))
+    # Out bias: [out_features (embed)]
+    out_bias_init = nnx.with_partitioning(nnx.initializers.zeros_init(), ("embed",))
 
     # Norm scales
     norm_scale_init = nnx.with_partitioning(nnx.initializers.ones_init(), ("norm",))