haliax 1.4.dev420__tar.gz → 1.4.dev439__tar.gz

{haliax-1.4.dev420 → haliax-1.4.dev439}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev420
+Version: 1.4.dev439
 Summary: Named Tensors for Legible Deep Learning in JAX
 Project-URL: Homepage, https://github.com/stanford-crfm/haliax
 Project-URL: Bug Tracker, https://github.com/stanford-crfm/haliax/issues/

{haliax-1.4.dev420 → haliax-1.4.dev439}/docs/api.md

@@ -4,6 +4,37 @@ that we use names (either strings or [haliax.Axis][] objects) to specify axes in
 arrays (see [haliax.zeros][] and [haliax.ones][]) as well as things like reductions (see [haliax.sum][] and
 [haliax.mean][]).
 
+## PyTree Helpers
+
+PyTrees are the lingua franca for composing state in JAX ecosystems. Haliax provides drop-in replacements for the
+[`jax.tree`][] helpers that are aware of [`NamedArray`][haliax.NamedArray] semantics. They preserve axis metadata across
+transformations while interoperating with standard JAX containers, so you can use them anywhere you would have reached
+for JAX's versions.
+
+Use these helpers whenever you need to map, flatten, or rebuild PyTrees that might include `NamedArray` instances:
+
+* [`haliax.tree.map`][] mirrors [`jax.tree.map`][] but forwards to Haliax's [`haliax.tree_util.tree_map`][] so axis names remain
+  intact.
+* [`haliax.tree.scan_aware_map`][] descends into [`haliax.nn.Stacked`][haliax.nn.Stacked] modules so that each layer is
+  transformed individually, effectively treating them as if they were unrolled when applying
+  [`haliax.tree_util.scan_aware_tree_map`][].
+* [`haliax.tree.flatten`][] / [`haliax.tree.unflatten`][] match the familiar flattening API while handling `NamedArray`
+  payloads safely.
+* [`haliax.tree.leaves`][] and [`haliax.tree.structure`][] provide direct access to the leaves and PyTree structure.
+
+All of these helpers accept the same `is_leaf` hook you might already use with JAX's utilities. They should be the first
+tools you reach for when you need deterministic tree transforms that understand named axes.
+
+::: haliax.tree.map
+::: haliax.tree.scan_aware_map
+::: haliax.tree.flatten
+::: haliax.tree.unflatten
+::: haliax.tree.leaves
+::: haliax.tree.structure
+
+[`jax.tree`]: https://jax.readthedocs.io/en/latest/_autosummary/jax.tree.html
+[`jax.tree.map`]: https://jax.readthedocs.io/en/latest/_autosummary/jax.tree.map.html
+
 ## Axis Types
 
 If you already speak NumPy or `jax.numpy`, think of Haliax as swapping positional axes (`axis=0`) for named axes

{haliax-1.4.dev420 → haliax-1.4.dev439}/src/haliax/__about__.py

@@ -3,4 +3,4 @@
 # SPDX-License-Identifier: Apache-2.0
 
 
-__version__ = "1.4.dev420"
+__version__ = "1.4.dev439"

{haliax-1.4.dev420 → haliax-1.4.dev439}/src/haliax/__init__.py

@@ -16,6 +16,7 @@ import haliax.nn as nn
 import haliax.quantization as quantization
 import haliax.random as random
 import haliax.state_dict as state_dict
+import haliax.tree as tree  # noqa: F401
 import haliax.tree_util as tree_util
 import haliax.util as util
 from .field import field

{haliax-1.4.dev420 → haliax-1.4.dev439}/src/haliax/nn/embedding.py

@@ -11,21 +11,31 @@ from jaxtyping import PRNGKeyArray
 
 import haliax as hax
 
+from .mup import AbstractEmbeddingReparam, ReparamEnabled, EmbeddingStandardParam
 from ..axis import Axis, AxisSpec, concat_axes
 from ..core import NamedArray
 from ..jax_utils import named_call
 from ..tree_util import resize_axis
 
 
-class Embedding(eqx.Module):
+class Embedding(eqx.Module, ReparamEnabled):
     weight: NamedArray
 
     # axes
     Vocab: Axis = eqx.field(static=True)
     Embed: AxisSpec = eqx.field(static=True)
+    reparam: AbstractEmbeddingReparam = eqx.field(static=True)
 
     @staticmethod
