Array API Support#

ezmsg-sigproc provides support for the Python Array API standard, enabling many transformers to work with arrays from different backends such as NumPy, CuPy, PyTorch, JAX, and MLX.

What is the Array API?#

The Array API is a standardized interface for array operations across different Python array libraries. By coding to this standard, ezmsg-sigproc transformers can process data regardless of which array library created it, enabling:

GPU acceleration via CuPy, PyTorch, or JAX tensors
Apple Silicon acceleration via MLX
Framework interoperability for integration with ML pipelines
Hardware flexibility without code changes

How It Works#

Compatible transformers use array-api-compat to detect the input array’s namespace and use the appropriate operations:

from array_api_compat import get_namespace

def _process(self, message: AxisArray) -> AxisArray:
    xp = get_namespace(message.data)  # numpy, cupy, torch, mlx.core, etc.
    result = xp.abs(message.data)     # Uses the correct backend
    return replace(message, data=result)

Usage Example#

Using Array API compatible transformers with CuPy for GPU acceleration:

import cupy as cp
from ezmsg.util.messages.axisarray import AxisArray
from ezmsg.sigproc.math.abs import AbsTransformer
from ezmsg.sigproc.math.clip import ClipTransformer, ClipSettings

# Create data on GPU
gpu_data = cp.random.randn(1000, 64).astype(cp.float32)
message = AxisArray(gpu_data, dims=["time", "ch"])

# Process entirely on GPU - no data transfer!
abs_transformer = AbsTransformer()
clip_transformer = ClipTransformer(ClipSettings(min=0.0, max=1.0))

result = clip_transformer(abs_transformer(message))
# result.data is still a CuPy array on GPU

Compatible Modules#

The following transformers fully support the Array API standard:

Math Operations#

Module	Description
`ezmsg.sigproc.math.abs`	Absolute value
`ezmsg.sigproc.math.clip`	Clip values to a range
`ezmsg.sigproc.math.log`	Logarithm with configurable base
`ezmsg.sigproc.math.scale`	Multiply by a constant
`ezmsg.sigproc.math.invert`	Compute 1/x
`ezmsg.sigproc.math.difference`	Subtract a constant (ConstDifferenceTransformer)

Signal Processing#

Module	Description
`ezmsg.sigproc.spectrum`	FFT-based spectrum (SpectrumTransformer)
`ezmsg.sigproc.aggregate`	Aggregate operations (AggregateTransformer, RangedAggregateTransformer)
`ezmsg.sigproc.diff`	Compute differences along an axis
`ezmsg.sigproc.transpose`	Transpose/permute array dimensions
`ezmsg.sigproc.linear`	Per-channel linear transform (scale + offset)

Coordinate Transforms#

Module	Description
`ezmsg.sigproc.coordinatespaces`	Cartesian/polar coordinate conversions

Composite Pipelines#

These CompositeProcessor pipelines chain Array API-aware steps together. When fed non-NumPy arrays, each step in the pipeline preserves the backend:

Module	Description
`ezmsg.sigproc.bandpower`	BandPowerTransformer (spectrogram + ranged aggregate)
`ezmsg.sigproc.singlebandpow`	RMSBandPowerTransformer (with explicit `backend` setting; only after initial IIR filter)

MLX on Apple Silicon#

MLX is an array library for Apple Silicon that provides GPU-accelerated operations with a NumPy-like API. ezmsg-sigproc’s Array API support enables MLX acceleration for spectral analysis and other pipelines without code changes to the transformers themselves.

Basic usage#

Pass MLX arrays in your AxisArray messages:

import mlx.core as mx
import numpy as np
from ezmsg.util.messages.axisarray import AxisArray
from ezmsg.sigproc.spectrum import SpectrumTransformer, SpectrumSettings

# Create data as MLX array
np_data = np.random.randn(1000, 64).astype(np.float32)
message = AxisArray(
    data=mx.array(np_data),
    dims=["time", "ch"],
    axes={"time": AxisArray.TimeAxis(fs=1000.0)},
)

proc = SpectrumTransformer(SpectrumSettings(axis="time"))
result = proc(message)
# result.data is an mlx.core.array

Lazy evaluation and `mx.eval`#

MLX uses lazy evaluation — computations are not executed until their results are needed. This allows MLX to fuse operations and optimize the computation graph. However, it means that timing code or downstream consumers may see artificially fast “processing” that is actually deferred.

To force evaluation, call mx.eval():

result = proc(message)
mx.eval(result.data)  # Forces computation to complete

Choosing where to evaluate#

A lazy backend only pays off while the graph stays lazy, so where you force evaluation is a property of the graph rather than of any one node. Wire an Materialize node at the point you want the barrier and pick a MaterializeMode:

from ezmsg.sigproc.materialize import Materialize, MaterializeSettings

# Bound the graph without stalling the caller
MAT = Materialize(MaterializeSettings(mode="async"))

sync (the default) blocks until the device finishes; async schedules the work and returns immediately; off leaves the data lazy. All three detach the pending graph except off, so async is usually the right choice when the guarantee you want is “the graph does not accumulate” rather than “the values are on the host now” — it costs no round-trip and leaves the CPU free to build the next message.

