haliax 1.4.dev392__tar.gz → 1.4.dev394__tar.gz

{haliax-1.4.dev392 → haliax-1.4.dev394}/AGENTS.md

@@ -74,3 +74,4 @@ repository. Follow these notes when implementing new features or fixing bugs.
 
 ## Documentation
 - Public functions and modules require docstrings. If behavior is non‑obvious, add examples in `docs/`.
+- For a concise overview of Haliax aimed at LLM agents, see [docs/primer.md](docs/primer.md).

{haliax-1.4.dev392 → haliax-1.4.dev394}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev392
+Version: 1.4.dev394
 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.dev394/docs/primer.md

@@ -0,0 +1,114 @@
+# Haliax Primer
+
+Haliax provides named tensors built on top of JAX.  This primer is written for LLM agents and other downstream libraries and collects the core ideas for quick reference.
+
+## Axes and Named Arrays
+
+Arrays are indexed by `Axis` objects. You can define them explicitly or generate several with `make_axes`.
+You may also specify shapes with a **shape dict**, mapping axis names to sizes.
+
+```python
+import haliax as hax
+from haliax import Axis
+
+Batch = Axis("batch", 4)
+Feature = Axis("feature", 8)
+# or: Batch, Feature = hax.make_axes(batch=4, feature=8)
+# using Axis objects
+x = hax.zeros((Batch, Feature))
+# or using a shape dict
+shape = {"batch": 4, "feature": 8}
+x = hax.zeros(shape)
+```
+
+Most functions accept either axes or shape dicts interchangeably.
+
+A tensor with named axes is a [`NamedArray`][haliax.NamedArray]. Elementwise operations mirror `jax.numpy` but accept named axes.
+
+## Indexing and Broadcasting
+
+Use axis names when slicing. Dictionaries are convenient for several axes:
+
+```python
+first = x["batch", 0]
+sub = x["batch", 1:3]
+# or with a dict
+first = x[{"batch": 0}]
+sub = x[{"batch": slice(1, 3)}]
+```
+
+Axes broadcast by matching names. `broadcast_axis` adds a new axis to an array:
+
+```python
+row = hax.arange(Feature)
+outer = row.broadcast_axis(Batch) * hax.arange(Batch)
+```
+
+See [Indexing and Slicing](indexing.md) and [Broadcasting](broadcasting.md) for details.
+
+## Rearranging Axes
+
+`rearrange` changes axis order and can merge or split axes using einops‑style syntax.  It is useful when interfacing with positional APIs.
+
+```python
+# transpose features and batch
+x_t = hax.rearrange(x, "batch feature -> feature batch")
+```
+
+More examples appear in [Rearrange](rearrange.md).
+
+## Matrix Multiplication
+
+`dot` contracts over named axes while preserving order independence.
+
+```python
+Weight = Axis("weight", 8)
+w = hax.ones((Feature, Weight))
+prod = hax.dot(x, w, axis=Feature)
+```
+
+For more complex contractions use [`einsum`][haliax.einsum].  See [Matrix Multiplication](matmul.md).
+
+## Scans and Folds
+
+Use [`scan`][haliax.scan] or [`fold`][haliax.fold] to apply a function along an axis with optional gradient checkpointing.
+
+```python
+Time = Axis("time", 10)
+sequence = hax.ones((Time, Feature))
+
+def add(prev, cur):
+    return prev + cur
+
+result = hax.fold(add, Time)(hax.zeros((Feature,)), sequence)
+```
+
+See [Scan and Fold](scan.md) for checkpointing policies and stacked modules.
+
+## Partitioning
+
+Arrays and modules can be distributed across devices by mapping named axes to mesh axes:
+
+```python
+with hax.axis_mapping({"batch": "data"}):
+    sharded = hax.shard(x)
+```
+
+The [Partitioning](partitioning.md) guide explains how to set up device meshes and shard arrays.
+
+## Typing Support
+
+Type annotations use `haliax.haxtyping` which extends `jaxtyping`:
+
+```python
+import haliax.haxtyping as ht
+
+def f(t: ht.Float[hax.NamedArray, "batch feature"]):
+    ...
+```
+
+See [Typing](typing.md) for matching runtime checks and dtype-aware annotations.
+
+---
+
+This primer highlights common patterns.  The [cheatsheet](cheatsheet.md) lists many additional conversions from JAX to Haliax.

{haliax-1.4.dev392 → haliax-1.4.dev394}/mkdocs.yml

