haliax 1.4.dev445__tar.gz → 1.4.dev447__tar.gz

haliax-1.4.dev447/.agents/refs.md

@@ -0,0 +1,50 @@
+# NamedRef & SliceRef Plan
+
+## Goals
+- Provide a `NamedRef` abstraction mirroring `jax.Ref` while preserving Haliax axis metadata.
+- Offer helper APIs (`hax.new_ref`, `hax.freeze`, `hax.ref.get/swap`) that operate on named refs and enforce shape/axis compatibility.
+- Enable "slice refs" that partially apply named indexers so subsequent operations only supply the remaining indices, including composing multiple slice refs.
+
+## Architectural Outline
+1. **Core data model**
+   - Introduce `NamedRef` in `src/haliax/core.py` (or a focused `src/haliax/ref.py` module) as a thin wrapper around `jax.Ref` plus axis metadata (likely reuse `NamedArray` named shape helpers).
+   - Register `NamedRef` as a PyTree node: treat the underlying `jax.Ref` handle as a leaf to avoid unwanted tracing semantics.
+   - Expose properties: `.axes`, `.named_shape`, `.shape`, `.dtype`, and `.value` (reads through `ref[...]`).
+
+2. **Creation and conversion helpers**
+   - Implement `haliax.new_ref(value, axes|named_shape)` returning `NamedRef`; allow existing `NamedArray` inputs to propagate their axis info.
+   - Add `haliax.freeze(ref)` returning a `NamedArray` while invalidating the underlying `jax.Ref`, preserving metadata.
+   - Provide `haliax.ref.get(ref, index)` / `haliax.ref.swap(ref, index, value)` that accept named indexers and return/require `NamedArray` values.
+
+3. **Indexing semantics**
+   - Reuse the existing NamedArray indexing utilities (`axis_spec_to_tuple`, `_normalize_indexers`, `_NamedIndexUpdateRef`) so `NamedRef.__getitem__` returns `NamedArray` views and `__setitem__` / `.set` accept either `NamedArray` or raw arrays (with shape checks).
+   - Enforce error messaging consistent with NamedArray when axes mismatch or out-of-bounds indices occur.
+
+4. **Slice refs**
+   - Add a `NamedSliceRef` (or similar) that wraps a base `NamedRef` plus a frozen partial index mapping.
+   - Support construction via `NamedRef.slice(selector)` where `selector` is any named index expression (dict, `Axis`, tuple, etc.).
+   - On read/write, merge the stored selector with the new selector: the stored selector is applied first, and new selectors can further index the result. Implement logic to resolve nested slices (e.g. range then integer index) by normalizing everything to positional indices via helper utilities.
+   - Ensure composition is associative: calling `.slice()` on a `NamedSliceRef` should produce another `NamedSliceRef` that collapses selectors appropriately.
+   - Preserve dynamic slice selectors (Haliax `dslice` or `pallas.dslice`) so their static size information is retained during trace-time shape inference.
+   - Provide an `unsliced()` helper to recover the original reference when staged selectors are no longer needed.
+   - Handle ellipsis and implicit axes; ensure merging respects axis ordering and catches conflicting assignments.
+
+5. **Integration & ergonomics**
+   - Audit parts of the codebase that manipulate mutable state (e.g. caching layers, training loops) to expose the new API where appropriate. Initially, surface `NamedRef` under the `haliax.ref` namespace without retrofitting all call sites.
+   - Provide a `NamedRef.unsafe_buffer_pointer()` passthrough for parity with `jax.Ref` diagnostics.
+
+6. **Testing strategy**
+   - Add `tests/test_named_ref.py` covering: creation from arrays, indexing/assignment (scalar, slices, advanced), `jit` usage, vmap plumbing, and ensuring restrictions (e.g. returning refs from jit) raise clear errors.
+   - Include dedicated tests for slice refs verifying merged indexing (simple axis, slices-of-slices, dict updates, ellipsis) and interactions under `jit`.
+   - Validate that gradients through pure functions using NamedRef internally behave like the array equivalents.
+
+7. **Documentation & guidelines**
+   - Document the new API in `docs/api.md` (table entry and short section) and, if needed, add a primer example in `docs/primer.md` showing plumbing state with NamedRef.
+   - Update contributor notes that future ref-based utilities must keep numerical tolerances intact (no relaxed tolerances), aligning with the existing testing guidelines.
+   - Call out any additional guidelines discovered during implementation (e.g., keeping dslice selectors intact for static shapes).
+   - Consider a `.playbooks/` entry if ref plumbing becomes a common flow after initial implementation.
+
+## Open Questions / Follow-ups
+- Confirm whether `NamedRef` needs to support asynchronous dispatch or pytree flattening beyond leaf behavior.
+- Decide if we want implicit conversion between `NamedRef` and `NamedArray` in certain helpers or require explicit `.value` reads for clarity.
+- Explore caching strategies or wrappers for `foreach`-style utilities once the base `NamedRef` and slice refs land.