-    def init(Vocab: Axis, Embed: AxisSpec, *, init_scale: float = 1, key, initializer_range: float | None = None):
+    def init(
+        Vocab: Axis,
+        Embed: AxisSpec,
+        *,
+        init_scale: float = 1,
+        key,
+        initializer_range: float | None = None,
+        reparam_cls: type[AbstractEmbeddingReparam] = EmbeddingStandardParam,
+    ):
         """
         Initialize an Embedding module.
 
@@ -41,13 +51,17 @@ class Embedding(eqx.Module):
             initializer_range: Deprecated. Use init_scale instead.
         """
         if initializer_range is not None:
-            warnings.warn("initializer_range is deprecated. Use init_std instead.", DeprecationWarning)
+            warnings.warn(
+                "initializer_range is deprecated. Use init_std instead.",
+                DeprecationWarning,
+            )
             init_scale = initializer_range
 
         all_axes = concat_axes(Vocab, Embed)
-        output_size = hax.axis_size(Embed)
-        weight = hax.random.truncated_normal(key, all_axes, -3, 3) * (init_scale / output_size)
-        return Embedding(weight=weight, Vocab=Vocab, Embed=Embed)
+        weight = hax.random.truncated_normal(key, all_axes, -3, 3) * (
+            init_scale * reparam_cls.init_scale(Vocab, Embed)
+        )
+        return Embedding(weight=weight, Vocab=Vocab, Embed=Embed, reparam=reparam_cls(Embed, Vocab))
 
     def __call__(self, input_ids: NamedArray, *, key: PRNGKeyArray | None = None):
         """Alias for `embed`. key is ignored."""
@@ -60,7 +74,7 @@ class Embedding(eqx.Module):
             input_ids: token IDs with shape > {Vocab}
         """
         input_embeds = self.weight.take(self.Vocab, input_ids)
-        return input_embeds
+        return input_embeds * self.reparam.active_scale
 
     def unembed(self, input_embeds: NamedArray):
         """
@@ -68,7 +82,7 @@ class Embedding(eqx.Module):
 
         Equivalent to `input_embeds.dot(self.weight, axis=self.Embed)`.
         """
-        return input_embeds.dot(self.weight, axis=self.Embed)
+        return input_embeds.dot(self.weight, axis=self.Embed) * self.reparam.unembed_active_scale
 
     def resize_embeddings(self, new_size: int, key: PRNGKeyArray | None = None):
         """

{haliax-1.4.dev420 → haliax-1.4.dev439}/src/haliax/nn/linear.py

@@ -6,6 +6,7 @@
 import dataclasses
 import math
 from functools import partial
+from typing import Optional
 
 import equinox as eqx
 import jax
@@ -17,7 +18,16 @@ from jaxtyping import PRNGKeyArray
 
 import haliax as hax
 
-from .._src.state_dict import Mod, ModuleWithStateDictSerialization
+
+from . import mup
+from .mup import AbstractLinearReparam, ReparamEnabled, LinearStandardParam
+from .._src.state_dict import (
+    Mod,
+    ModuleWithStateDictSerialization,
+    StateDict,
+    default_eqx_module_from_state_dict,
+    default_eqx_module_to_state_dict,
+)
 from ..axis import Axis, AxisSpec
 from ..core import NamedArray
 from ..jax_utils import named_call
@@ -26,7 +36,7 @@ from ..quantization import DotGeneralOp
 from ..util import ensure_tuple
 
 
-class Linear(ModuleWithStateDictSerialization):
+class Linear(ModuleWithStateDictSerialization, ReparamEnabled):
     """A named Linear layer. This module allows you to specify multiple named axes for both input
     and output, which is occasionally useful."""
 
@@ -35,6 +45,7 @@ class Linear(ModuleWithStateDictSerialization):
 
     In: AxisSpec = eqx.field(static=True)
     Out: AxisSpec = eqx.field(static=True)
+    reparam: AbstractLinearReparam = eqx.field(static=True)
     dot_general: DotGeneralOp = eqx.field(default_factory=DotGeneralOp.default)
 
     @staticmethod
@@ -47,6 +58,7 @@ class Linear(ModuleWithStateDictSerialization):
         out_first: bool = True,
         dot_general: DotGeneralOp | None = None,
         init_scale: float = 1.0,