Two things are worth knowing before placing one:

Evaluating an output also evaluates the state computed alongside it. MLX evaluates a multi-output primitive once and materializes all of its outputs. The Metal kernels in this package emit filter state as a second output of the same kernel as the filtered signal (ewma_mlx_metal, sosfilt_mlx_metal), so a single barrier on the output is enough — the stateful transformers do not need, and deliberately do not have, their own internal mx.eval calls.

Evaluating an empty message forces nothing. A branch that withholds output on some cycles — a binning stage fed sub-bin chunks, for instance — leaves its upstream graph un-evaluated on exactly those cycles. Place the barrier upstream of any stage that can emit a zero-length message, not downstream of it.

A CompositeProcessor can carry its own barrier rather than relying on a separate node, by overriding _post_process and delegating to materialize_array:

class BandPowerTransformer(CompositeProcessor[BandPowerSettings, AxisArray, AxisArray]):
    @staticmethod
    def _initialize_processors(settings):
        return {
            "spectrogram": SpectrogramTransformer(settings=settings.spectrogram_settings),
            "aggregate": RangedAggregateTransformer(...),
        }

    def _post_process(self, result: AxisArray | None) -> AxisArray | None:
        if result is not None:
            materialize_array(result.data, self.settings.materialize)
        return result

BandPowerTransformer and RMSBandPowerTransformer do this, each exposing a materialize setting that defaults to ASYNC. That is a safety valve — it keeps the lazy graph from accumulating when nothing downstream evaluates it — but it is a default, not a decision the node makes for you. A graph that already materializes downstream should set OFF; a caller that needs the values on the host immediately should set SYNC.

Note

The _post_process hook is defined on CompositeProcessor in ezmsg-baseproc. It runs after the entire processor chain completes and receives the final output. Routing through materialize_array keeps MLX an optional dependency: it is a no-op on every other backend, and on a host without MLX installed at all.

A node should not force evaluation that its settings do not describe. An unconditional mx.eval inside a transformer is a device round-trip its caller did not ask for and cannot remove, and it compounds: a chain of n such nodes pays n stalls per message no matter how few barriers the graph actually needs. The stateful transformers in this package therefore have none — see “Choosing where to evaluate” above for why their state does not need one.

MLX quirks#

MLX’s Array API coverage is nearly complete but has a few gaps that ezmsg-sigproc works around internally:

Feature	MLX status	Workaround
`fft(norm=...)`	Not supported	Manual normalization (`/ n`, `/ sqrt(n)`)
`fftshift(axes=int)`	Needs tuple	Always pass `axes=(idx,)`
`fftfreq` / `rfftfreq`	Not available	Computed with NumPy (metadata only)
`dtype.kind`	No `.kind` attribute	`is_complex_dtype()` helper in `ezmsg.sigproc.util.array`
Window functions	Not available	Computed with NumPy, converted via `xp.asarray()`
`nan*` functions	Not available	Falls back to NumPy automatically
Boolean indexing	Not supported	Avoided in hot paths; used only in NumPy metadata code
Slice with `np.int64`	Rejected	Slice bounds cast to Python `int`

These workarounds are handled inside the transformers — user code does not need to account for them.

Limitations#

Some operations remain NumPy-only due to lack of Array API equivalents:

SciPy operations: Butterworth filtering (scipy.signal.sosfilt) and other scipy-dependent steps. Use AsArrayTransformer to convert between backends at pipeline boundaries (see RMSBandPowerTransformer for an example).
Random number generation: Modules using np.random (e.g., denormalize)
Trapezoidal integration: np.trapezoid has no Array API equivalent. RangedAggregateTransformer falls back to NumPy transparently.
Memory layout: np.require for contiguous array optimization

Metadata arrays (axis labels, coordinates) always remain as NumPy arrays since they are not performance-critical.

Adding Array API Support#

When contributing new transformers, follow this pattern:

from array_api_compat import get_namespace
from ezmsg.baseproc import BaseTransformer
from ezmsg.util.messages.axisarray import AxisArray
from ezmsg.util.messages.util import replace

class MyTransformer(BaseTransformer[MySettings, AxisArray, AxisArray]):
    def _process(self, message: AxisArray) -> AxisArray:
        xp = get_namespace(message.data)

        # Use xp instead of np for array operations
        result = xp.sqrt(xp.abs(message.data))

        return replace(message, data=result)

Key guidelines:

Call get_namespace(message.data) at the start of _process (or _reset_state for stateful transformers).
Use xp.function_name instead of np.function_name for all operations on message.data.
Note that some functions have different names: - np.concatenate → xp.concat - np.transpose → xp.permute_dims
Keep metadata operations (axis labels, etc.) as NumPy.

When a backend lacks a function (e.g., MLX has no nanmean), fall back gracefully:

func_name = "mean"
if hasattr(xp, func_name):
    result = getattr(xp, func_name)(data, axis=axis_idx)
else:
    result = np.mean(np.asarray(data), axis=axis_idx)

For CompositeProcessor subclasses that may produce MLX output, add a _post_process override to call mx.eval() (see the MLX section above).
Use portable helpers from ezmsg.sigproc.util.array when needed: is_complex_dtype, is_float_dtype, xp_asarray.