{haliax-1.4.dev445 → haliax-1.4.dev447}/PKG-INFO

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

@@ -179,6 +179,42 @@ See also the section on [Indexing and Slicing](indexing.md).
 ::: haliax.take
 ::: haliax.updated_slice
 
+## Mutable References
+
+See also the section on [Mutable References](mutable-refs.md).
+
+JAX provides [`jax.Ref`][], a mutable array reference that can be read or written in place while remaining compatible
+with transformations such as `jax.jit` or `jax.grad`. Haliax mirrors that API with [`haliax.NamedRef`][haliax.NamedRef],
+which carries axis metadata so you can keep using named indexing when plumbing state through your programs.
+
+You introduce a new reference with [`haliax.new_ref`][haliax.new_ref]. The returned object behaves much like a
+[`NamedArray`][haliax.NamedArray] for indexing purposes: `ref[{"batch": 0}]` reads a slice, and assignments like
+`ref[{"token": slice(1, 3)}] = update` perform in-place updates on the underlying buffer. If you need to stage part of a
+reference for repeated use, call [`NamedRef.slice`][haliax.NamedRef.slice] to create a *slice ref*. Slice refs remember a
+partial indexing expression so you only supply the remaining axes during reads or writes:
+
+```python
+Cache = hax.Axis("layers", 24)
+Head = hax.Axis("head", 8)
+cache = hax.zeros((Cache, Head))
+cache_ref = hax.new_ref(cache)
+
+# Pin the layer axis once so subsequent lookups only specify the head coordinate.
+layer_ref = cache_ref.slice({"layers": slice(4, 8)})
+layer_ref[{"layers": 0, "head": 3}] = 1.0  # updates layer 4, head 3 in the original cache
+```
+
+When you are done mutating a reference, call [`haliax.freeze`][haliax.freeze] to invalidate it and recover a final
+`NamedArray` snapshot. You can also perform atomic-style updates with [`haliax.swap`][haliax.swap] (or the functional
+helpers under `haliax.ref`).
+
+::: haliax.NamedRef
+::: haliax.new_ref
+::: haliax.freeze
+::: haliax.swap
+
+[`jax.Ref`]: https://docs.jax.dev/en/latest/array_refs.html
+
 #### Dynamic Slicing
 
 ::: haliax.dslice

haliax-1.4.dev447/docs/mutable-refs.md

