haliax 1.4.dev438__tar.gz → 1.4.dev441__tar.gz

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

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

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

{haliax-1.4.dev438 → haliax-1.4.dev441}/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.dev438 → haliax-1.4.dev441}/src/haliax/nn/mup.py

@@ -168,7 +168,7 @@ class EmbeddingStandardParam(AbstractEmbeddingReparam):
 
     @property
     def unembed_active_scale(self):
-        return 1 / hax.axis_size(self.Embed)
+        return 1
 
 
 class EmbeddingMup(AbstractEmbeddingReparam):

haliax-1.4.dev441/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",
+]