haliax 1.4.dev420__tar.gz → 1.4.dev438__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev420
+Version: 1.4.dev438
 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.dev438}/src/haliax/__about__.py

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

{haliax-1.4.dev420 → haliax-1.4.dev438}/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.dev438}/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.dev438/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.dev438/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}"

haliax-1.4.dev438/tests/test_mup_embedding.py

@@ -0,0 +1,48 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import jax.numpy as jnp
+import jax.random as jrandom
+import pytest
+
+import haliax as hax
+from haliax.nn import Embedding
+from haliax.nn.mup import EmbeddingMup
+
+
+@pytest.mark.parametrize("init_scale", [0.5, 1.0])
+def test_mup_embedding_init_matches_embedding(init_scale: float):
+    Vocab = hax.Axis("V", 8)
+    Embed = (hax.Axis("E", 4),)
+
+    key = jrandom.PRNGKey(0)
+
+    baseline = Embedding.init(Vocab, Embed, key=key, init_scale=init_scale)
+    mup = Embedding.init(Vocab, Embed, key=key, init_scale=init_scale, reparam_cls=EmbeddingMup)
+
+    assert mup.weight.axes == baseline.weight.axes
+    scale_factor = hax.axis_size(Embed)
+    assert jnp.allclose(mup.weight.array, baseline.weight.array * scale_factor)
+
+
+def test_mup_embedding_unembedding_scale():
+    Vocab = hax.Axis("V", 6)
+    Embed = (hax.Axis("E", 3),)
+
+    weight = hax.ones(hax.concat_axis_specs(Vocab, Embed))
+    layer = Embedding(weight=weight, Vocab=Vocab, Embed=Embed, reparam=EmbeddingMup(Embed, Vocab))
+
+    scale = layer.reparam.unembed_active_scale
+    assert scale == pytest.approx(1.0 / hax.axis_size(Embed))
+    assert layer.reparam.active_scale == pytest.approx(1.0)
+    assert jnp.allclose((layer.weight * scale).array, weight.array * scale)
+
+    Batch = hax.Axis("B", 2)
+    inputs = hax.ones((Batch, *Embed))
+
+    logits = layer.unembed(inputs)
+    expected = hax.ones((Batch, Vocab))
+
+    assert logits.axes == expected.axes
+    assert jnp.allclose(logits.array, expected.array)

haliax-1.4.dev438/tests/test_mup_linear.py

@@ -0,0 +1,120 @@
+# Copyright 2025 The Levanter Authors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+
+import dataclasses
+
+import jax.numpy as jnp
+import jax.random as jrandom
+import pytest
+
+import haliax as hax
+from haliax.nn import Linear
+from haliax.nn.mup import InputLinearMup, LinearStandardParam, HiddenLinearMup, OutputLinearMup
+
+
+@pytest.mark.parametrize("out_first", [True, False])
+def test_mup_linear_init_matches_linear_axes(out_first: bool):
+    In = (hax.Axis("I", 3), hax.Axis("J", 2))
+    Out = (hax.Axis("O", 5),)
+    key_linear, key_mup = jrandom.split(jrandom.PRNGKey(0))
+
+    linear = Linear.init(In, Out, key=key_linear, out_first=out_first)
+    mup = Linear.init(In, Out, key=key_mup, out_first=out_first, reparam_cls=InputLinearMup)
+
+    assert linear.weight.axes == mup.weight.axes
+    if linear.bias is not None:
+        assert mup.bias is not None
+        assert linear.bias.axes == mup.bias.axes
+
+
+def test_mup_linear_call_matches_linear():
+    Batch = hax.Axis("B", 2)
+    In = (hax.Axis("I", 3),)
+    Out = (hax.Axis("O", 4),)
+
+    weight = hax.ones(hax.concat_axis_specs(Out, In)) * 0.5
+    bias = hax.full(Out, 0.25)
+
+    linear = Linear(weight, bias, In, Out, reparam=LinearStandardParam(In, Out))
+    mup = Linear(weight, bias, In, Out, reparam=InputLinearMup(In, Out))
+
+    inputs = hax.full(hax.concat_axis_specs(Batch, In), 2.0)
+
+    expected = linear(inputs)
+    actual = mup(inputs)
+
+    assert actual.axes == expected.axes
+    assert jnp.allclose(actual.array, expected.array)
+
+
+@pytest.mark.parametrize("out_first", [True, False])
+def test_hidden_linear_init_matches_linear_scaling(out_first: bool):
+    In = hax.Axis("I", 6)
+    Out = hax.Axis("O", 5)
+    key = jrandom.PRNGKey(0)
+
+    linear = Linear.init(In, Out, key=key, use_bias=False, out_first=out_first)
+    hidden = Linear.init(
+        In,
+        Out,
+        key=key,
+        use_bias=False,
+        out_first=out_first,
+        reparam_cls=HiddenLinearMup,
+    )
+
+    assert jnp.allclose(hidden.weight.array, linear.weight.array)
+
+
+def test_output_linear_scales_activation_by_input_size():
+    Batch = hax.Axis("B", 2)
+    In = hax.Axis("I", 3)
+    Out = hax.Axis("O", 2)
+
+    layer = Linear.init(In, Out, key=jrandom.PRNGKey(2), use_bias=False, reparam_cls=OutputLinearMup)
+    layer = dataclasses.replace(layer, weight=hax.ones(layer.weight.axes))
+
+    inputs = hax.full((Batch, In), 2.0)
+    expected = hax.full(hax.concat_axis_specs(Batch, layer.Out), 2.0)
+
+    actual = layer(inputs)
+
+    assert actual.axes == expected.axes
+    assert jnp.allclose(actual.array, expected.array)
+
+
+def test_output_linear_state_dict_scales_weight():
+    In = hax.Axis("I", 4)
+    Out = hax.Axis("O", 3)
+
+    layer = Linear.init(In, Out, key=jrandom.PRNGKey(3), use_bias=False, reparam_cls=OutputLinearMup)
+    layer = dataclasses.replace(layer, weight=hax.ones(layer.weight.axes))
+
+    state = layer.to_state_dict()
+    assert jnp.allclose(state["weight"], layer.weight.array * layer.reparam.active_scale)
+
+    template = Linear.init(In, Out, key=jrandom.PRNGKey(4), use_bias=False, reparam_cls=OutputLinearMup)
+    restored = template.from_state_dict(state)
+
+    assert jnp.allclose(restored.weight.array, layer.weight.array)
+
+
+def test_input_linear_behaves_like_base_linear():
+    Batch = hax.Axis("B", 2)
+    In = hax.Axis("I", 3)
+    Out = hax.Axis("O", 4)
+
+    weight = hax.ones((Out, In)) * 0.1
+    bias = hax.zeros(Out)
+
+    linear = Linear(weight, bias, In, Out, reparam=LinearStandardParam(In, Out))
+    input_linear = Linear(weight, bias, In, Out, reparam=InputLinearMup(In, Out))
+
+    inputs = hax.random.normal(jrandom.PRNGKey(5), (Batch, In))
+
+    expected = linear(inputs)
+    actual = input_linear(inputs)
+
+    assert jnp.allclose(actual.array, expected.array)