+        reparam_cls: type[AbstractLinearReparam] = LinearStandardParam,
     ) -> "Linear":
         """
         Args:
@@ -59,14 +71,13 @@ class Linear(ModuleWithStateDictSerialization):
             init_scale: float: The scale to use for initialization. We scale init by 1/sqrt(Input.size)*init_scale
         """
         joint_spec = hax.concat_axis_specs(Out, In) if out_first else hax.concat_axis_specs(In, Out)
-        input_size = hax.axis_size(In)
-        weight = hax.random.truncated_normal(key, joint_spec, -3, 3) * (init_scale / math.sqrt(input_size))
+        weight = hax.random.truncated_normal(key, joint_spec, -3, 3) * (init_scale * reparam_cls.init_scale(In, Out))
         bias = hax.zeros(Out) if use_bias else None
 
         if dot_general is None:
             dot_general = DotGeneralOp.default()
 
-        return Linear(weight, bias, In, Out, dot_general=dot_general)
+        return Linear(weight, bias, In, Out, dot_general=dot_general, reparam=reparam_cls(In, Out))
 
     @named_call
     def __call__(self, inputs, *, key: PRNGKeyArray | None = None):
@@ -76,7 +87,11 @@ class Linear(ModuleWithStateDictSerialization):
             key: Not used, but there for compat with other modules
         """
         del key
-        q = inputs.dot(self.weight, axis=self.In, dot_general=self.dot_general)
+        q = inputs.dot(
+            self.weight * self.reparam.active_scale,
+            axis=self.In,
+            dot_general=self.dot_general,
+        )
         q = hax.auto_sharded(q)
 
         if self.bias is not None:
@@ -137,6 +152,32 @@ class Linear(ModuleWithStateDictSerialization):
         else:
             return self.weight.axes[-len(self.Out) :] != self.Out
 
+    def to_state_dict(self, prefix: Optional[str] = None) -> StateDict:
+        scaled = dataclasses.replace(self, weight=self.weight * self.reparam.active_scale)
+        return default_eqx_module_to_state_dict(scaled, prefix)
+
+    def from_state_dict(self: Mod, state_dict: StateDict, prefix: Optional[str] = None) -> Mod:
+        unscaled = default_eqx_module_from_state_dict(self, state_dict, prefix)
+        return dataclasses.replace(unscaled, weight=unscaled.weight / self.reparam.active_scale)
+
+    @staticmethod
+    def input_reparam(use_mup: bool = True) -> type[AbstractLinearReparam]:
+        """Return the reparameterization class for an input linear layer."""
+
+        return mup.InputLinearMup if use_mup else mup.LinearStandardParam
+
+    @staticmethod
+    def hidden_reparam(use_mup: bool = True) -> type[AbstractLinearReparam]:
+        """Return the reparameterization class for a hidden linear layer."""
+
+        return mup.HiddenLinearMup if use_mup else mup.LinearStandardParam
+
+    @staticmethod
+    def output_reparam(use_mup: bool = True) -> type[AbstractLinearReparam]:
+        """Return the reparameterization class for an output linear layer."""
+
+        return mup.OutputLinearMup if use_mup else mup.LinearStandardParam
+
 
 class MoELinear(eqx.Module):
     """A named Linear layer for MoE. This module allows you to specify multiple named axes for both input
