haliax 1.4.dev340__tar.gz → 1.4.dev342__tar.gz

{haliax-1.4.dev340 → haliax-1.4.dev342}/.github/workflows/run_quick_levanter_tests.yaml

@@ -17,19 +17,21 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install flake8 pytest
-          pip install --upgrade "jax[cpu]==0.4.35" "jaxlib[cpu]==0.4.35"
+          pip install "jax[cpu]==0.5.3" "jaxlib[cpu]==0.5.3" .[dev]
 
       - name: Install Levanter from source
         run: |
           cd ..
           git clone https://github.com/stanford-crfm/levanter.git
           cd levanter
-          pip install -e .
+          pip install -e .[tests]
+          pip install -r tests/requirements.txt
+          # i don't know why this is necessary
+          pip install tensorboardX
       - name: Install Haliax on top
         run: |
           # install second since levanter will install a built version of haliax
           cd ../haliax
-          pip install "jax[cpu]==0.4.35" "jaxlib[cpu]==0.4.35" .[dev]
       - name: Test levanter with pytest
         run: |
           cd ../levanter

{haliax-1.4.dev340 → haliax-1.4.dev342}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev340
+Version: 1.4.dev342
 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.dev340 → haliax-1.4.dev342}/docs/nn.md

@@ -18,7 +18,10 @@ of the same module (e.g. transformer blocks), which is a common pattern in deep
 ::: haliax.nn.Dropout
 
 ### Normalization
+
+::: haliax.nn.normalization.LayerNormBase
 ::: haliax.nn.LayerNorm
+::: haliax.nn.RmsNorm
 
 ### Meta
 

{haliax-1.4.dev340 → haliax-1.4.dev342}/docs/state-dict.md

