sharding added

prishajain1 · prishajain1 · commit f9cc8f82814c · 2026-01-16T00:12:32.000+05:30
diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py
@@ -684,7 +684,7 @@ def __init__(self, rngs: nnx.Rngs, dim: int, dim_out: Optional[int] = None, mult
         dtype=dtype,
         param_dtype=weights_dtype,
         precision=precision,
-        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "embed")),
+        kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", None)),
         bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),
     )
     self.act = get_activation(activation_fn)

Original file line number	Diff line number	Diff line change
`@@ -684,7 +684,7 @@ def __init__(self, rngs: nnx.Rngs, dim: int, dim_out: Optional[int] = None, mult`
`684`	`684`	`dtype=dtype,`
`685`	`685`	`param_dtype=weights_dtype,`
`686`	`686`	`precision=precision,`
`687`		`- kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", "embed")),`
	`687`	`+ kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", None)),`
`688`	`688`	`bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),`
`689`	`689`	`)`
`690`	`690`	`self.act = get_activation(activation_fn)`