add mup configuration

src/haliax/nn/attention.py

@@ -33,6 +33,7 @@ def dot_product_attention_weights(
     bias: Optional[NamedArray] = None,
     attention_dtype: Optional[jnp.dtype] = None,
     precision: PrecisionLike = None,
+    use_mup: bool = False,
 ) -> NamedArray:
     """
     NamedArray version of dot product attention. Computes the logits for the attention weights. Note that the
@@ -51,7 +52,10 @@ def dot_product_attention_weights(
     # cf https://github.com/google/flax/blob/509bf97ea272e130d932920f45307ac98947d994/flax/linen/attention.py#L40
 
     orig_dtype = query.dtype
-    query = query / jnp.sqrt(query.axis_size(Key))
+    if use_mup:
+        query = query / query.axis_size(Key)
+    else:
+        query = query / jnp.sqrt(query.axis_size(Key))
 
     if attention_dtype is not None:
         query = query.astype(attention_dtype)