@@ -122,6 +122,7 @@ For example, below is the implementation of `to_state_dict()` in [levanter.model
 In this class, we want to preserve HF compatibility by saving untied output embeddings. (We chose not to implement
 non-weight-tied embeddings.)
 
+
 ```python
 from typing import Optional
 
@@ -155,8 +156,69 @@ to join the prefix to the keys of the state dict.
 
 ```
 
+#### Flattening and Unflattening
+
+Haliax differs from many NN frameworks, and PyTorch in particular, in supporting multiple axes as inputs and outputs
+for linear transformations and layer norms. This means that a Linear layer's weight might have shape `(Heads, Dim, Out)`
+rather than just `(In, Out)` (where `In = Heads * Dim`). To facilitate compatibility with PyTorch, we provide
+two functions, `flatten_modules_for_export` and `unflatten_modules_from_export`, that can be used to convert
+modules to and from a format that is compatible with PyTorch. These functions are used internally by
+`to_torch_compatible_state_dict` and `from_torch_compatible_state_dict` and we expose them for advanced users.
+
+If you are adding a new module, you can plug into this system by inheriting from [haliax.state_dict.ModuleWithStateDictSerialization][]
+and overriding the `flatten_for_export` and `unflatten_from_export` methods. Here is an example from [haliax.nn.LayerNorm][]:
+
+```python
+Mod = TypeVar("Mod")
+
+class LayerNormBase(ModuleWithStateDictSerialization):
+    def flatten_for_export(self: Mod) -> Mod:
+        if isinstance(self.axis, hax.Axis):
+            return self
+
+        if self.weight is not None:
+            weight = self.weight.flatten("__OUT")
+        else:
+            weight = None
+
+        if self.bias is not None:
+            bias = self.bias.flatten("__OUT")
+        else:
+            bias = None
+
+        return dataclasses.replace(
+            self, weight=weight, bias=bias, axis=hax.flatten_axes(self.axis, "__OUT")
+        )
+
+    def unflatten_from_export(self: Mod, template: Mod) -> Mod:
+        if template.axis == self.axis:
+            return self
+
+        if self.weight is not None:
+            assert isinstance(self.axis, hax.Axis), "Cannot unflatten weight with non-axis axis"
+            weight = hax.unflatten_axis(self.weight, self.axis, template.axis)
+        else:
+            weight = None
+
+        if self.bias is not None:
+            assert isinstance(self.axis, hax.Axis), "Cannot unflatten weight with non-axis axis"
+            bias = hax.unflatten_axis(self.bias, self.axis, template.axis)
+
+        else:
+            bias = None
+
+        return dataclasses.replace(
+            self, weight=weight, bias=bias, axis=template.axis
+        )
+```
+
+The code is a bit boilerplate-y but the idea is to find articulated axes in arrays and flatten them, while updating
+any Axis members to match the new shape.
+
 ## API Reference
 
+### Types
+
 ::: haliax.state_dict.StateDict
 ::: haliax.state_dict.ModuleWithStateDictSerialization
 
@@ -173,5 +235,5 @@ to join the prefix to the keys of the state dict.
 
 ::: haliax.state_dict.from_torch_compatible_state_dict
 ::: haliax.state_dict.to_torch_compatible_state_dict
-::: haliax.state_dict.flatten_linear_layers
-::: haliax.state_dict.unflatten_linear_layers
+::: haliax.state_dict.flatten_modules_for_export
+::: haliax.state_dict.unflatten_modules_from_export

haliax-1.4.dev342/src/haliax/__about__.py

@@ -0,0 +1 @@
+__version__ = "1.4.dev342"

{haliax-1.4.dev340 → haliax-1.4.dev342}/src/haliax/_src/state_dict.py

@@ -18,6 +18,7 @@ from haliax._src.util import index_where
 from haliax.axis import Axis
 from haliax.core import NamedArray, flatten_axes, named
 from haliax.jax_utils import is_jax_array_like, is_scalarish
+from haliax.tree_util import scan_aware_tree_map
 
 
 try:
@@ -32,14 +33,16 @@ T = TypeVar("T")
 
 
 def from_torch_compatible_state_dict(
-    t: T, state_dict: StateDict, *, unflatten_linear: bool = True, prefix: Optional[str] = None
+    t: T, state_dict: StateDict, *, unflatten: bool = True, prefix: Optional[str] = None
 ) -> T:
     """
     Convert a state dict to a tree that is compatible with the structure of `t`.
 
-    This applies [haliax.state_dict.from_state_dict][] followed by [haliax.state_dict.unflatten_linear_layers][].
+    If unflatten is true, then the weights in the state dict are assumed to have been flattened (as by flatten_modules_for_export).
+
+    This applies [haliax.state_dict.from_state_dict][]
     """
-    if unflatten_linear:
+    if unflatten:
         t = _flatten_to_unflatten(t, state_dict, prefix)
     else:
         t = from_state_dict(t, state_dict, prefix=prefix)
@@ -47,6 +50,45 @@ def from_torch_compatible_state_dict(
     return t
 
 
+def flatten_modules_for_export(t: T) -> T:
+    """
+    Flatten all modules in a tree for export to torch.
+    """
+
+    def _flatten_module(module):
+        if isinstance(module, ModuleWithStateDictSerialization):
+            module = module.flatten_for_export()
+            module = scan_aware_tree_map(
+                _flatten_module,
+                module,
+                is_leaf=lambda x: x is not module and isinstance(x, ModuleWithStateDictSerialization),
+            )
+        return module
+
+    return scan_aware_tree_map(_flatten_module, t, is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization))
+
+
+def unflatten_modules_from_export(t: T, template: T) -> T:
+    """
+    Unflatten all modules in a tree after import from torch.
+    """
+
+    def _unflatten_module(module, template):
+        if isinstance(module, ModuleWithStateDictSerialization):
+            module = module.unflatten_from_export(template)
+            module = scan_aware_tree_map(
+                _unflatten_module,
+                module,
+                template,
+                is_leaf=lambda x: x is not module and isinstance(x, ModuleWithStateDictSerialization),
+            )
+        return module
+
+    return scan_aware_tree_map(
+        _unflatten_module, t, template, is_leaf=lambda x: isinstance(x, ModuleWithStateDictSerialization)
+    )
+
+
 def _flatten_to_unflatten(t, state_dict, prefix):
     """
     Flatten the torch compatible state_dict before loading into t, and then recover the unflattened layers.
@@ -59,9 +101,9 @@ def _flatten_to_unflatten(t, state_dict, prefix):
         return jnp.zeros(struct.shape, struct.dtype)
 
     t = jax.tree.map(_dt_struct_to_array, t)
-    flat_t = flatten_linear_layers(t)
+    flat_t = flatten_modules_for_export(t)
     flat_t = from_state_dict(flat_t, state_dict, prefix=prefix)
-    t = unflatten_linear_layers(t, flat_t)
+    t = unflatten_modules_from_export(flat_t, t)
     return t
 
 
@@ -103,6 +145,25 @@ class ModuleWithStateDictSerialization(eqx.Module):
         """Returns a dict mapping eqx.Module keys to torch keys that need to be renamed for serialization"""
         return {}
 
+    def flatten_for_export(self: Mod) -> Mod:
+        """
+        Flatten articulated named arrays for export to torch. In general this method should, for a linear layer, flatten
+        all input axes into a single axis, and all output axes into a single axis. You can do whatever else
+        you want to support pytorch-compatible serialization if you want.
+
+        This method is less general than to_state_dict and is only called when using to_torch_compatible_state_dict.
+        """
+        return self
+
+    def unflatten_from_export(self: Mod, template: Mod) -> Mod:
+        """
+        Unflatten the module after import from torch.
+
+        Template has the proper structure (e.g. articulated named axes) but the values are meaningless.
+        """
+        del template
+        return self
+
 
 def from_state_dict(tree: T, state_dict: StateDict, prefix: Optional[str] = None) -> T:
     """
@@ -372,78 +433,3 @@ def load_state_dict(path):
     """
     state_dict = safetensors.numpy.load_file(path)
     return state_dict
-
-
-def flatten_linear_layers(tree: T) -> T:
-    """
-    In PyTorch, linear layers are stored as a 2d weight matrix and a 1d bias vector. In Haliax,
-    linear layers can have arbitrary dimensions, grouped into input and output axes. This function
-    flattens the linear layers in a tree to be compatible with PyTorch-style state dicts.
-
-    :param tree:
-    """
-    from haliax.nn import Linear
-
-    def _flatten_linear(layer):
-        if not isinstance(layer, Linear):
-            return layer
-
-        weight = layer.weight
-        bias = layer.bias
-
-        new_Out: Axis = flatten_axes(layer.Out, "__OUT__")
-        new_In: Axis = flatten_axes(layer.In, "__IN__")
-
-        if weight.array is not None:
-            out_first = layer._out_first
-            weight = weight.flatten_axes(layer.Out, new_Out).flatten_axes(layer.In, new_In)
-
-            if out_first:
-                weight = weight.rearrange((..., "__OUT__", "__IN__"))
-            else:
-                weight = weight.rearrange((..., "__IN__", "__OUT__"))
-
-        if isinstance(bias, NamedArray):
-            bias = bias.flatten_axes(layer.Out, new_Out)
-
-        return dataclasses.replace(layer, weight=weight, bias=bias, In=new_In, Out=new_Out)  # type: ignore
-
-    return jax.tree.map(_flatten_linear, tree, is_leaf=lambda x: isinstance(x, Linear))
-
-
-def unflatten_linear_layers(template: T, tree_with_flattened_linears: T) -> T:
-    """
-    Unflattens linear layers in a tree that was flattened with [haliax.state_dict.flatten_linear_layers][].
-    Template has the same structure as the tree that was flattened, but with the original (unflattened)
-    linear layers.
-
-    Returns:
-        The same tree as `tree_with_flattened_linears`, but with the linear layers unflattened to match
-        the structure of `template`.
-    """
-
-    from haliax.nn import Linear
-
-    def _unflatten_linear(template, flattened):
-        assert isinstance(template, Linear) == isinstance(flattened, Linear)
-
-        if not isinstance(template, Linear):
-            return flattened
-
-        weight = flattened.weight
-        bias = flattened.bias
-
-        if weight.array is not None:
-            weight = weight.unflatten_axis("__OUT__", template.Out).unflatten_axis("__IN__", template.In)
-            weight = weight.rearrange(template.weight.axes)
-
-        if isinstance(bias, NamedArray):
-            bias = bias.unflatten_axis("__OUT__", template.Out)
-            assert template.bias is not None, "Flattened bias but template has no bias"
-            bias = bias.rearrange(template.bias.axes)
-
-        return dataclasses.replace(template, weight=weight, bias=bias)  # type: ignore
-
-    return jax.tree.map(
-        _unflatten_linear, template, tree_with_flattened_linears, is_leaf=lambda x: isinstance(x, Linear)
-    )

{haliax-1.4.dev340 → haliax-1.4.dev342}/src/haliax/nn/__init__.py

@@ -36,14 +36,11 @@ from .embedding import Embedding
 from .linear import Linear
 from .loss import binary_cross_entropy_loss, cross_entropy_loss, cross_entropy_loss_and_log_normalizers, reduce_loss
 from .mlp import MLP
-from .normalization import LayerNorm, log_softmax, logsumexp, softmax, standardize
+from .normalization import LayerNorm, RmsNorm, log_softmax, logsumexp, softmax, standardize
 from .pool import max_pool, mean_pool, min_pool
 from .scan import BlockSeq, Stacked
 
 
-# TODO: support where in softmax, etc
-
-
 def one_hot(x: NamedArray | int, class_axis: Axis, *, dtype=None) -> NamedArray:
     """
     Convert an integer to a one-hot vector. This is basically a generalization of [jax.nn.one_hot][]
@@ -87,6 +84,7 @@ __all__ = [
     "LayerNorm",
     "Linear",
     "Embedding",
+    "RmsNorm",
     "Stacked",
     "BlockSeq",
     "MLP",

{haliax-1.4.dev340 → haliax-1.4.dev342}/src/haliax/nn/linear.py

@@ -1,3 +1,4 @@
+import dataclasses
 import math
 from typing import Optional
 
@@ -6,13 +7,14 @@ from jax.random import PRNGKey
 
 import haliax as hax
 
+from .._src.state_dict import Mod, ModuleWithStateDictSerialization
 from ..axis import AxisSpec
 from ..core import NamedArray
 from ..jax_utils import named_call
 from ..quantization import DotGeneralOp
 
 
-class Linear(eqx.Module):
+class Linear(ModuleWithStateDictSerialization):
     """A named Linear layer. This module allows you to specify multiple named axes for both input
     and output, which is occasionally useful."""
 
@@ -71,6 +73,47 @@ class Linear(eqx.Module):
 
         return q
 
+    def flatten_for_export(self: Mod) -> Mod:
+        if isinstance(self.Out, hax.Axis) and isinstance(self.In, hax.Axis):
+            return self
+
+        weight = self.weight
+        bias = self.bias
+
+        new_Out = hax.flatten_axes(self.Out, "__OUT__")
+        new_In = hax.flatten_axes(self.In, "__IN__")
+
+        if weight.array is not None:
+            out_first = self._out_first
+            weight = weight.flatten_axes(self.Out, new_Out).flatten_axes(self.In, new_In)
+
+            if out_first:
+                weight = weight.rearrange((..., "__OUT__", "__IN__"))
+            else:
+                weight = weight.rearrange((..., "__IN__", "__OUT__"))
+
+        if isinstance(bias, NamedArray):
+            bias = bias.flatten_axes(self.Out, new_Out)
+
+        return dataclasses.replace(self, weight=weight, bias=bias, In=new_In, Out=new_Out)
+
+    def unflatten_from_export(self: Mod, template: Mod) -> Mod:
+        weight = self.weight
+        bias = self.bias
+
+        if (template.In, template.Out) == (self.In, self.Out):
+            return self
+
+        if weight.array is not None:
+            weight = weight.unflatten_axis("__OUT__", template.Out).unflatten_axis("__IN__", template.In)
+            weight = weight.rearrange(template.weight.axes)
+
+        if isinstance(bias, NamedArray):
+            bias = bias.unflatten_axis("__OUT__", template.Out)
+            bias = bias.rearrange(template.bias.axes)
+
+        return dataclasses.replace(template, weight=weight, bias=bias)
+
     @property
     def _out_first(self):
         """

haliax-1.4.dev342/src/haliax/nn/normalization.py

@@ -0,0 +1,169 @@
+import dataclasses
+from abc import abstractmethod
+from typing import Optional, TypeVar
+
+import equinox as eqx
+from jax import nn as jnn
+from jax import numpy as jnp
+
+import haliax
+import haliax as hax
+
+from .._src.state_dict import Mod, ModuleWithStateDictSerialization
+from ..axis import AxisSelection, AxisSpec
+from ..core import NamedArray
+from ..types import Scalar
+from ..wrap import unwrap_namedarrays, wrap_axiswise_call, wrap_reduction_call
+
+
+A = TypeVar("A", Scalar, NamedArray, jnp.ndarray)
+
+
+class LayerNormBase(ModuleWithStateDictSerialization):
+    axis: AxisSpec = eqx.static_field()
+    weight: Optional[NamedArray]
+    bias: Optional[NamedArray]
+    eps: float = eqx.static_field(default=1e-5)
+    dtype: Optional[jnp.dtype] = eqx.field(default=None, static=True)
+
+    @abstractmethod
+    def __call__(self, x: NamedArray) -> NamedArray:
+        pass
+
+    @classmethod
+    def init(
+        cls,
+        axis: AxisSpec,
+        eps: float = 1e-5,
+        *,
+        use_weight: bool = True,
+        use_bias: bool = True,
+        dtype: Optional[jnp.dtype] = None,
+    ):
+        if use_weight:
+            weight = hax.ones(axis)
+        else:
+            weight = None
+
+        if use_bias:
+            bias = hax.zeros(axis)
+        else:
+            bias = None
+
+        return cls(axis, weight, bias, eps, dtype)
+
+    def flatten_for_export(self: Mod) -> Mod:
+        if isinstance(self.axis, hax.Axis):
+            return self
+
+        if self.weight is not None:
+            weight = self.weight.flatten("__OUT")
+        else:
+            weight = None
+
+        if self.bias is not None:
+            bias = self.bias.flatten("__OUT")
+        else:
+            bias = None
+
+        return dataclasses.replace(self, weight=weight, bias=bias, axis=hax.flatten_axes(self.axis, "__OUT"))
+
+    def unflatten_from_export(self: Mod, template: Mod) -> Mod:
+        if template.axis == self.axis:
+            return self
+
+        if self.weight is not None:
+            assert isinstance(self.axis, hax.Axis), "Cannot unflatten weight with non-axis axis"
+            weight = hax.unflatten_axis(self.weight, self.axis, template.axis)
+        else:
+            weight = None
+
+        if self.bias is not None:
+            assert isinstance(self.axis, hax.Axis), "Cannot unflatten weight with non-axis axis"
+            bias = hax.unflatten_axis(self.bias, self.axis, template.axis)
+
+        else:
+            bias = None
+
+        return dataclasses.replace(self, weight=weight, bias=bias, axis=template.axis)
+
+
+class LayerNorm(LayerNormBase):
+    r"""
+    Normalises the input along the specified axis (or axes), using the mean and variance of the
+    input along that axis.
+    """
+    axis: AxisSpec = eqx.field(static=True)
+    weight: Optional[NamedArray]
+    bias: Optional[NamedArray]
+
+    eps: float = eqx.field(default=1e-5, static=True)
+    dtype: Optional[jnp.dtype] = eqx.field(default=None, static=True)
+
+    def __call__(self, x: NamedArray) -> NamedArray:
+        dtype = x.dtype
+        mean = x.mean(self.axis)
+        var = x.var(self.axis)
+        inv = hax.rsqrt(var + self.eps)
+        out = (x - mean) * inv
+        out = out.astype(dtype)
+
+        if self.weight is not None:
+            out = self.weight * out
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+
+class RmsNorm(LayerNormBase):
+    r"""
+    Implements RMS normalization, which normalizes the input by dividing by the root mean square of the input.
+    """
+
+    def __call__(self, x: NamedArray) -> NamedArray:
+        in_dtype = x.dtype
+        x = x.astype(self.dtype)
+        var = hax.mean(hax.square(x), axis=self.axis)
+        inv = hax.rsqrt(var + self.eps)
+        out = x * inv
+        out = out.astype(in_dtype)
+
+        if self.weight is not None:
+            out = self.weight * out
+        if self.bias is not None:
+            out = out + self.bias
+        return out
+
+
+def logsumexp(a: A, axis: Optional[AxisSelection] = None) -> A:
+    # TODO: logsumexp indirectly supports where via `b`. we should support it directly
+    return wrap_reduction_call(jnn.logsumexp, a, axis=axis, single_axis_only=False, supports_where=False)
+
+
+# TODO: support where in softmax, etc
+
+
+def softmax(a: A, axis: Optional[AxisSelection] = None) -> A:
+    return wrap_axiswise_call(jnn.softmax, a, axis=axis, single_axis_only=False)
+
+
+def log_softmax(a: A, axis: Optional[AxisSelection] = None) -> A:
+    return wrap_axiswise_call(jnn.log_softmax, a, axis=axis, single_axis_only=False)
+
+
+def standardize(
+    x: NamedArray,
+    axis: AxisSpec,
+    *,
+    mean: Optional[NamedArray] = None,
+    variance: Optional[NamedArray] = None,
+    epsilon: float = 1e-5,
+    where: Optional[NamedArray] = None,
+) -> NamedArray:
+    """Analogous to [jax.nn.standardize][], but with support for NamedArrays."""
+    x, mean, variance, where = haliax.broadcast_arrays(x, mean, variance, where)  # type: ignore
+    raw_x, mean, variance, where = unwrap_namedarrays(x, mean, variance, where)
+    axis_indices = x._lookup_indices(axis)
+
+    plain = jnn.standardize(raw_x, axis_indices, mean=mean, variance=variance, epsilon=epsilon, where=where)
+    return NamedArray(plain, x.axes)

{haliax-1.4.dev340 → haliax-1.4.dev342}/src/haliax/state_dict.py

@@ -8,14 +8,14 @@ from haliax.types import FilterSpec
 from ._src.state_dict import (
     ModuleWithStateDictSerialization,
     StateDict,
-    flatten_linear_layers,
+    flatten_modules_for_export,
     from_state_dict,
     from_torch_compatible_state_dict,
     load_state_dict,
     save_state_dict,
     to_numpy_state_dict,
     to_state_dict,
-    unflatten_linear_layers,
+    unflatten_modules_from_export,
     with_prefix,
 )
 
@@ -24,23 +24,23 @@ T = TypeVar("T")
 
 
 def to_torch_compatible_state_dict(
-    t: T, *, flatten_linear: bool = True, prefix: Optional[str] = None, filter: FilterSpec = is_jax_array_like
+    t: T, *, flatten: bool = True, prefix: Optional[str] = None, filter: FilterSpec = is_jax_array_like
 ) -> StateDict:
     """
     Convert a tree to a state dict that is compatible with torch-style state dicts.
 
-    This applies [haliax.state_dict.flatten_linear_layers][] followed by [haliax.state_dict.to_state_dict][]
+    This applies the same logic as [to_state_dict][] but also uses [haliax.state_dict.ModuleWithStateDictSerialization.flatten_for_export][] to flatten
 
     Args:
         t: The tree to convert
-        flatten_linear: Whether to flatten linear layers
+        flatten: Whether to flatten axes using flatten_for_export
         prefix: The prefix to use for the state dict keys
         filter: The filter to use for selecting which nodes to include in the state dict. By default, this includes only
             array-like objects (e.g. JAX and NumPy arrays).
     """
     t = equinox.filter(t, filter)
-    if flatten_linear:
-        t = flatten_linear_layers(t)
+    if flatten:
+        t = flatten_modules_for_export(t)
     return to_numpy_state_dict(t, prefix=prefix)
 
 
@@ -50,11 +50,11 @@ __all__ = [
     "load_state_dict",
     "save_state_dict",
     "from_state_dict",
-    "flatten_linear_layers",
-    "unflatten_linear_layers",
     "with_prefix",
     "to_state_dict",
     "to_numpy_state_dict",
     "StateDict",
     "to_torch_compatible_state_dict",
+    "flatten_modules_for_export",
+    "unflatten_modules_from_export",
 ]

{haliax-1.4.dev340 → haliax-1.4.dev342}/src/haliax/tree_util.py

@@ -30,17 +30,10 @@ def tree_map(fn, tree, *rest, is_leaf=None):
 
 def scan_aware_tree_map(fn, tree, *rest, is_leaf=None):
     """
-    Version of [haliax.tree_util.tree_map][] that is aware of the scan-layer pattern, specifically as implmeneted
+    Version of [haliax.tree_util.tree_map][] that is aware of the scan-layer pattern, specifically as implemented
     in hax.nn.Stacked. This function will (implicitly) apply the transform to each layer in each Stacked module
     (using vmap). If there are no Stacked modules in the tree, this function is equivalent to [haliax.tree_util.tree_map][].
 
-    Args:
-        fn:
-        tree:
-        *rest:
-        is_leaf:
-
-    Returns:
     """
     old_is_leaf = is_leaf
     if is_leaf is None:
@@ -55,7 +48,7 @@ def scan_aware_tree_map(fn, tree, *rest, is_leaf=None):
             new_inner = haliax.vmap(mapped_fn, x.Block)(x.stacked, *[r.stacked for r in rest])
             return dataclasses.replace(x, stacked=new_inner)  # type: ignore
         else:
-            return fn(x)
+            return fn(x, *rest)
 
     return tree_map(rec_fn, tree, *rest, is_leaf=is_leaf)
 

{haliax-1.4.dev340 → haliax-1.4.dev342}/tests/test_state_dict.py

@@ -1,3 +1,4 @@
+import dataclasses
 from typing import Any
 
 import equinox as eqx
@@ -6,9 +7,10 @@ import jax.numpy as jnp
 import pytest
 
 import haliax as hax
+from haliax._src.state_dict import flatten_modules_for_export, unflatten_modules_from_export
 from haliax.nn import Linear
-from haliax.nn.scan import _stack_state_dict, _unstack_state_dict
-from haliax.state_dict import flatten_linear_layers, from_state_dict, to_state_dict, unflatten_linear_layers
+from haliax.nn.scan import Stacked, _stack_state_dict, _unstack_state_dict
+from haliax.state_dict import from_state_dict, to_state_dict
 
 
 @pytest.mark.parametrize("out_dims_first", [True, False])
@@ -24,7 +26,7 @@ def test_flatten_linear_layers(out_dims_first: bool):
     else:
         assert linear.weight.axes == (H, W, D, B)
 
-    flat_linear = flatten_linear_layers(linear)
+    flat_linear = linear.flatten_for_export()
 
     flat_state_dict = to_state_dict(flat_linear)
     if out_dims_first:
@@ -36,7 +38,7 @@ def test_flatten_linear_layers(out_dims_first: bool):
 
     # now unflatten it
     linear2 = Linear.init((H, W), (D, B), key=jax.random.PRNGKey(1), use_bias=True, out_first=out_dims_first)
-    new_linear = unflatten_linear_layers(linear2, flat_linear)
+    new_linear = flat_linear.unflatten_from_export(linear2)
 
     if out_dims_first:
         assert new_linear.weight.axes == (D, B, H, W)
@@ -127,3 +129,49 @@ def test_to_from_state_dict():
     m2 = from_state_dict(m2, state_dict)
     assert jnp.all(m2.a == a)
     assert jnp.all(m2.b == b)
+
+
+def test_export_layer_norm():
+    D = hax.Axis("D", 10)
+    E = hax.Axis("E", 20)
+    layer_norm = hax.nn.LayerNorm.init((D, E), eps=1e-5, use_weight=True, use_bias=True)
+
+    flat_layer_norm = layer_norm.flatten_for_export()
+
+    flat_state_dict = to_state_dict(flat_layer_norm)
+
+    assert flat_state_dict["weight"].shape == (D.size * E.size,)
+    assert flat_state_dict["bias"].shape == (D.size * E.size,)
+    assert flat_state_dict["weight"].dtype == flat_state_dict["bias"].dtype == layer_norm.weight.dtype
+
+    # now unflatten it
+    layer_norm2 = hax.nn.LayerNorm.init((D, E), eps=1e-5, use_weight=True, use_bias=True)
+    # ensure we have different weights
+    layer_norm2 = dataclasses.replace(layer_norm2, weight=layer_norm2.weight + 1, bias=layer_norm2.bias + 1)
+
+    new_layer_norm = flat_layer_norm.unflatten_from_export(layer_norm2)
+
+    assert layer_norm == new_layer_norm
+
+
+def test_stacked_layer_norm():
+    L = hax.Axis("L", 4)
+    D = hax.Axis("D", 10)
+    E = hax.Axis("E", 20)
+
+    norms = Stacked.init(L, hax.nn.LayerNorm)((D, E), eps=1e-5, use_weight=True, use_bias=True)
+
+    norms_flat = flatten_modules_for_export(norms)
+
+    flat_state_dict = to_state_dict(norms_flat)
+
+    assert flat_state_dict["0.weight"].shape == (D.size * E.size,)
+    assert flat_state_dict["0.bias"].shape == (D.size * E.size,)
+    assert flat_state_dict["1.weight"].shape == (D.size * E.size,)
+
+    # now unflatten it
+    norms2 = Stacked.init(L, hax.nn.LayerNorm)((D, E), eps=1e-5, use_weight=True, use_bias=True)
+
+    new_norms = unflatten_modules_from_export(norms_flat, norms2)
+
+    assert norms == new_norms

haliax-1.4.dev340/src/haliax/__about__.py

@@ -1 +0,0 @@
-__version__ = "1.4.dev340"

haliax-1.4.dev340/src/haliax/nn/normalization.py

@@ -1,87 +0,0 @@
-from typing import Optional, TypeVar
-
-import equinox as eqx
-from jax import nn as jnn
-from jax import numpy as jnp
-
-import haliax
-import haliax as hax
-
-from ..axis import AxisSelection, AxisSpec
-from ..core import NamedArray
-from ..types import Scalar
-from ..wrap import unwrap_namedarrays, wrap_axiswise_call, wrap_reduction_call
-
-
-A = TypeVar("A", Scalar, NamedArray, jnp.ndarray)
-
-
-class LayerNorm(eqx.Module):
-    r"""
-    Normalises the input along the specified axis (or axes), using the mean and variance of the
-    input along that axis.
-    """
-    axis: AxisSpec = eqx.static_field()
-    weight: Optional[NamedArray]
-    bias: Optional[NamedArray]
-
-    eps: float = eqx.static_field(default=1e-5)
-
-    @staticmethod
-    def init(axis: AxisSpec, eps: float = 1e-5, use_weight: bool = True, use_bias: bool = True):
-        if use_weight:
-            weight = hax.ones(axis)
-        else:
-            weight = None
-        if use_bias:
-            bias = hax.zeros(axis)
-        else:
-            bias = None
-
-        return LayerNorm(axis, weight, bias, eps)
-
-    def __call__(self, x: NamedArray) -> NamedArray:
-        mean = x.mean(self.axis)
-        var = x.var(self.axis)
-        inv = hax.rsqrt(var + self.eps)
-        out = (x - mean) * inv
-
-        if self.weight is not None:
-            out = self.weight * out
-        if self.bias is not None:
-            out = out + self.bias
-        return out
-
-
-def logsumexp(a: A, axis: Optional[AxisSelection] = None) -> A:
-    # TODO: logsumexp indirectly supports where via `b`. we should support it directly
-    return wrap_reduction_call(jnn.logsumexp, a, axis=axis, single_axis_only=False, supports_where=False)
-
-
-# TODO: support where in softmax, etc
-
-
-def softmax(a: A, axis: Optional[AxisSelection] = None) -> A:
-    return wrap_axiswise_call(jnn.softmax, a, axis=axis, single_axis_only=False)
-
-
-def log_softmax(a: A, axis: Optional[AxisSelection] = None) -> A:
-    return wrap_axiswise_call(jnn.log_softmax, a, axis=axis, single_axis_only=False)
-
-
-def standardize(
-    x: NamedArray,
-    axis: AxisSpec,
-    *,
-    mean: Optional[NamedArray] = None,
-    variance: Optional[NamedArray] = None,
-    epsilon: float = 1e-5,
-    where: Optional[NamedArray] = None,
-) -> NamedArray:
-    """Analogous to [jax.nn.standardize][], but with support for NamedArrays."""
-    x, mean, variance, where = haliax.broadcast_arrays(x, mean, variance, where)  # type: ignore
-    raw_x, mean, variance, where = unwrap_namedarrays(x, mean, variance, where)
-    axis_indices = x._lookup_indices(axis)
-
-    plain = jnn.standardize(raw_x, axis_indices, mean=mean, variance=variance, epsilon=epsilon, where=where)
-    return NamedArray(plain, x.axes)