ezmsg.sigproc.asarray#

Convert AxisArray data to a target array backend.

This module provides a transformer that converts AxisArray payloads between array backends (NumPy, MLX, PyTorch, CuPy, JAX). Useful for wiring a conversion step between nodes — e.g., numpy → MLX before a GPU-accelerated filter, or MLX → numpy before a scipy-dependent node.

Note

This module supports the Array API standard, enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.

Classes

class ArrayBackend(*values)#

Bases: StrEnum

numpy = 'numpy'#
torch = 'torch'#
class AsArray(*args, settings=None, **kwargs)[source]#

Bases: BaseTransformerUnit[AsArraySettings, AxisArray, AxisArray, AsArrayTransformer]

Parameters:

settings (Settings | None)

SETTINGS#

alias of AsArraySettings

class AsArraySettings(backend: ArrayBackend = <ArrayBackend.numpy: 'numpy'>, dtype: str | None = None, mlx_cache_limit_mb: float | None = 512.0)[source]#

Bases: Settings

Parameters:
backend: ArrayBackend = 'numpy'#

Target array backend.

dtype: str | None = None#

Target dtype as a string (e.g. “float32”, “float64”). None keeps the original dtype.

__init__(backend=ArrayBackend.numpy, dtype=None, mlx_cache_limit_mb=512.0)#
Parameters:
Return type:

None

mlx_cache_limit_mb: float | None = 512.0#

Cap the MLX buffer cache (MiB) for the process that runs this node.

Applied only when backend is MLX, once per process, on the first message. None leaves MLX’s default, which is the size of the machine.

Sizing: one distinct message shape costs roughly 50x the message payload in cached buffers – about 20 intermediates across a typical chain, each keeping its own size class. Measured cache for a steady-state chain, against samples x channels x 4 bytes per message: 46x at 256 ch x 1200 samples, 67x at 256 ch x 300. So:

limit_MiB ~= 50 * message_MiB * (distinct shapes to keep hot)

A 256-channel, 300-sample float32 message is 0.29 MiB, so ~15 MiB per shape and the 512 MiB default holds ~26 distinct shapes. Steady state needs only one; the rest is headroom for the varying-length messages a stall produces. Raise it if the graph legitimately cycles through many shapes; the floor is one working set (~1x the 50 * term), and 0 disables caching entirely at a measured 40% throughput cost.

This is a process-global MLX setting, so it is shared with anything else using MLX in the same process. The default suits streaming graphs; large offline batch work in the same process may want it raised or set to None.

class AsArrayTransformer(*args, settings=None, **kwargs)[source]#

Bases: BaseTransformer[AsArraySettings, AxisArray, AxisArray]

Parameters:

settings (SettingsType)