haliax 1.4.dev412__tar.gz → 1.4.dev413__tar.gz

{haliax-1.4.dev412 → haliax-1.4.dev413}/PKG-INFO

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

@@ -6,19 +6,89 @@ arrays (see [haliax.zeros][] and [haliax.ones][]) as well as things like reducti
 
 ## Axis Types
 
-There are four types related to [haliax.Axis][] you will see in the API reference:
-
-* [haliax.Axis][]: This is the main type for representing axes. It is a dataclass with a name and a size.
-* [haliax.AxisSelector][]: This is a type alias for either [haliax.Axis][] or a `str`. This type is used when we want
-  one axis and the size can be inferred from the inputs.
-* [haliax.AxisSpec][]: This is a type alias for either [haliax.Axis][] or a tuple/list of [haliax.Axis][]. This type is
-  used when we want one or more axes and the sizes cannot be inferred from the inputs, for instance when creating arrays.
-* [haliax.AxisSelection][]: This is a type alias for either [haliax.AxisSelector][] or a tuple/list of [haliax.AxisSelector][].
-    This type is used when we want one or more axes and the sizes can be inferred from the inputs, for instance when
-    reducing an array.
-
-Occasionally, an axis size can be inferred in some circumstances but not others. When this happens, we still use
-`AxisSelector` but document the behavior in the docstring. A RuntimeError will be raised if the size cannot be inferred.
+If you already speak NumPy or `jax.numpy`, think of Haliax as swapping positional axes (`axis=0`) for named axes
+(`axis="batch"`). The type hints in this section describe the different ways those named axes can be provided to the API.
+They appear throughout the documentation and in signatures so that you can quickly tell which forms are accepted.
+
+| Name | Accepts | When to use it | Example |
+| --- | --- | --- | --- |
+| [`Axis`][haliax.Axis] | `Axis(name: str, size: int)` | Define a named dimension with an explicit size | `Batch = Axis("batch", 32)` |
+| [`AxisSelector`][haliax.AxisSelector] | `Axis` or `str` | Refer to an existing axis whose size can be inferred from the arrays you pass in | `x.sum(axis="batch")` |
+| [`AxisSpec`][haliax.AxisSpec] | `dict[str, int]`, `Axis`, or a sequence of `Axis` objects | Create or reshape arrays when the axis sizes must be provided | `hax.zeros((Batch, Feature))` |
+| [`AxisSelection`][haliax.AxisSelection] | `dict[str, int | None]`, `AxisSpec`, or a sequence of `AxisSelector` values | Work with one or more existing axes (reductions, indexing helpers, flattening, …) | `x.sum(axis=("batch", Feature))` |
+
+### Axis
+
+An [`Axis`][haliax.Axis] is the fundamental building block: it is a tiny dataclass that stores a name and a size.  You can
+construct one directly or use [`haliax.make_axes`][] to generate several at a time.
+
+```python
+import haliax as hax
+from haliax import Axis
+
+Batch = Axis("batch", 32)
+Feature = Axis("feature", 128)
+x = hax.ones((Batch, Feature))
+print(Batch.name, Batch.size)
+```
+
+Using `Axis` objects keeps array creation explicit and gives reusable handles you can share between different tensors.
+Equality compares both the name and size so you get guardrails when wiring pieces together.
+
+### AxisSelector
+
+An [`AxisSelector`][haliax.AxisSelector] accepts either an `Axis` object or just the axis name as a string.  It is used
+whenever a function can read the axis size from one of its arguments.  This mirrors how NumPy lets you pass `axis=0` when
+reducing an array:
+
+```python
+total = x.sum(axis=Batch)       # using the Axis handle
+same_total = x.sum(axis="batch")  # using only the name
+```
+
+Strings are convenient when you only care about the name, but `Axis` objects still work so you can keep using the handles
+you created earlier.  If an axis with that name is missing, Haliax raises a `ValueError`.
+
+### AxisSpec
+
+An [`AxisSpec`][haliax.AxisSpec] is used when Haliax needs full size information to create or reshape an array.  You can
+provide a shape dictionary (sometimes called a "shape dict") that maps names to sizes, or a sequence of `Axis` objects:
+
+```python
+shape = {"batch": 32, "feature": 128}
+y = hax.zeros(shape)                 # using a shape dict
+z = hax.zeros((Batch, Feature))      # using the Axis objects directly
+```
+
+Both forms describe the same layout.  Python dictionaries preserve insertion order, so the ordering in a shape dict matches
+the order that axes appear in the array.  Sequences must contain `Axis` objects (not plain strings) because Haliax cannot
+otherwise know the axis sizes.
+
+### AxisSelection
+
+[`AxisSelection`][haliax.AxisSelection] generalizes the previous aliases so you can talk about several axes at once.  It
+shows up in reductions, indexing helpers, axis-mapping utilities, and anywhere you might have written `axis=(0, 1)` in
+NumPy.  You may supply:
+
+* a sequence mixing `Axis` objects and strings, e.g. `("batch", Feature)` when reducing two axes,
+* an `AxisSpec`, which is handy when you already have a tuple of `Axis` objects, or
+* a "partial shape dict" where the values are either sizes or `None` to indicate "any size".  Dictionaries are useful when
+  you only care about a subset of axes or want to assert a particular size.
+
+```python
+# Reduce over two axes using a tuple of selectors.
+scalar = x.sum(axis=("batch", Feature))
+
+# Ask for the axes by name and optionally pin sizes.
+x.resolve_axis({"batch": None, "feature": None})  # returns {"batch": 32, "feature": 128}
+
+from haliax.axis import selects_axis
+assert selects_axis((Batch, "feature"), {"batch": None, "feature": 128})
+```
+
+Occasionally, an axis size can be inferred in some circumstances but not others. When this happens we still use
+`AxisSelector` (or `AxisSelection` for multiple axes) but document the behavior in the docstring. A `RuntimeError` will be
+raised if the size cannot be inferred.
 
 ::: haliax.Axis
 ::: haliax.AxisSelector

