sharding

prishajain1 · prishajain1 · commit 022520191cf0 · 2026-01-14T23:32:40.000+05:30
diff --git a/src/maxdiffusion/models/attention_flax.py b/src/maxdiffusion/models/attention_flax.py
@@ -717,7 +717,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
`@@ -717,7 +717,7 @@ def __init__(self, rngs: nnx.Rngs, dim: int, dim_out: Optional[int] = None, mult`
`717`	`717`	`param_dtype=weights_dtype,`
`718`	`718`	`precision=precision,`
`719`	`719`	`kernel_init=nnx.with_partitioning(nnx.initializers.lecun_normal(), ("embed", None)),`
`720`		`- bias_init=nnx.with_partitioning(nnx.initializers.zeros, ("embed",)),`
	`720`	`+ bias_init=nnx.with_partitioning(nnx.initializers.zeros, (None,)),`
`721`	`721`	`)`
`722`	`722`	`self.act = get_activation(activation_fn)`
`723`	`723`	`self.net_2 = nnx.Linear(`