haliax 1.4.dev411__tar.gz → 1.4.dev413__tar.gz

{haliax-1.4.dev411 → haliax-1.4.dev413}/.agents/projects/api_parity.md

@@ -161,24 +161,19 @@ APIs that don't translate well to named tensors are intentionally omitted here.
 - [ ] `vstack`
 
 ## fft
-- [ ] `fft`
-- [ ] `fft2`
-- [ ] `fftfreq`
-- [ ] `fftn`
-- [ ] `fftshift`
-- [ ] `hfft`
-- [ ] `ifft`
-- [ ] `ifft2`
-- [ ] `ifftn`
-- [ ] `ifftshift`
-- [ ] `ihfft`
-- [ ] `irfft`
-- [ ] `irfft2`
-- [ ] `irfftn`
-- [ ] `rfft`
-- [ ] `rfft2`
-- [ ] `rfftfreq`
-- [ ] `rfftn`
+Multi-dimensional transforms like `fftn` or `rfft2` are handled by passing a
+mapping of axes to :func:`haliax.fft` and friends.
+
+- [x] `fft`
+- [x] `fftfreq`
+- [x] `fftshift`
+- [x] `hfft`
+- [x] `ifft`
+- [x] `ifftshift`
+- [x] `ihfft`
+- [x] `irfft`
+- [x] `rfft`
+- [x] `rfftfreq`
 
 ## linalg
 - [ ] `cholesky`

{haliax-1.4.dev411 → haliax-1.4.dev413}/.playbooks/wrap-non-named.md

@@ -45,6 +45,14 @@ def sum(a, axis=None):
 ## Harder Cases
 Some functions need bespoke handling. For example `jnp.unique` returns several arrays and may change shape unpredictably. There is no generic helper, so you will need to manually map between `NamedArray` axes and the outputs. Use the lower level utilities in `haliax.wrap` for broadcasting and axis lookup.
 
+## Axis-aware functions
+For JAX functions that accept an ``axis`` or ``axes`` argument, prefer a single
+``axis`` parameter of type :class:`haliax.AxisSelection` that also accepts an
+ordered mapping of axes to sizes. A mapping dispatches to the corresponding
+``n``‑dimensional JAX primitive (e.g. :func:`jax.numpy.fft.fftn`) and allows
+resizing the transformed axes by specifying integer lengths or ``Axis``
+instances.
+
 ## Testing
 Add tests to ensure that named and unnamed calls produce the same results and that axis names are preserved or removed correctly.
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: haliax
-Version: 1.4.dev411
+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.dev411 → 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:
+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.
 
-* [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.
+| 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))` |
 
-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.
+### 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
@@ -289,6 +359,49 @@ These are all more or less directly from JAX's NumPy API.
 ::: haliax.triu
 ::: haliax.where
 
+### FFT
+
+All FFT helpers accept an ``axis`` argument which may be a single axis, its
+name, or an ordered mapping from axes to output sizes.  Passing a mapping
+dispatches to the ``n``‑dimensional variants in :mod:`jax.numpy.fft`.
+
+For example::
+
+    import jax.numpy as jnp
+    import haliax as hax
+
+    T = hax.Axis("time", 8)
+    signal = hax.arange(T, dtype=jnp.float32)
+
+    # operate along a single axis specified by name
+    hax.fft(signal, axis="time")
+
+    # resize by passing an Axis object
+    hax.fft(signal, axis=hax.Axis("time", 16))
+
+    X, Y = hax.make_axes(X=4, Y=6)
+    image = hax.arange((X, Y), dtype=jnp.float32)
+
+    # transform across several axes in order by passing a sequence
+    hax.fft(image, axis=("X", "Y"))
+
+    # selectively resize axes by providing a mapping
+    hax.fft(image, axis={"X": None, "Y": hax.Axis("Y", 10)})
+
+    # mappings can cover just a subset of axes when only partial resizing is needed
+    hax.fft(image, axis={"Y": 10})
+
+::: haliax.fft
+::: haliax.ifft
+::: haliax.hfft
+::: haliax.ihfft
+::: haliax.rfft
+::: haliax.irfft
+::: haliax.fftfreq
+::: haliax.rfftfreq
+::: haliax.fftshift
+::: haliax.ifftshift
+
 
 
 ## Named Array Reference

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.dev411 → haliax-1.4.dev413}/src/haliax/__about__.py

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

{haliax-1.4.dev411 → haliax-1.4.dev413}/src/haliax/__init__.py

@@ -96,6 +96,18 @@ from .ops import (
     bincount,
     where,
 )
+from .fft import (
+    fft,
+    fftfreq,
+    fftshift,
+    hfft,
+    ifft,
+    ifftshift,
+    ihfft,
+    irfft,
+    rfft,
+    rfftfreq,
+)
 from .partitioning import auto_sharded, axis_mapping, fsdp, named_jit, shard, shard_with_axis_mapping
 from .specialized_fns import top_k
 from .types import Scalar
@@ -1214,6 +1226,16 @@ __all__ = [
     "clip",
     "tril",
     "triu",
+    "fft",
+    "ifft",
+    "hfft",
+    "ihfft",
+    "rfft",
+    "irfft",
+    "fftfreq",
+    "rfftfreq",
+    "fftshift",
+    "ifftshift",
     "add",
     "arctan2",
     "bitwise_and",