@@ -100,3 +100,5 @@ nav:
   - Serialization: 'state-dict.md'
   - API Reference: 'api.md'
   - FAQ: 'faq.md'
+  - LLMs:
+      - "LLM Primer": 'primer.md'

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

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

{haliax-1.4.dev392 → haliax-1.4.dev394}/src/haliax/core.py

@@ -32,6 +32,7 @@ from .axis import (
     dslice,
     eliminate_axes,
     selects_axis,
+    _check_size_consistency,
 )
 from .types import GatherScatterModeStr, IntScalar, PrecisionLike, Scalar
 
@@ -1649,39 +1650,62 @@ def _broadcast_axes(
 def broadcast_to(
     a: NamedOrNumeric, axes: AxisSpec, ensure_order: bool = True, enforce_no_extra_axes: bool = True
 ) -> NamedArray:
-    """
-    Broadcasts a so that it has the given axes.
-     If ensure_order is True (default), then the returned array will have the same axes in the same order as the given
-     axes. Otherwise, the axes may not be moved if they are already in the array. The axes may not be contiguous however
+    """Broadcast ``a`` so that it has the given axes.
+
+    If ``ensure_order`` is ``True`` (default) then the returned array's axes are
+    arranged in the same order as ``axes``. Otherwise existing axes may remain in
+    their current order, though they may still be moved to the front if new axes
+    are added.
 
-    If enforce_no_extra_axes is True and the array has axes that are not in axes, then a ValueError is raised.
+    If ``enforce_no_extra_axes`` is ``True`` and ``a`` has axes that are not in
+    ``axes`` then a ``ValueError`` is raised.
     """
-    axes = axis_spec_to_tuple(axes)
+
+    axes_dict = axis_spec_to_shape_dict(axes)
+    axes_tuple = axis_spec_to_tuple(axes)
 
     if not isinstance(a, NamedArray):
         a = named(jnp.asarray(a), ())
 
     assert isinstance(a, NamedArray)  # mypy gets confused
 
-    if a.axes == axes:
-        return a
+    a_axes_dict = axis_spec_to_shape_dict(a.axes)
+
+    # fill in missing sizes and check for mismatches
+    for name, sz in list(axes_dict.items()):
+        if sz is None:
+            if name not in a_axes_dict:
+                raise ValueError(
+                    f"Cannot broadcast: size for axis '{name}' is unspecified and it does not exist in array"
+                )
+            axes_dict[name] = a_axes_dict[name]
+        elif name in a_axes_dict:
+            _check_size_consistency(axes, a.axes, name, sz, a_axes_dict[name])
 
-    to_add = tuple(ax for ax in axes if ax not in a.axes)
+    extra_axis_names = [ax.name for ax in a.axes if ax.name not in axes_dict]
+    if enforce_no_extra_axes and extra_axis_names:
+        raise ValueError(
+            f"Cannot broadcast {a.shape} to {axes_dict}: extra axes present {extra_axis_names}"
+        )
 
-    all_axes = to_add + a.axes
+    axes_names_in_a = {ax.name for ax in a.axes}
+    to_add = tuple(
+        Axis(axis_name(ax), axes_dict[axis_name(ax)])
+        for ax in axes_tuple
+        if axis_name(ax) not in axes_names_in_a
+    )
 
-    if enforce_no_extra_axes and len(all_axes) != len(axes):
-        raise ValueError(f"Cannot broadcast {a.shape} to {axis_spec_to_shape_dict(axes)}: extra axes present")
+    all_axes = to_add + a.axes
 
-    extra_axes = tuple(ax for ax in a.axes if ax not in axes)
+    extra_axes = tuple(ax for ax in a.axes if ax.name not in axes_dict)
 
-    # broadcast whatever we need to the front and reorder
     a_array = jnp.broadcast_to(a.array, [ax.size for ax in all_axes])
     a = NamedArray(a_array, all_axes)
 
-    # if the new axes are already in the right order, then we're done
-    if ensure_order and not _is_subsequence(axes, all_axes):
-        a = a.rearrange(axes + extra_axes)
+    axes_tuple_complete = tuple(Axis(axis_name(ax), axes_dict[axis_name(ax)]) for ax in axes_tuple)
+
+    if ensure_order and not _is_subsequence(axes_tuple_complete, all_axes):
+        a = a.rearrange(axes_tuple_complete + extra_axes)
 
     return typing.cast(NamedArray, a)
 

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

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