haliax 1.4.dev411__tar.gz → 1.4.dev412__tar.gz

{haliax-1.4.dev411 → haliax-1.4.dev412}/.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.dev412}/.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.dev412}/PKG-INFO

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

@@ -289,6 +289,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.dev411 → haliax-1.4.dev412}/src/haliax/__about__.py

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

{haliax-1.4.dev411 → haliax-1.4.dev412}/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",

haliax-1.4.dev412/src/haliax/fft.py

@@ -0,0 +1,365 @@
+"""Named wrappers around :mod:`jax.numpy.fft`.
+
+These functions mirror the behaviour of their :mod:`jax.numpy.fft` counterparts
+while accepting named axes.  Instead of separate ``fftn``/``fft2`` variants we
+provide a single ``fft`` family of functions whose ``axis`` argument controls
+which axes are transformed.
+
+The ``axis`` parameter can be one of:
+
+* ``None`` – operate on the last axis.
+* ``str`` – name of an existing axis in the input.
+* :class:`~haliax.Axis` – specifies both the axis to transform (by name) and the
+  desired FFT length.  The output axis is replaced by the provided ``Axis``.
+* ``dict`` – mapping from axis selectors (names or ``Axis`` objects) to optional
+  sizes.  A value of ``None`` uses the existing axis length.  The mapping order
+  determines the order of transforms and dispatches to the ``n``‑dimensional
+  variants in :mod:`jax.numpy.fft`.
+
+Example
+-------
+
+```python
+X, Y = hax.make_axes(X=4, Y=6)
+arr = hax.arange((X, Y))
+
+# 1D transform along ``Y``
+hax.fft(arr, axis="Y")
+
+# 2D transform across both axes
+hax.fft(arr, axis={"X": None, "Y": None})
+
+# Resize the ``Y`` axis before transforming
+hax.fft(arr, axis={"Y": Axis("Y", 8)})
+```
+"""
+
+from __future__ import annotations
+
+from typing import Mapping, MutableSequence, Sequence
+
+import jax.numpy.fft as jfft
+
+from .axis import Axis, AxisSelector, AxisSelection
+from .core import NamedArray
+
+
+AxisSizeLike = int | Axis | None
+AxisMapping = Mapping[AxisSelector, AxisSizeLike]
+
+
+def _single_axis(a: NamedArray, axis: AxisSelector | None):
+    if axis is None:
+        idx = a.ndim - 1
+        ax = a.axes[idx]
+        n = None
+    elif isinstance(axis, Axis):
+        idx = a.axis_indices(axis.name)
+        if idx is None:
+            raise ValueError(f"Axis {axis} not found in {a.axes}")
+        ax = axis
+        n = axis.size
+    else:
+        idx = a.axis_indices(axis)
+        if idx is None:
+            raise ValueError(f"Axis {axis} not found in {a.axes}")
+        ax = a.axes[idx]
+        n = None
+    return idx, ax, n
+
+
+def _multi_axis(a: NamedArray, axis: AxisMapping):
+    axes_idx: MutableSequence[int] = []
+    sizes: MutableSequence[int] = []
+    new_axes = list(a.axes)
+    for key, val in axis.items():
+        idx = a.axis_indices(key)
+        if idx is None:
+            raise ValueError(f"Axis {key} not found in {a.axes}")
+        ax = a.axes[idx]
+        if isinstance(val, Axis):
+            size = val.size
+            new_axes[idx] = val
+        elif val is None:
+            size = ax.size
+        else:
+            size = int(val)
+            new_axes[idx] = ax.resize(size)
+        axes_idx.append(idx)
+        sizes.append(size)
+    return list(axes_idx), list(sizes), new_axes
+
+
+def fft(
+    a: NamedArray,
+    axis: AxisSelector | Sequence[AxisSelector] | AxisMapping | None = None,
+    norm: str | None = None,
+) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.fft`.
+
+    See module level documentation for the behaviour of the ``axis`` argument.
+    """
+
+    if isinstance(axis, Mapping):
+        axes_idx, sizes, new_axes = _multi_axis(a, axis)
+        out = jfft.fftn(a.array, s=tuple(sizes), axes=tuple(axes_idx), norm=norm)
+        return NamedArray(out, tuple(new_axes))
+    elif isinstance(axis, Sequence) and not isinstance(axis, (str, Axis)):
+        axes_idx, sizes, new_axes = _multi_axis(a, {ax: None for ax in axis})
+        out = jfft.fftn(a.array, s=tuple(sizes), axes=tuple(axes_idx), norm=norm)
+        return NamedArray(out, tuple(new_axes))
+    else:
+        idx, new_axis, n = _single_axis(a, axis)  # type: ignore[arg-type]
+        out = jfft.fft(a.array, n=n, axis=idx, norm=norm)
+        axes = list(a.axes)
+        axes[idx] = new_axis
+        return NamedArray(out, tuple(axes))
+
+
+def ifft(
+    a: NamedArray,
+    axis: AxisSelector | Sequence[AxisSelector] | AxisMapping | None = None,
+    norm: str | None = None,
+) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.ifft`."""
+
+    if isinstance(axis, Mapping):
+        axes_idx, sizes, new_axes = _multi_axis(a, axis)
+        out = jfft.ifftn(a.array, s=tuple(sizes), axes=tuple(axes_idx), norm=norm)
+        return NamedArray(out, tuple(new_axes))
+    elif isinstance(axis, Sequence) and not isinstance(axis, (str, Axis)):
+        axes_idx, sizes, new_axes = _multi_axis(a, {ax: None for ax in axis})
+        out = jfft.ifftn(a.array, s=tuple(sizes), axes=tuple(axes_idx), norm=norm)
+        return NamedArray(out, tuple(new_axes))
+    else:
+        idx, new_axis, n = _single_axis(a, axis)  # type: ignore[arg-type]
+        out = jfft.ifft(a.array, n=n, axis=idx, norm=norm)
+        axes = list(a.axes)
+        axes[idx] = new_axis
+        return NamedArray(out, tuple(axes))
+
+
+def rfft(
+    a: NamedArray,
+    axis: AxisSelector | Sequence[AxisSelector] | AxisMapping | None = None,
+    norm: str | None = None,
+) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.rfft`."""
+
+    if isinstance(axis, Mapping):
+        axes_idx, sizes, new_axes = _multi_axis(a, axis)
+        out = jfft.rfftn(a.array, s=tuple(sizes), axes=tuple(axes_idx), norm=norm)
+        last_idx = axes_idx[-1]
+        last_in = sizes[-1]
+        new_axes[last_idx] = new_axes[last_idx].resize(last_in // 2 + 1)
+        return NamedArray(out, tuple(new_axes))
+    elif isinstance(axis, Sequence) and not isinstance(axis, (str, Axis)):
+        axes_idx, sizes, new_axes = _multi_axis(a, {ax: None for ax in axis})
+        out = jfft.rfftn(a.array, s=tuple(sizes), axes=tuple(axes_idx), norm=norm)
+        last_idx = axes_idx[-1]
+        last_in = sizes[-1]
+        new_axes[last_idx] = new_axes[last_idx].resize(last_in // 2 + 1)
+        return NamedArray(out, tuple(new_axes))
+    else:
+        idx, ax, n = _single_axis(a, axis)  # type: ignore[arg-type]
+        out = jfft.rfft(a.array, n=n, axis=idx, norm=norm)
+        length = n if n is not None else ax.size
+        new_axis = ax.resize(length // 2 + 1)
+        axes = list(a.axes)
+        axes[idx] = new_axis
+        return NamedArray(out, tuple(axes))
+
+
+def _single_axis_irfft(a: NamedArray, axis: AxisSelector | None):
+    if axis is None:
+        idx = a.ndim - 1
+        in_ax = a.axes[idx]
+        length = (in_ax.size - 1) * 2
+        out_ax = in_ax.resize(length)
+        n = None
+    elif isinstance(axis, Axis):
+        idx = a.axis_indices(axis.name)
+        if idx is None:
+            raise ValueError(f"Axis {axis} not found in {a.axes}")
+        out_ax = axis
+        n = axis.size
+    else:
+        idx = a.axis_indices(axis)
+        if idx is None:
+            raise ValueError(f"Axis {axis} not found in {a.axes}")
+        in_ax = a.axes[idx]
+        length = (in_ax.size - 1) * 2
+        out_ax = in_ax.resize(length)
+        n = None
+    return idx, out_ax, n
+
+
+def _multi_axis_irfft(a: NamedArray, axis: AxisMapping):
+    axes_idx: MutableSequence[int] = []
+    sizes: MutableSequence[int] = []
+    new_axes = list(a.axes)
+    items = list(axis.items())
+    for i, (key, val) in enumerate(items):
+        idx = a.axis_indices(key)
+        if idx is None:
+            raise ValueError(f"Axis {key} not found in {a.axes}")
+        ax = a.axes[idx]
+        if isinstance(val, Axis):
+            size = val.size
+            new_axes[idx] = val
+        elif val is None:
+            if i == len(items) - 1:
+                size = (ax.size - 1) * 2
+            else:
+                size = ax.size
+            new_axes[idx] = ax.resize(size)
+        else:
+            size = int(val)
+            new_axes[idx] = ax.resize(size)
+        axes_idx.append(idx)
+        sizes.append(size)
+    return list(axes_idx), list(sizes), new_axes
+
+
+def irfft(
+    a: NamedArray,
+    axis: AxisSelector | Sequence[AxisSelector] | AxisMapping | None = None,
+    norm: str | None = None,
+) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.irfft`."""
+
+    if isinstance(axis, Mapping):
+        axes_idx, sizes, new_axes = _multi_axis_irfft(a, axis)
+        out = jfft.irfftn(a.array, s=tuple(sizes), axes=tuple(axes_idx), norm=norm)
+        return NamedArray(out, tuple(new_axes))
+    elif isinstance(axis, Sequence) and not isinstance(axis, (str, Axis)):
+        axes_idx, sizes, new_axes = _multi_axis_irfft(a, {ax: None for ax in axis})
+        out = jfft.irfftn(a.array, s=tuple(sizes), axes=tuple(axes_idx), norm=norm)
+        return NamedArray(out, tuple(new_axes))
+    else:
+        idx, out_ax, n = _single_axis_irfft(a, axis)  # type: ignore[arg-type]
+        out = jfft.irfft(a.array, n=n, axis=idx, norm=norm)
+        axes = list(a.axes)
+        axes[idx] = out_ax
+        return NamedArray(out, tuple(axes))
+
+
+def hfft(
+    a: NamedArray,
+    axis: AxisSelector | Sequence[AxisSelector] | AxisMapping | None = None,
+    norm: str | None = None,
+) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.hfft`.
+
+    Only a single axis is supported; passing a dictionary with more than one
+    entry will raise an error.
+    """
+
+    if isinstance(axis, Mapping):
+        if len(axis) != 1:
+            raise ValueError("hfft only supports a single axis")
+        key, val = next(iter(axis.items()))
+        if isinstance(val, Axis):
+            axis = val
+        elif val is None:
+            axis = key
+        else:
+            name = key.name if isinstance(key, Axis) else key
+            axis = Axis(name, int(val))
+
+    idx, out_ax, n = _single_axis_irfft(a, axis)  # type: ignore[arg-type]
+    out = jfft.hfft(a.array, n=n, axis=idx, norm=norm)
+    axes = list(a.axes)
+    axes[idx] = out_ax
+    return NamedArray(out, tuple(axes))
+
+
+def ihfft(
+    a: NamedArray,
+    axis: AxisSelector | Sequence[AxisSelector] | AxisMapping | None = None,
+    norm: str | None = None,
+) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.ihfft`.
+
+    Only a single axis is supported; passing a dictionary with more than one
+    entry will raise an error.
+    """
+
+    if isinstance(axis, Mapping):
+        if len(axis) != 1:
+            raise ValueError("ihfft only supports a single axis")
+        key, val = next(iter(axis.items()))
+        if isinstance(val, Axis):
+            axis = val
+        elif val is None:
+            axis = key
+        else:
+            name = key.name if isinstance(key, Axis) else key
+            axis = Axis(name, int(val))
+
+    idx, ax, n = _single_axis(a, axis)  # type: ignore[arg-type]
+    out = jfft.ihfft(a.array, n=n, axis=idx, norm=norm)
+    length = n if n is not None else ax.size // 2 + 1
+    new_axis = ax.resize(length)
+    axes = list(a.axes)
+    axes[idx] = new_axis
+    return NamedArray(out, tuple(axes))
+
+
+def fftshift(x: NamedArray, axes: AxisSelection | None = None) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.fftshift`."""
+
+    if axes is None:
+        out = jfft.fftshift(x.array)
+    else:
+        idxs = x.axis_indices(axes)
+        if isinstance(idxs, tuple):
+            if any(i is None for i in idxs):
+                raise ValueError(f"Axis {axes} not found in {x.axes}")
+        elif idxs is None:
+            raise ValueError(f"Axis {axes} not found in {x.axes}")
+        out = jfft.fftshift(x.array, axes=idxs)
+    return NamedArray(out, x.axes)
+
+
+def ifftshift(x: NamedArray, axes: AxisSelection | None = None) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.ifftshift`."""
+
+    if axes is None:
+        out = jfft.ifftshift(x.array)
+    else:
+        idxs = x.axis_indices(axes)
+        if isinstance(idxs, tuple):
+            if any(i is None for i in idxs):
+                raise ValueError(f"Axis {axes} not found in {x.axes}")
+        elif idxs is None:
+            raise ValueError(f"Axis {axes} not found in {x.axes}")
+        out = jfft.ifftshift(x.array, axes=idxs)
+    return NamedArray(out, x.axes)
+
+
+def fftfreq(axis: Axis, d: float = 1.0) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.fftfreq`."""
+
+    return NamedArray(jfft.fftfreq(axis.size, d), (axis,))
+
+
+def rfftfreq(axis: Axis, d: float = 1.0) -> NamedArray:
+    """Named version of :func:`jax.numpy.fft.rfftfreq`."""
+
+    new_axis = axis.resize(axis.size // 2 + 1)
+    return NamedArray(jfft.rfftfreq(axis.size, d), (new_axis,))
+
+
+__all__ = [
+    "fft",
+    "ifft",
+    "rfft",
+    "irfft",
+    "hfft",
+    "ihfft",
+    "fftfreq",
+    "rfftfreq",
+    "fftshift",
+    "ifftshift",
+]

haliax-1.4.dev412/tests/test_fft.py

@@ -0,0 +1,81 @@
+import jax.numpy as jnp
+import jax.numpy.fft as jfft
+
+import haliax as hax
+from haliax import Axis
+
+
+def test_fft_axis_options():
+    N = Axis("n", 8)
+    x = hax.arange(N, dtype=jnp.float32)
+
+    # string axis
+    assert jnp.allclose(hax.fft(x, axis="n").array, jfft.fft(x.array))
+
+    # Axis object with resize
+    N2 = Axis("n", 16)
+    f = hax.fft(x, axis=N2)
+    assert f.axes[0] == N2
+    assert jnp.allclose(f.array, jfft.fft(x.array, n=16))
+
+    r = hax.rfft(x)
+    assert r.axes[0].size == 5
+    assert jnp.allclose(r.array, jfft.rfft(x.array))
+
+    ir = hax.irfft(r)
+    assert ir.axes[0].size == 8
+    assert jnp.allclose(ir.array, jfft.irfft(jfft.rfft(x.array)))
+
+    h = hax.hfft(r)
+    assert h.axes[0].size == 8
+    assert jnp.allclose(h.array, jfft.hfft(jfft.rfft(x.array)))
+
+    ih = hax.ihfft(x)
+    assert ih.axes[0].size == 5
+    assert jnp.allclose(ih.array, jfft.ihfft(x.array))
+
+
+def test_fft_freq_and_shift():
+    N = Axis("n", 8)
+    x = hax.arange(N)
+
+    f = hax.fftfreq(N)
+    assert f.axes == (N,)
+    assert jnp.allclose(f.array, jfft.fftfreq(8))
+
+    rf = hax.rfftfreq(N)
+    assert rf.axes[0].size == 5
+    assert jnp.allclose(rf.array, jfft.rfftfreq(8))
+
+    shifted = hax.fftshift(x)
+    assert jnp.allclose(shifted.array, jfft.fftshift(x.array))
+    unshifted = hax.ifftshift(shifted)
+    assert jnp.allclose(unshifted.array, x.array)
+
+
+def test_fft_multi_axis():
+    X = Axis("x", 4)
+    Y = Axis("y", 6)
+    Z = Axis("z", 8)
+    arr = hax.arange((X, Y, Z), dtype=jnp.float32)
+
+    f = hax.fft(arr, axis={"y": None, "z": None})
+    assert jnp.allclose(f.array, jfft.fftn(arr.array, axes=(1, 2)))
+
+    f_seq = hax.fft(arr, axis=("y", "z"))
+    assert jnp.allclose(f_seq.array, jfft.fftn(arr.array, axes=(1, 2)))
+    assert f_seq.axes == f.axes
+
+    rf = hax.rfft(arr, axis={"y": None, "z": None})
+    assert rf.axes[2].size == 5
+    assert jnp.allclose(rf.array, jfft.rfftn(arr.array, axes=(1, 2)))
+
+    irf = hax.irfft(rf, axis={"y": None, "z": Z})
+    assert jnp.allclose(
+        irf.array, jfft.irfftn(jfft.rfftn(arr.array, axes=(1, 2)), s=(Y.size, Z.size), axes=(1, 2))
+    )
+
+    # resizing via dict values
+    f2 = hax.fft(arr, axis={"y": 4, "z": None})
+    assert f2.axes[1].size == 4
+    assert jnp.allclose(f2.array, jfft.fftn(arr.array, s=(4, 8), axes=(1, 2)))