Fix WAN transformer partitioning for bias and kernel init

Author: csgoogle

src/maxdiffusion/models/attention_flax.py

@@ -979,7 +979,7 @@ def __init__(
         precision=precision,
         bias_init=nnx.with_partitioning(
             nnx.initializers.zeros,
-            ("embed",),
+            ("heads",),
         ),
     )
 
@@ -993,7 +993,7 @@ def __init__(
         precision=precision,
         bias_init=nnx.with_partitioning(
             nnx.initializers.zeros,
-            ("embed",),
+            ("heads",),
         ),
     )
 
@@ -1007,7 +1007,7 @@ def __init__(
         precision=precision,
         bias_init=nnx.with_partitioning(
             nnx.initializers.zeros,
-            ("embed",),
+            ("heads",),
         ),
     )
 
@@ -1021,7 +1021,7 @@ def __init__(
         precision=precision,
         bias_init=nnx.with_partitioning(
             nnx.initializers.zeros,
-            ("heads",),
+            ("embed",),
         ),
     )
 
@@ -1332,11 +1332,13 @@ def setup(self):
         precision=self.precision,
     )
 
+    proj_attn_kernel_axes = ("heads", "embed")
+
     self.proj_attn = nn.Dense(
         self.query_dim,
-        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), kernel_axes),
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), proj_attn_kernel_axes),
         use_bias=True,
-        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("heads",)),
+        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("embed",)),
         dtype=self.dtype,
         param_dtype=self.weights_dtype,
         name="i_proj",
@@ -1345,9 +1347,9 @@ def setup(self):
 
     self.encoder_proj_attn = nn.Dense(
         self.query_dim,
-        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), kernel_axes),
+        kernel_init=nn.with_logical_partitioning(nn.initializers.lecun_normal(), proj_attn_kernel_axes),
         use_bias=True,
-        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("heads",)),
+        bias_init=nn.with_logical_partitioning(nn.initializers.zeros, ("embed",)),
         dtype=self.dtype,
         param_dtype=self.weights_dtype,
         name="e_proj",

src/maxdiffusion/models/wan/transformers/transformer_wan.py

@@ -193,11 +193,11 @@ def __init__(
         kernel_init=nnx.with_partitioning(
             nnx.initializers.xavier_uniform(),
             (
-                "mlp",
                 "embed",
+                "mlp",
             ),
         ),
-        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("mlp",)),
     )
 
   def __call__(self, x: jax.Array) -> jax.Array:
@@ -249,8 +249,8 @@ def __init__(
         kernel_init=nnx.with_partitioning(
             nnx.initializers.xavier_uniform(),
             (
+                 "mlp",
                 "embed",
-                "mlp",
             ),
         ),
     )