sharding

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

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