@@ -197,7 +238,10 @@ class MoELinear(eqx.Module):
         dim_numbers = jax.lax.RaggedDotDimensionNumbers(
             dot_dimension_numbers=(
                 # contracting
-                (ensure_tuple(inputs.axis_indices(self.In)), ensure_tuple(self.weight.axis_indices(self.In))),
+                (
+                    ensure_tuple(inputs.axis_indices(self.In)),
+                    ensure_tuple(self.weight.axis_indices(self.In)),
+                ),
                 # batch
                 ((), ()),
             ),

haliax-1.4.dev439/src/haliax/nn/mup.py

@@ -0,0 +1,206 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import math
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+
+import haliax as hax
+import equinox as eqx
+
+from ..axis import AxisSpec
+
+
+class AbstractReparam(ABC):
+    """Abstract base class for abc-parameterization rules.
+
+    Defines the interface for active scaling of parameters (a),
+    computing initialization scales (b), and learning rate scaling (c)
+
+    See: https://arxiv.org/abs/2011.14522
+    """
+
+    @staticmethod
+    @abstractmethod
+    def init_scale(In: AxisSpec, Out: AxisSpec):
+        """Return the scaling factor for initializing weights
+        given input and output axes."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def lr_scale(self):
+        """Return the learning-rate scaling factor."""
+        raise NotImplementedError
+
+    @property
+    @abstractmethod
+    def active_scale(self):
+        """Return the scaling applied to activations."""
+        raise NotImplementedError
+
+
+@dataclass
+class AbstractLinearReparam(AbstractReparam):
+    """Base class for linear-layer reparameterizations.
+
+    Stores input and output axis specifications, and inherits
+    the reparameterization interface.
+    """
+
+    In: AxisSpec
+    Out: AxisSpec
+
+
+class LinearStandardParam(AbstractLinearReparam):
+    """Standard (non-muP) parameterization for linear layers.
+
+    Uses the usual fan-in scaling for initialization and
+    leaves learning rate and activation scaling unchanged.
+    """
+
+    @staticmethod
+    def init_scale(In: AxisSpec, Out: AxisSpec):
+        return 1 / math.sqrt(hax.axis_size(In))
+
+    @property
+    def active_scale(self):
+        return 1
+
+    @property
+    def lr_scale(self):
+        return 1
+
+
+class InputLinearMup(AbstractLinearReparam):
+    """muP-style parameterization for input linear layers.
+
+    Uses no scaling on initialization or learning rate.
+    See: https://arxiv.org/abs/2011.14522 (Maximal Update Parametrization)
+    """
+
+    @staticmethod
+    def init_scale(In: AxisSpec, Out: AxisSpec):
+        return 1
+
+    @property
+    def active_scale(self):
+        return 1
+
+    @property
+    def lr_scale(self):
+        return 1
+
+
+class HiddenLinearMup(AbstractLinearReparam):
+    """muP-style parameterization for hidden linear layers.
+
+    Applies fan-in scaling at initialization and scales
+    learning rate inversely with layer width.
+    """
+
+    @staticmethod
+    def init_scale(In: AxisSpec, Out: AxisSpec):
+        return 1 / math.sqrt(hax.axis_size(In))
+
+    @property
+    def active_scale(self):
+        return 1
+
+    @property
+    def lr_scale(self):
+        return 1 / hax.axis_size(self.In)
+
+
+class OutputLinearMup(AbstractLinearReparam):
+    """muP-style parameterization for output linear layers.
+
+    Uses unit initialization and applies inverse-width
+    scaling to the output activations.
+    """
+
+    @staticmethod
+    def init_scale(In: AxisSpec, Out: AxisSpec):
+        return 1
+
+    @property
+    def active_scale(self):
+        return 1 / hax.axis_size(self.In)
+
+    @property
+    def lr_scale(self):
+        return 1
+
+
+@dataclass
+class AbstractEmbeddingReparam(AbstractReparam):
+    """Base class for embedding-layer reparameterizations.
+
+    Defines the interface for both embedding and unembedding
+    scaling rules.
+    """
+
+    Embed: AxisSpec
+    Vocab: AxisSpec
+
+    @property
+    @abstractmethod
+    def unembed_active_scale(self):
+        """Scaling factor applied when unembedding embeddings."""
+        raise NotImplementedError
+
+
+class EmbeddingStandardParam(AbstractEmbeddingReparam):
+    """Standard embedding parameterization."""
+
+    @staticmethod
+    def init_scale(In: AxisSpec, Out: AxisSpec):
+        return 1 / hax.axis_size(Out)
+
+    @property
+    def active_scale(self):
+        return 1
+
+    @property
+    def lr_scale(self):
+        return 1
+
+    @property
+    def unembed_active_scale(self):
+        return 1 / hax.axis_size(self.Embed)
+
+
+class EmbeddingMup(AbstractEmbeddingReparam):
+    """muP-style parameterization for embeddings.
+
+    Keeps initialization and learning-rate scaling neutral,
+    but applies inverse-width scaling to unembedding outputs for tied weights.
+    See: https://www.cerebras.ai/blog/the-practitioners-guide-to-the-maximal-update-parameterization
+    """
+
+    @staticmethod
+    def init_scale(In: AxisSpec, Out: AxisSpec):
+        return 1
+
+    @property
+    def active_scale(self):
+        return 1
+
+    @property
+    def lr_scale(self):
+        return 1
+
+    @property
+    def unembed_active_scale(self):
+        return 1 / hax.axis_size(self.Embed)
+
+
+class ReparamEnabled(ABC):
+    """Mixin for modules that support reparameterization.
+
+    Stores an abstract `reparam` attribute that specifies
+    how initialization and scaling are handled.
+    """
+
+    reparam: eqx.AbstractVar[AbstractReparam]

haliax-1.4.dev439/src/haliax/tree.py

@@ -0,0 +1,59 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Convenience wrappers for :mod:`haliax.tree_util` that mirror :mod:`jax.tree`."""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Iterable, Sequence, TypeVar
+
+from . import tree_util
+
+T = TypeVar("T")
+
+
+def map(fn: Callable[..., T], tree: Any, *rest: Any, is_leaf: Callable[[Any], bool] | None = None) -> Any:
+    """Alias for :func:`haliax.tree_util.tree_map` matching :func:`jax.tree.map`."""
+
+    return tree_util.tree_map(fn, tree, *rest, is_leaf=is_leaf)
+
+
+def scan_aware_map(fn: Callable[..., T], tree: Any, *rest: Any, is_leaf: Callable[[Any], bool] | None = None) -> Any:
+    """Alias for :func:`haliax.tree_util.scan_aware_tree_map` with :mod:`jax.tree` style naming."""
+
+    return tree_util.scan_aware_tree_map(fn, tree, *rest, is_leaf=is_leaf)
+
+
+def flatten(tree: Any, *, is_leaf: Callable[[Any], bool] | None = None) -> tuple[Sequence[Any], Any]:
+    """Alias for :func:`haliax.tree_util.tree_flatten` matching :func:`jax.tree.flatten`."""
+
+    return tree_util.tree_flatten(tree, is_leaf=is_leaf)
+
+
+def unflatten(treedef: Any, leaves: Iterable[Any]) -> Any:
+    """Alias for :func:`haliax.tree_util.tree_unflatten` matching :func:`jax.tree.unflatten`."""
+
+    return tree_util.tree_unflatten(treedef, leaves)
+
+
+def leaves(tree: Any, *, is_leaf: Callable[[Any], bool] | None = None) -> Sequence[Any]:
+    """Alias for :func:`haliax.tree_util.tree_leaves` matching :func:`jax.tree.leaves`."""
+
+    return tree_util.tree_leaves(tree, is_leaf=is_leaf)
+
+
+def structure(tree: Any, *, is_leaf: Callable[[Any], bool] | None = None) -> Any:
+    """Alias for :func:`haliax.tree_util.tree_structure` matching :func:`jax.tree.structure`."""
+
+    return tree_util.tree_structure(tree, is_leaf=is_leaf)
+
+
+__all__ = [
+    "map",
+    "scan_aware_map",
+    "flatten",
+    "unflatten",
+    "leaves",
+    "structure",
+]

haliax-1.4.dev439/tests/test_mup_coordinate_check.py

@@ -0,0 +1,164 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+"""Coordinate check for µP modules built on Haliax primitives."""
+
+from __future__ import annotations
+
+import dataclasses
+from typing import Any, Iterable
+
+import equinox as eqx
+import jax
+import jax.random as jrandom
+
+import haliax as hax
+from haliax import Axis, NamedArray
+from haliax.nn import Linear, activations
+from haliax.nn.mup import InputLinearMup, HiddenLinearMup, OutputLinearMup
+
+
+class TinyMLP(eqx.Module):
+    """Minimal 2-hidden-layer MLP composed of Haliax Linear variants."""
+
+    first: Linear
+    second: Linear
+    third: Linear
+
+    @staticmethod
+    def init(width: int, *, key: jax.Array, use_mup: bool) -> TinyMLP:
+        in_axis = Axis("in", 2)
+        hidden = Axis("hidden", width)
+        hidden2 = hidden.alias("hidden2")
+        out_axis = Axis("out", 1)
+
+        k1, k2, k3 = jrandom.split(key, 3)
+
+        if use_mup:
+            first = Linear.init((in_axis,), hidden, key=k1, reparam_cls=InputLinearMup)
+            second = Linear.init(hidden, hidden2, key=k2, reparam_cls=HiddenLinearMup)
+            third = Linear.init(hidden2, (out_axis,), key=k3, reparam_cls=OutputLinearMup)
+        else:
+            first = Linear.init((in_axis,), hidden, key=k1)
+            second = Linear.init(hidden, hidden2, key=k2)
+            third = Linear.init(hidden2, (out_axis,), key=k3)
+
+        return TinyMLP(first=first, second=second, third=third)
+
+    def __call__(self, x: NamedArray) -> NamedArray:
+        h = activations.relu(self.first(x))
+        h = activations.relu(self.second(h))
+        return self.third(h)
+
+
+def _loss_fn(params: TinyMLP, x: NamedArray, y: NamedArray) -> jax.Array:
+    preds = params(x)
+    diff = preds - y
+    return hax.mean(diff * diff).scalar()
+
+
+_loss_and_grad = eqx.filter_jit(eqx.filter_value_and_grad(_loss_fn))
+_loss_value = jax.jit(_loss_fn)
+
+
+def _apply_sgd(module: TinyMLP, grads: TinyMLP, *, base_lr: float, use_mup: bool) -> TinyMLP:
+    def update_linear(layer: Linear, grad_layer: Linear) -> Linear:
+        lr_scale = layer.reparam.lr_scale
+        new_weight = layer.weight - (base_lr * lr_scale) * grad_layer.weight
+        if layer.bias is None or grad_layer.bias is None:
+            new_bias = layer.bias
+        else:
+            new_bias = layer.bias - base_lr * grad_layer.bias
+        return dataclasses.replace(layer, weight=new_weight, bias=new_bias)
+
+    return TinyMLP(
+        first=update_linear(module.first, grads.first),
+        second=update_linear(module.second, grads.second),
+        third=update_linear(module.third, grads.third),
+    )
+
+
+def _make_dataset(key: jax.Array, *, n_points: int = 2048) -> tuple[NamedArray, NamedArray]:
+    data_axis = Axis("data", n_points)
+    feature_axis = Axis("in", 2)
+    out_axis = Axis("out", 1)
+
+    xy = jrandom.uniform(key, (n_points, 2), minval=-1.0, maxval=1.0)
+    inputs = hax.named(xy, (data_axis, feature_axis))
+    targets = hax.named(xy[:, :1], (data_axis, out_axis))
+    return inputs, targets
+
+
+def _run_once(
+    key: jax.Array,
+    *,
+    width: int,
+    use_mup: bool,
+    steps: int = 120,
+    batch_size: int = 256,
+    base_lr: float = 3e-3,
+) -> float:
+    data_key, model_key = jrandom.split(key)
+    inputs, targets = _make_dataset(data_key)
+    params = TinyMLP.init(width, key=model_key, use_mup=use_mup)
+
+    def train_step(state: TinyMLP, xb: NamedArray, yb: NamedArray):
+        loss, grads = _loss_and_grad(state, xb, yb)
+        new_state = _apply_sgd(state, grads, base_lr=base_lr, use_mup=use_mup)
+        return new_state, loss
+
+    data_axis = inputs.axes[0]
+    n = data_axis.size
+
+    state = params
+    for t in range(steps):
+        start = (t * batch_size) % n
+        end = start + batch_size
+        batch_idx = (data_axis, slice(start, end))
+        xb = inputs[batch_idx]
+        yb = targets[batch_idx]
+        state, _ = train_step(state, xb, yb)
+
+    final_loss = _loss_value(state, inputs, targets)
+    return float(final_loss)
+
+
+def _span(values: Iterable[float]) -> float:
+    seq = list(values)
+    return max(seq) - min(seq)
+
+
+def coord_check(
+    widths: tuple[int, ...] = (32, 128, 512),
+    *,
+    steps: int = 120,
+    base_lr: float = 3e-3,
+) -> dict[str, Any]:
+    seed = 0
+    keys = jrandom.split(jrandom.PRNGKey(seed), len(widths))
+    mup_losses = [
+        _run_once(key, width=width, use_mup=True, steps=steps, base_lr=base_lr) for key, width in zip(keys, widths)
+    ]
+    ctrl_losses = [
+        _run_once(key, width=width, use_mup=False, steps=steps, base_lr=base_lr) for key, width in zip(keys, widths)
+    ]
+
+    return {
+        "widths": list(widths),
+        "mup_losses": mup_losses,
+        "ctrl_losses": ctrl_losses,
+        "mup_span": _span(mup_losses),
+        "ctrl_span": _span(ctrl_losses),
+    }
+
+
+def test_mup_coordinate_check_is_width_invariant():
+    result = coord_check(widths=(32, 128, 512), steps=120, base_lr=3e-3)
+    mup_span = result["mup_span"]
+    ctrl_span = result["ctrl_span"]
+
+    if ctrl_span < 1e-5:
+        assert mup_span <= ctrl_span + 1e-6, f"μP not at least as invariant: {result}"
+    else:
+        assert mup_span <= 0.6 * ctrl_span, f"μP did not improve width invariance enough.\n{result}"