Composite mlx example

.github/workflows/python-tests.yml

@@ -34,4 +34,4 @@ jobs:
           uv tool run ruff check --output-format=github src
 
       - name: Run tests
-        run: uv run pytest tests
+        run: uv run pytest tests --benchmark-skip

.gitignore

@@ -147,3 +147,4 @@ cython_debug/
 src/ezmsg/sigproc/__version__.py
 uv.lock
 *.local.json
+tmp/

docs/source/guides/explanations/array_api.rst

@@ -3,7 +3,7 @@ Array API Support
 
 ezmsg-sigproc provides support for the `Python Array API standard
 <https://data-apis.org/array-api/>`_, enabling many transformers to work with
-arrays from different backends such as NumPy, CuPy, PyTorch, and JAX.
+arrays from different backends such as NumPy, CuPy, PyTorch, JAX, and MLX.
 
 What is the Array API?
 ----------------------
@@ -12,7 +12,8 @@ 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 or PyTorch tensors
+- **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
 
@@ -27,7 +28,7 @@ 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, etc.
+        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)
 
@@ -90,14 +91,16 @@ Signal Processing
 
    * - Module
      - Description
+   * - :mod:`ezmsg.sigproc.spectrum`
+     - FFT-based spectrum (SpectrumTransformer)
+   * - :mod:`ezmsg.sigproc.aggregate`
+     - Aggregate operations (AggregateTransformer, RangedAggregateTransformer)
    * - :mod:`ezmsg.sigproc.diff`
      - Compute differences along an axis
    * - :mod:`ezmsg.sigproc.transpose`
      - Transpose/permute array dimensions
    * - :mod:`ezmsg.sigproc.linear`
      - Per-channel linear transform (scale + offset)
-   * - :mod:`ezmsg.sigproc.aggregate`
-     - Aggregate operations (AggregateTransformer only)
 
 Coordinate Transforms
 ^^^^^^^^^^^^^^^^^^^^^
@@ -111,17 +114,163 @@ Coordinate Transforms
    * - :mod:`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:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Module
+     - Description
+   * - :mod:`ezmsg.sigproc.bandpower`
+     - BandPowerTransformer (spectrogram + ranged aggregate)
+   * - :mod:`ezmsg.sigproc.singlebandpow`
+     - RMSBandPowerTransformer (with explicit ``backend`` setting; only after initial IIR filter)
+
+MLX on Apple Silicon
+--------------------
+
+`MLX <https://github.com/ml-explore/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:
+
+.. code-block:: python
+
+    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()``:
+
+.. code-block:: python
+
+    result = proc(message)
+    mx.eval(result.data)  # Forces computation to complete
+
+For ``CompositeProcessor`` pipelines (like ``BandPowerTransformer``), you can
+override ``_post_process`` to call ``mx.eval()`` automatically so that every
+output is fully materialized:
+
+.. code-block:: python
+
+    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:
+                try:
+                    import mlx.core as mx
+
+                    if isinstance(result.data, mx.array):
+                        mx.eval(result.data)
+                except ImportError:
+                    pass
+            return result
+
+This pattern is used by ``BandPowerTransformer`` and ``RMSBandPowerTransformer``.
+It ensures downstream consumers (ezmsg Units, visualization, logging) receive
+fully evaluated arrays without needing to know about MLX internals.
+It also provides a safety valve so the lazy graph does not accumulate if the graph
+is not evaluated at the right time downstream.
+
+.. 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. The ``try``/``except ImportError`` pattern
+    keeps MLX as an optional dependency.
+
+MLX quirks
+^^^^^^^^^^
+
+MLX's Array API coverage is nearly complete but has a few gaps that
+ezmsg-sigproc works around internally:
+
+.. list-table::
+   :header-rows: 1
+   :widths: 25 35 40
+
+   * - 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``)
-- **SciPy operations**: Filtering (``scipy.signal.lfilter``), FFT, wavelets
-- **Advanced indexing**: Some slicing operations for metadata handling
-- **Memory layout**: ``np.require`` for contiguous array optimization (NumPy only)
+- **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) typically remain as NumPy arrays
+Metadata arrays (axis labels, coordinates) always remain as NumPy arrays
 since they are not performance-critical.
 
 Adding Array API Support
@@ -147,10 +296,27 @@ When contributing new transformers, follow this pattern:
 
 Key guidelines:
 
-1. Call ``get_namespace(message.data)`` at the start of ``_process``
-2. Use ``xp.function_name`` instead of ``np.function_name``
+1. Call ``get_namespace(message.data)`` at the start of ``_process`` (or
+   ``_reset_state`` for stateful transformers).
+2. Use ``xp.function_name`` instead of ``np.function_name`` for all operations
+   on ``message.data``.
 3. Note that some functions have different names:
    - ``np.concatenate`` → ``xp.concat``
    - ``np.transpose`` → ``xp.permute_dims``
-4. Keep metadata operations (axis labels, etc.) as NumPy
-5. Use in-place operations (``/=``, ``*=``) where possible for efficiency
+4. Keep metadata operations (axis labels, etc.) as NumPy.
+5. When a backend lacks a function (e.g., MLX has no ``nanmean``), fall back
+   gracefully:
+
+   .. code-block:: python
+
+       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)
+
+6. For ``CompositeProcessor`` subclasses that may produce MLX output, add
+   a ``_post_process`` override to call ``mx.eval()`` (see the MLX section
+   above).
+7. Use portable helpers from ``ezmsg.sigproc.util.array`` when needed:
+   ``is_complex_dtype``, ``is_float_dtype``, ``xp_asarray``.

pyproject.toml

@@ -9,12 +9,12 @@ authors = [
 ]
 license = "MIT"
 readme = "README.md"
-requires-python = ">=3.10.15"
+requires-python = ">=3.10"
 dynamic = ["version"]
 dependencies = [
   "array-api-compat>=1.11.1",
-  "ezmsg-baseproc>=1.3.0",
-  "ezmsg>=3.6.0",
+  "ezmsg[axisarray]>=3.7.2",
+  "ezmsg-baseproc>=1.5.0",
   "mlx>=0.18.0; sys_platform == 'darwin' and platform_machine == 'arm64'",
   "numba>=0.61.0",
   "numpy>=1.26.0",
@@ -29,6 +29,7 @@ dev = [
   "pre-commit>=4.2.0",
   "jupyter>=1.1.1",
   "scipy-stubs>=1.15.3.0",
+  "matplotlib>=3.10.8",
   {include-group = "lint"},
   {include-group = "test"},
   {include-group = "docs"},
@@ -39,6 +40,7 @@ lint = [
 test = [
   "frozendict>=2.4.4",
   "pytest-asyncio>=0.24.0",
+  "pytest-benchmark>=5.2.3",
   "pytest-cov>=5.0.0",
   "pytest>=8.3.3",
 ]
@@ -70,7 +72,10 @@ packages = ["src/ezmsg"]
 asyncio_mode = "auto"
 asyncio_default_fixture_loop_scope = "function"
 norecursedirs = "tests/helpers"
-addopts = "-p no:warnings"
+addopts = "-p no:warnings --benchmark-group-by=group,param:n_channels -m 'not sosfilt'"
+markers = [
+  "sosfilt: marks sosfilt development benchmarks (deselected by default, run with: -m sosfilt)",
+]
 
 [tool.ruff]
 line-length = 120