haliax-1.4.dev413/docs/primer.md

@@ -0,0 +1,183 @@
+# 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.
+
+### Ways to Describe Axes
+
+If you're used to `axis=0` style code in NumPy or JAX, think of Haliax as swapping those positional indices for names
+like `axis="batch"`. The API hints refer to a few helper aliases; this table shows how they map back to familiar
+concepts:
+
+| Name | Accepts | Typical use | Example |
+| --- | --- | --- | --- |
+| [`Axis`][haliax.Axis] | `Axis(name: str, size: int)` | Define a named dimension with a fixed size | `Batch = Axis("batch", 32)` |
+| [`AxisSelector`][haliax.AxisSelector] | `Axis` or `str` | Refer to an existing axis when the size can be inferred from the argument | `x.sum(axis="batch")` |
+| [`AxisSpec`][haliax.AxisSpec] | `dict[str, int]`, `Axis`, or a sequence of `Axis` objects | Supply complete shape information (array creation, reshaping) | `hax.zeros((Batch, Feature))` |
+| [`AxisSelection`][haliax.AxisSelection] | `dict[str, int | None]`, `AxisSpec`, or a sequence of `AxisSelector` values | Work with one or more existing axes (reductions, indexing helpers, flattening, …) | `x.sum(axis=("batch", Feature))` |
+
+The following sections expand on each alias with quick references and NumPy-style parallels.
+
+#### `Axis`: reusable named dimensions
+
+An [`Axis`][haliax.Axis] stores a `name` and a `size`. Create them directly or let
+[`haliax.make_axes`][] build a handful at once. Because axes compare by both name and size, they act as reusable handles and
+catch many wiring mistakes early.
+
+```python
+Batch = Axis("batch", 32)
+Feature = Axis("feature", 128)
+x = hax.ones((Batch, Feature))
+print(Batch.name, Batch.size)
+```
+
+#### `AxisSelector`: when the size is already known
+
+Many functions already see the array whose axes you're referencing (e.g. reductions). In those cases you can pass either the
+`Axis` object or simply the axis name as a string. Haliax resolves the name against the array, similar to `axis=0` in NumPy.
+
+```python
+total = x.sum(axis=Batch)       # use the Axis handle
+same_total = x.sum(axis="batch")  # or just the name
+```
+
+If you reference an axis name that isn't present, Haliax raises a `ValueError`.
+
+#### `AxisSpec`: describing complete shapes
+
+When Haliax needs explicit sizes—creating arrays, reshaping, broadcasting to a new axis—you provide an [`AxisSpec`][haliax.AxisSpec].
+Shape dictionaries keep things close to standard Python, while sequences require actual `Axis` objects so the sizes stay explicit.
+
+```python
+shape = {"batch": 32, "feature": 128}
+y = hax.zeros(shape)                 # using a shape dict
+z = hax.zeros((Batch, Feature))      # or a sequence of Axis objects
+```
+
+Python dictionaries preserve insertion order, so the layout in a shape dict matches the order of axes in the resulting array.
+
+#### `AxisSelection`: several axes at once
+
+[`AxisSelection`][haliax.AxisSelection] is the plural form used by reductions, indexing helpers, and flattening utilities. Supply a
+tuple mixing `Axis` objects and strings, reuse an existing `AxisSpec`, or pass a partial shape dict where values are either
+sizes or `None` for "any size".
+
+```python
+scalar = x.sum(axis=("batch", Feature))
+
+# Ask for the axes by name and optionally pin sizes.
+x.resolve_axis({"batch": None, "feature": None})  # returns {"batch": 32, "feature": 128}
+```
+
+Partial shape dicts shine when you only care about a subset of axes or want assertions about their sizes. If the axis size
+cannot be inferred from the provided arguments, Haliax raises a `RuntimeError`.
+
+## 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.dev412 → haliax-1.4.dev413}/src/haliax/__about__.py

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

haliax-1.4.dev412/docs/primer.md

@@ -1,114 +0,0 @@
-# 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.