@@ -0,0 +1,133 @@
+# Mutable References
+
+JAX now exposes [`jax.Ref`](https://jax.readthedocs.io/en/latest/notebooks/refs.html), a mutable buffer that
+remains compatible with transformations such as `jax.jit`, `jax.grad`, and friends. Haliax mirrors that idea with
+[`haliax.NamedRef`][haliax.NamedRef], preserving axis metadata so you can keep using named indexing while wiring state
+through pure computations.
+
+This page walks through the common idioms you will encounter when working with named references, including ways to
+slice references, compose partial selectors, and freeze or swap out the underlying data when you are done.
+
+## Creating a `NamedRef`
+
+Use [`haliax.new_ref`][haliax.new_ref] to wrap an existing tensor. The helper accepts either a
+[`NamedArray`][haliax.NamedArray] or a raw JAX array plus an axis specification.
+
+```python
+import haliax as hax
+import jax
+
+Batch, Feature = hax.make_axes(Batch=4, Feature=8)
+weights = hax.random.normal(jax.random.PRNGKey(0), (Batch, Feature))
+weights_ref = hax.new_ref(weights)
+
+# You can still inspect the named shape.
+assert weights_ref.axes == weights.axes
+assert weights_ref.shape == {"Batch": 4, "Feature": 8}
+```
+
+Inside a pure function you can allocate fresh references to stage intermediate results:
+
+```python
+@jax.jit
+def normalize(x: hax.NamedArray) -> hax.NamedArray:
+    ref = hax.new_ref(x)
+    ref[{"Batch": slice(None)}] = ref[{"Batch": slice(None)}] - hax.mean(x, axis=Batch)
+    return ref[...]
+```
+
+References follow JAX's rules: you may not return the ref itself from inside a transformed function, pass the same ref
+argument multiple times, or close over a ref that is also passed as a parameter. Those restrictions keep aliasing under
+control so transformations can reason about side effects.
+
+## Reading and writing
+
+`NamedRef` behaves like a lightweight view of the underlying buffer. Reading always returns a `NamedArray`, while writing
+accepts either bare scalars/JAX arrays or another `NamedArray` with matching axes.
+
+```python
+logits_ref = hax.new_ref(hax.zeros((Batch, Feature)))
+
+# Read a slice.
+logits = logits_ref[{"Batch": 0}]
+
+# Write back in-place.
+logits_ref[{"Batch": 0}] = hax.random.uniform(jax.random.PRNGKey(1), logits.axes)
+
+# Mutations are visible through other aliasing refs.
+assert logits_ref[{"Batch": 0}].array is not None
+```
+
+If you prefer a functional style, use [`haliax.ref.get`][haliax.ref.get] or [`haliax.swap`][haliax.swap]. The latter
+swaps in a new value and returns the previous contents, similar to `dict.setdefault`.
+
+## Slice references
+
+Plumbing state often requires staging a subset of the axes. `NamedRef.slice` lets you pre-apply a named indexer to a
+reference and reuse it later without repeating the prefix. We call the result a *slice ref*.
+
+```python
+Cache = hax.Axis("layer", 24)
+Head = hax.Axis("head", 8)
+cache = hax.zeros((Cache, Head))
+cache_ref = hax.new_ref(cache)
+
+# Focus on a subset of layers.
+window = cache_ref.slice({"layer": slice(4, 8)})
+
+# Indexing the slice ref automatically splices the prefix into the base ref.
+window[{"layer": 0, "head": 3}] = 1.0  # updates layer 4, head 3 in the original buffer
+# `.value()` reads the staged slice.
+window_value = window.value()
+```
+
+Slice refs compose. Applying `.slice(...)` to an existing slice ref merges the new selection with the previous one, so
+`cache_ref.slice({"layer": slice(1, 4)}).slice({"layer": 0})` updates the second layer of the original cache. Integer and
+slice prefixes are supported; advanced indexing still needs to happen at call time (`slice` with `NamedArray` selectors or
+lists is not yet folded into the prefix).
+
+If you need to recover the original reference (without any staged prefixes), call [`NamedRef.unsliced`][haliax.NamedRef.unsliced].
+
+The merging logic mirrors `NamedArray.__getitem__`: ellipses expand to the remaining axes, implicit dimensions are
+filled in order, and conflicting assignments raise errors. See [Indexing and Slicing](indexing.md) for a refresher.
+
+## Working with transformations
+
+Because `NamedRef` is registered as a PyTree node, you can nest it inside larger pytree structures or pass it through
+`jax.tree_map` for bookkeeping. Only the underlying `jax.Ref` is considered a leaf, so axis metadata is treated as static
+structure and does not trigger recompilation.
+
+Keep in mind the JAX restrictions on mutable functions:
+
+- Do not return a ref from inside `jax.jit`, `jax.grad`, `jax.vmap`, `jax.lax.scan`, etc.
+- Avoid closing over a ref you also pass into the function (JAX disallows those aliasing patterns).
+- Use `jax.lax.stop_gradient` if you are plumbing values that should not affect differentiation.
+
+Pure functions that create refs internally continue to compose normally with autodiff and vectorization. If your ref is
+only used for bookkeeping (e.g., collecting auxiliary statistics), stop gradients before writing into it:
+
+```python
+def collect_metrics(x, stats_ref):
+    y = hax.sin(x)
+    stats_ref[{"Batch": slice(None)}] += jax.lax.stop_gradient(y)
+    return y
+```
+
+## Freezing and swapping
+
+When you are done mutating a reference, call [`haliax.freeze`][haliax.freeze] to convert it back into an immutable
+`NamedArray`.
+
+```python
+final = hax.freeze(cache_ref)
+assert isinstance(final, hax.NamedArray)
+```
+
+[`haliax.swap`][haliax.swap] provides an atomic-style update: it returns the previous value while storing the new one.
+The helper integrates with slice refs, so you can swap just a subset of the buffer.
+
+```python
+prev = hax.swap(cache_ref, {"layer": slice(0, 2)}, hax.ones((Cache.resize(2), Head)))
+```
+
+See [`tests/test_named_ref.py`](tests/test_named_ref.py) for runnable examples that exercise the API.

{haliax-1.4.dev445 → haliax-1.4.dev447}/mkdocs.yml

@@ -88,6 +88,7 @@ nav:
   - Named Arrays:
     - Broadcasting: 'broadcasting.md'
     - Indexing and Slicing: 'indexing.md'
+    - Mutable References: 'mutable-refs.md'
     - Rearrange: 'rearrange.md'
     - Matrix Multiplication: 'matmul.md'
     - Type Annotations: 'typing.md'

{haliax-1.4.dev445 → haliax-1.4.dev447}/src/haliax/__about__.py

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

{haliax-1.4.dev445 → haliax-1.4.dev447}/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.ref as ref  # noqa: F401
 import haliax.tree as tree  # noqa: F401
 import haliax.tree_util as tree_util
 import haliax.util as util
@@ -68,6 +69,7 @@ from .core import (
     unflatten_axis,
     updated_slice,
 )
+from .ref import NamedRef, freeze, get, new_ref, swap
 from .haxtyping import Named
 from .hof import fold, map, scan, vmap
 from .jax_utils import tree_checkpoint_name
@@ -1350,5 +1352,10 @@ __all__ = [
     "NamedArrayAxes",
     "NamedArrayAxesSpec",
     "Named",
+    "NamedRef",
+    "new_ref",
+    "freeze",
+    "swap",
+    "get",
     "set_mesh",
 ]

{haliax-1.4.dev445 → haliax-1.4.dev447}/src/haliax/jax_utils.py

@@ -17,6 +17,7 @@ from jax import numpy as jnp
 from jax import random as jrandom
 from jax.experimental.multihost_utils import host_local_array_to_global_array
 from jax.sharding import PartitionSpec
+from jax._src.state.indexing import Slice
 from jax.ad_checkpoint import checkpoint_name
 from jax.typing import DTypeLike
 from jaxtyping import PRNGKeyArray
@@ -218,13 +219,7 @@ def is_in_jit():
 
 
 def is_pallas_dslice(x: object) -> bool:
-    try:
-        from jax.experimental.pallas import dslice as pdslice
-    except ImportError:
-        return False
-
-    _PALLAS_DSLICE_TYPE = type(pdslice(0, 1))
-    return isinstance(x, _PALLAS_DSLICE_TYPE)
+    return isinstance(x, Slice)
 
 
 def is_scalarish(x):

{haliax-1.4.dev445 → haliax-1.4.dev447}/src/haliax/nn/scan.py

@@ -150,6 +150,11 @@ class BlockFoldable(Protocol[M]):
         """
         ...
 
+    def get_layer(self, index: int) -> M:
+        """Return the ``index``th layer of the folded module."""
+
+        ...
+
 
 class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
     """
@@ -377,6 +382,11 @@ class BlockSeq(ModuleWithStateDictSerialization, Generic[M]):
 
         return state_dict
 
+    def get_layer(self, index: int) -> M:
+        """Return the ``index``th block in this sequential container."""
+
+        return self.blocks[index]
+
     @property
     def _output_ckpt_name(self):
         return f"BlockSeq[{self.Block}, {self.blocks[0].__class__.__name__}].outputs"
@@ -442,6 +452,12 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
     Block: Axis = eqx.field(static=True)
     gradient_checkpointing: ScanCheckpointPolicy = eqx.field(static=True)
 
+    @property
+    def Layers(self) -> Axis:
+        """Alias for :attr:`Block` used by some downstream code."""
+
+        return self.Block
+
     @classmethod
     def init(
         cls,
@@ -700,6 +716,23 @@ class Stacked(ModuleWithStateDictSerialization, Generic[M]):
         unstacked_leaves = tuple(zip(*unstacked_leaves))
         return tuple(map(lambda x: jax.tree_util.tree_unflatten(structure, x), unstacked_leaves))
 
+    def get_layer(self, index: int) -> M:
+        """Return the ``index``th layer of this stacked module."""
+
+        def select_leaf(leaf):
+            if isinstance(leaf, haliax.NamedArray):
+                if haliax.selects_axis(leaf.axes, self.Block):
+                    return leaf[self.Block, index]
+                else:
+                    return leaf
+            if is_jax_or_hax_array_like(leaf):
+                if getattr(leaf, "shape", ()) and leaf.shape[0] == self.Block.size:
+                    return leaf[index]
+                return leaf
+            return leaf
+
+        return haliax.tree_util.tree_map(select_leaf, self.stacked)
+
     def to_state_dict(self, prefix: str | None = None) -> StateDict:
         # this method needs to "devectorize" the blocks, so that we have a list of blocks h.0.FOO, h.1.FOO, etc.
         # first just do the normal thing with our own dict, which we'll post-process