embed -> none

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -717,7 +717,7 @@ def __init__(
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
-        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "mlp")),
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), (None, "mlp")),
         bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("mlp",)),
     )
     self.act = get_activation(activation_fn)
@@ -729,8 +729,8 @@ def __init__(
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
-        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("mlp","embed")),
-        bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("mlp", None)),
+        bias_init=nnx.with_partitioning(nnx.initializers.zeros, (None,)),
     )
 
   def __call__(self, hidden_states: Array) -> Array: