bias_init added in attention_flax.py

Author: prishajain1

src/maxdiffusion/models/attention_flax.py

@@ -888,14 +888,26 @@ def __init__(
     if self.added_kv_proj_dim is not None:
       self.add_k_proj = nnx.Linear(
           self.added_kv_proj_dim, self.inner_dim, rngs=rngs,
-          dtype=dtype, param_dtype=weights_dtype, precision=precision
+          dtype=dtype, param_dtype=weights_dtype, precision=precision,
+          bias_init=nnx.with_partitioning(
+              nnx.initializers.zeros,
+              ("embed",), 
+          ),
       )
       self.add_v_proj = nnx.Linear(
           self.added_kv_proj_dim, self.inner_dim, rngs=rngs,
-          dtype=dtype, param_dtype=weights_dtype, precision=precision
+          dtype=dtype, param_dtype=weights_dtype, precision=precision,
+          bias_init=nnx.with_partitioning(
+              nnx.initializers.zeros,
+              ("embed",),
+          ),
       )
       self.norm_added_k = nnx.RMSNorm(
-          num_features=self.inner_dim, rngs=rngs, epsilon=eps, dtype=dtype, param_dtype=weights_dtype
+          num_features=self.inner_dim, rngs=rngs, epsilon=eps, dtype=dtype, param_dtype=weights_dtype,
+          scale_init=nnx.with_partitioning(
+              nnx.initializers.ones,
+              ("norm",),
+          ),
       )
 
   def _apply_rope(self, xq: jax.Array, xk: jax.Array, freqs_cis: jax.Array) -> Tuple[jax.Array, jax.Array]:

src/maxdiffusion/pyconfig.py

@@ -195,10 +195,6 @@ def user_init(raw_keys):
     max_utils.write_config_raw_keys_for_gcs(raw_keys)
 
     raw_keys["logical_axis_rules"] = _lists_to_tuples(raw_keys["logical_axis_rules"])
-    logical_axis_rules = list(raw_keys["logical_axis_rules"])
-    logical_axis_rules.append(('bias', 'tensor'))
-    logical_axis_rules.append(('attn2', 'add_k_proj', 'bias', 'tensor'))
-    raw_keys["logical_axis_rules"] = tuple(logical_axis_rules)
     # Verify qkv is sharded across sequence.
     if raw_keys["attention"] == "ring":
       logical_axis_rules = list(raw_keys["logical_axis_rules"])