ezmsg-sigproc 2.8.0__tar.gz → 2.9.0__tar.gz

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg-sigproc
-Version: 2.8.0
+Version: 2.9.0
 Summary: Timeseries signal processing implementations in ezmsg
 Author-email: Griffin Milsap <griffin.milsap@gmail.com>, Preston Peranich <pperanich@gmail.com>, Chadwick Boulay <chadwick.boulay@gmail.com>, Kyle McGraw <kmcgraw@blackrockneuro.com>
 License-Expression: MIT

ezmsg_sigproc-2.9.0/docs/source/guides/explanations/array_api.rst

@@ -0,0 +1,156 @@
+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.
+
+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 or PyTorch tensors
+- **Framework interoperability** for integration with ML pipelines
+- **Hardware flexibility** without code changes
+
+How It Works
+------------
+
+Compatible transformers use `array-api-compat <https://github.com/data-apis/array-api-compat>`_
+to detect the input array's namespace and use the appropriate operations:
+
+.. code-block:: python
+
+    from array_api_compat import get_namespace
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)  # numpy, cupy, torch, 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:
+
+.. code-block:: python
+
+    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
+^^^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Module
+     - Description
+   * - :mod:`ezmsg.sigproc.math.abs`
+     - Absolute value
+   * - :mod:`ezmsg.sigproc.math.clip`
+     - Clip values to a range
+   * - :mod:`ezmsg.sigproc.math.log`
+     - Logarithm with configurable base
+   * - :mod:`ezmsg.sigproc.math.scale`
+     - Multiply by a constant
+   * - :mod:`ezmsg.sigproc.math.invert`
+     - Compute 1/x
+   * - :mod:`ezmsg.sigproc.math.difference`
+     - Subtract a constant (ConstDifferenceTransformer)
+
+Signal Processing
+^^^^^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Module
+     - Description
+   * - :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
+^^^^^^^^^^^^^^^^^^^^^
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Module
+     - Description
+   * - :mod:`ezmsg.sigproc.coordinatespaces`
+     - Cartesian/polar coordinate conversions
+
+Limitations
+-----------
+
+Some operations remain NumPy-only due to lack of Array API equivalents:
+
+- **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)
+
+Metadata arrays (axis labels, coordinates) typically remain as NumPy arrays
+since they are not performance-critical.
+
+Adding Array API Support
+------------------------
+
+When contributing new transformers, follow this pattern:
+
+.. code-block:: python
+
+    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:
+
+1. Call ``get_namespace(message.data)`` at the start of ``_process``
+2. Use ``xp.function_name`` instead of ``np.function_name``
+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

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/docs/source/guides/sigproc/content-sigproc.rst

@@ -4,7 +4,8 @@ ezmsg-sigproc
 Timeseries signal processing implementations in ezmsg, leveraging numpy and scipy.
 Most of the methods and classes in this extension are intended to be used in building signal processing pipelines.
 They use :class:`ezmsg.util.messages.axisarray.AxisArray` as the primary data structure for passing signals between components.
-The message's data are expected to be a numpy array.
+The message's data are typically NumPy arrays, though many transformers support the
+:doc:`Array API standard <../explanations/array_api>` for use with CuPy, PyTorch, and other backends.
 
 .. note:: Some generators might yield valid :class:`AxisArray` messages with ``.data`` size of 0.
 This may occur when the generator receives inadequate data to produce a valid output, such as when windowing or buffering.
@@ -21,3 +22,4 @@ This may occur when the generator receives inadequate data to produce a valid ou
     base
     units
     processors
+    ../explanations/array_api

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.8.0'
-__version_tuple__ = version_tuple = (2, 8, 0)
+__version__ = version = '2.9.0'
+__version_tuple__ = version_tuple = (2, 9, 0)
 
 __commit_id__ = commit_id = None

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/aggregate.py

@@ -1,3 +1,12 @@
+"""
+Aggregation operations over arrays.
+
+.. note::
+    :obj:`AggregateTransformer` supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+    :obj:`RangedAggregateTransformer` currently requires NumPy arrays.
+"""
+
 import typing
 
 import ezmsg.core as ez

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/coordinatespaces.py

@@ -3,6 +3,10 @@ Coordinate space transformations for streaming data.
 
 This module provides utilities and ezmsg nodes for transforming between
 Cartesian (x, y) and polar (r, theta) coordinate systems.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
 """
 
 from enum import Enum
@@ -11,6 +15,7 @@ from typing import Tuple
 import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
+from array_api_compat import get_namespace, is_array_api_obj
 from ezmsg.baseproc import (
     BaseTransformer,
     BaseTransformerUnit,
@@ -20,14 +25,24 @@ from ezmsg.util.messages.axisarray import AxisArray, replace
 # -- Utility functions for coordinate transformations --
 
 
+def _get_namespace_or_numpy(*args: npt.ArrayLike):
+    """Get array namespace if any arg is an array, otherwise return numpy."""
+    for arg in args:
+        if is_array_api_obj(arg):
+            return get_namespace(arg)
+    return np
+
+
 def polar2z(r: npt.ArrayLike, theta: npt.ArrayLike) -> npt.ArrayLike:
     """Convert polar coordinates to complex number representation."""
-    return r * np.exp(1j * theta)
+    xp = _get_namespace_or_numpy(r, theta)
+    return r * xp.exp(1j * theta)
 
 
 def z2polar(z: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
     """Convert complex number to polar coordinates (r, theta)."""
-    return np.abs(z), np.angle(z)
+    xp = _get_namespace_or_numpy(z)
+    return xp.abs(z), xp.atan2(xp.imag(z), xp.real(z))
 
 
 def cart2z(x: npt.ArrayLike, y: npt.ArrayLike) -> npt.ArrayLike:
@@ -37,7 +52,8 @@ def cart2z(x: npt.ArrayLike, y: npt.ArrayLike) -> npt.ArrayLike:
 
 def z2cart(z: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
     """Convert complex number to Cartesian coordinates (x, y)."""
-    return np.real(z), np.imag(z)
+    xp = _get_namespace_or_numpy(z)
+    return xp.real(z), xp.imag(z)
 
 
 def cart2pol(x: npt.ArrayLike, y: npt.ArrayLike) -> Tuple[npt.ArrayLike, npt.ArrayLike]:
@@ -90,6 +106,7 @@ class CoordinateSpacesTransformer(BaseTransformer[CoordinateSpacesSettings, Axis
     """
 
     def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
         axis = self.settings.axis or message.dims[-1]
         axis_idx = message.get_axis_idx(axis)
 
@@ -116,9 +133,9 @@ class CoordinateSpacesTransformer(BaseTransformer[CoordinateSpacesSettings, Axis
             out_a, out_b = pol2cart(component_a, component_b)
 
         # Stack results back along the same axis
-        result = np.stack([out_a, out_b], axis=axis_idx)
+        result = xp.stack([out_a, out_b], axis=axis_idx)
 
-        # Update axis labels if present
+        # Update axis labels if present (use numpy for string labels)
         axes = message.axes
         if axis in axes and hasattr(axes[axis], "data"):
             if self.settings.mode == CoordinateMode.CART2POL:

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/diff.py

@@ -1,6 +1,15 @@
+"""
+Compute differences along an axis.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
 import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
+from array_api_compat import get_namespace
 from ezmsg.baseproc import (
     BaseStatefulTransformer,
     BaseTransformerUnit,
@@ -39,23 +48,25 @@ class DiffTransformer(BaseStatefulTransformer[DiffSettings, AxisArray, AxisArray
                     self.state.last_time = ax_info.data[0] - 0.001
 
     def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
         axis = self.settings.axis or message.dims[0]
         ax_idx = message.get_axis_idx(axis)
 
-        diffs = np.diff(
-            np.concatenate((self.state.last_dat, message.data), axis=ax_idx),
+        diffs = xp.diff(
+            xp.concat((self.state.last_dat, message.data), axis=ax_idx),
             axis=ax_idx,
         )
         # Prepare last_dat for next iteration
         self.state.last_dat = slice_along_axis(message.data, slice(-1, None), axis=ax_idx)
-        # Scale by fs if requested. This convers the diff to a derivative. e.g., diff of position becomes velocity.
+        # Scale by fs if requested. This converts the diff to a derivative. e.g., diff of position becomes velocity.
         if self.settings.scale_by_fs:
             ax_info = message.get_axis(axis)
             if hasattr(ax_info, "data"):
+                # ax_info.data is typically numpy for metadata, so use np.diff here
                 dt = np.diff(np.concatenate(([self.state.last_time], ax_info.data)))
                 # Expand dt dims to match diffs
                 exp_sl = (None,) * ax_idx + (Ellipsis,) + (None,) * (message.data.ndim - ax_idx - 1)
-                diffs /= dt[exp_sl]
+                diffs /= xp.asarray(dt[exp_sl])
                 self.state.last_time = ax_info.data[-1]  # For next iteration
             else:
                 diffs /= ax_info.gain

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/linear.py

@@ -1,20 +1,22 @@
-"""Apply a linear transformation: output = scale * input + offset.
+"""
+Apply a linear transformation: output = scale * input + offset.
 
 Supports per-element scale and offset along a specified axis.
-Uses Array API for compatibility with numpy, cupy, torch, etc.
-
 For full matrix transformations, use :obj:`AffineTransformTransformer` instead.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
 """
 
 import ezmsg.core as ez
 import numpy as np
 import numpy.typing as npt
 from array_api_compat import get_namespace
+from ezmsg.baseproc import BaseStatefulTransformer, BaseTransformerUnit, processor_state
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from .base import BaseStatefulTransformer, BaseTransformerUnit, processor_state
-
 
 class LinearTransformSettings(ez.Settings):
     scale: float | list[float] | npt.ArrayLike = 1.0

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/math/abs.py

@@ -1,12 +1,16 @@
-"""Take the absolute value of the data."""
-# TODO: Array API
+"""
+Take the absolute value of the data.
 
-import numpy as np
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
+from array_api_compat import get_namespace
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ..base import BaseTransformer, BaseTransformerUnit
-
 
 class AbsSettings:
     pass
@@ -14,7 +18,8 @@ class AbsSettings:
 
 class AbsTransformer(BaseTransformer[None, AxisArray, AxisArray]):
     def _process(self, message: AxisArray) -> AxisArray:
-        return replace(message, data=np.abs(message.data))
+        xp = get_namespace(message.data)
+        return replace(message, data=xp.abs(message.data))
 
 
 class Abs(BaseTransformerUnit[None, AxisArray, AxisArray, AbsTransformer]): ...  # SETTINGS = None

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/math/add.py

@@ -5,12 +5,11 @@ import typing
 from dataclasses import dataclass, field
 
 import ezmsg.core as ez
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.baseproc.util.asio import run_coroutine_sync
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ..base import BaseTransformer, BaseTransformerUnit
-
 # --- Constant Addition (single input) ---
 
 

ezmsg_sigproc-2.9.0/src/ezmsg/sigproc/math/clip.py

@@ -0,0 +1,48 @@
+"""
+Clips the data to be within the specified range.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
+import ezmsg.core as ez
+from array_api_compat import get_namespace
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
+from ezmsg.util.messages.axisarray import AxisArray
+from ezmsg.util.messages.util import replace
+
+
+class ClipSettings(ez.Settings):
+    min: float | None = None
+    """Lower clip bound. If None, no lower clipping is applied."""
+
+    max: float | None = None
+    """Upper clip bound. If None, no upper clipping is applied."""
+
+
+class ClipTransformer(BaseTransformer[ClipSettings, AxisArray, AxisArray]):
+    def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
+        return replace(
+            message,
+            data=xp.clip(message.data, self.settings.min, self.settings.max),
+        )
+
+
+class Clip(BaseTransformerUnit[ClipSettings, AxisArray, AxisArray, ClipTransformer]):
+    SETTINGS = ClipSettings
+
+
+def clip(min: float | None = None, max: float | None = None) -> ClipTransformer:
+    """
+    Clips the data to be within the specified range.
+
+    Args:
+        min: Lower clip bound. If None, no lower clipping is applied.
+        max: Upper clip bound. If None, no upper clipping is applied.
+
+    Returns:
+        :obj:`ClipTransformer`.
+    """
+    return ClipTransformer(ClipSettings(min=min, max=max))

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/math/difference.py

@@ -1,16 +1,22 @@
-"""Take the difference between 2 signals or between a signal and a constant value."""
+"""
+Take the difference between 2 signals or between a signal and a constant value.
+
+.. note::
+    :obj:`ConstDifferenceTransformer` supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+    :obj:`DifferenceProcessor` (two-input difference) currently requires NumPy arrays.
+"""
 
 import asyncio
 import typing
 from dataclasses import dataclass, field
 
 import ezmsg.core as ez
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.baseproc.util.asio import run_coroutine_sync
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ..base import BaseTransformer, BaseTransformerUnit
-
 
 class ConstDifferenceSettings(ez.Settings):
     value: float = 0.0

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/math/invert.py

@@ -1,10 +1,15 @@
-"""1/data transformer."""
+"""
+Compute the multiplicative inverse (1/x) of the data.
 
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
+
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ..base import BaseTransformer, BaseTransformerUnit
-
 
 class InvertTransformer(BaseTransformer[None, AxisArray, AxisArray]):
     def _process(self, message: AxisArray) -> AxisArray:

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/math/log.py

@@ -1,13 +1,17 @@
-"""Take the logarithm of the data."""
+"""
+Take the logarithm of the data.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
 
-# TODO: Array API
 import ezmsg.core as ez
-import numpy as np
+from array_api_compat import get_namespace
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ..base import BaseTransformer, BaseTransformerUnit
-
 
 class LogSettings(ez.Settings):
     base: float = 10.0
@@ -19,10 +23,17 @@ class LogSettings(ez.Settings):
 
 class LogTransformer(BaseTransformer[LogSettings, AxisArray, AxisArray]):
     def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
         data = message.data
-        if self.settings.clip_zero and np.any(data <= 0) and np.issubdtype(data.dtype, np.floating):
-            data = np.clip(data, a_min=np.finfo(data.dtype).tiny, a_max=None)
-        return replace(message, data=np.log(data) / np.log(self.settings.base))
+        if self.settings.clip_zero:
+            # Check if any values are <= 0 and dtype is floating point
+            has_non_positive = bool(xp.any(data <= 0))
+            is_floating = xp.isdtype(data.dtype, "real floating")
+            if has_non_positive and is_floating:
+                # Use smallest_normal (Array API equivalent of numpy's finfo.tiny)
+                min_val = xp.finfo(data.dtype).smallest_normal
+                data = xp.clip(data, min_val, None)
+        return replace(message, data=xp.log(data) / xp.log(self.settings.base))
 
 
 class Log(BaseTransformerUnit[LogSettings, AxisArray, AxisArray, LogTransformer]):

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/math/scale.py

@@ -1,11 +1,16 @@
-"""Scale the data by a constant factor."""
+"""
+Scale the data by a constant factor.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+"""
 
 import ezmsg.core as ez
+from ezmsg.baseproc import BaseTransformer, BaseTransformerUnit
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
-from ..base import BaseTransformer, BaseTransformerUnit
-
 
 class ScaleSettings(ez.Settings):
     scale: float = 1.0

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/src/ezmsg/sigproc/transpose.py

@@ -1,7 +1,17 @@
+"""
+Transpose or permute array dimensions.
+
+.. note::
+    This module supports the :doc:`Array API standard </guides/explanations/array_api>`,
+    enabling use with NumPy, CuPy, PyTorch, and other compatible array libraries.
+    Memory layout optimization (C/F order) only applies to NumPy arrays.
+"""
+
 from types import EllipsisType
 
 import ezmsg.core as ez
 import numpy as np
+from array_api_compat import get_namespace, is_numpy_array
 from ezmsg.baseproc import (
     BaseStatefulTransformer,
     BaseTransformerUnit,
@@ -84,6 +94,7 @@ class TransposeTransformer(BaseStatefulTransformer[TransposeSettings, AxisArray,
         return super().__call__(message)
 
     def _process(self, message: AxisArray) -> AxisArray:
+        xp = get_namespace(message.data)
         if self.state.axes_ints is None:
             # No transpose required
             if self.settings.order is None:
@@ -91,15 +102,19 @@ class TransposeTransformer(BaseStatefulTransformer[TransposeSettings, AxisArray,
                 # Note: We should not be able to reach here because it should be shortcutted at passthrough.
                 msg_out = message
             else:
-                # If the memory is already contiguous in the correct order, np.require won't do anything.
-                msg_out = replace(
-                    message,
-                    data=np.require(message.data, requirements=self.settings.order.upper()[0]),
-                )
+                # Memory layout optimization only applies to numpy arrays
+                if is_numpy_array(message.data):
+                    msg_out = replace(
+                        message,
+                        data=np.require(message.data, requirements=self.settings.order.upper()[0]),
+                    )
+                else:
+                    msg_out = message
         else:
             dims_out = [message.dims[ix] for ix in self.state.axes_ints]
-            data_out = np.transpose(message.data, axes=self.state.axes_ints)
-            if self.settings.order is not None:
+            data_out = xp.permute_dims(message.data, axes=self.state.axes_ints)
+            if self.settings.order is not None and is_numpy_array(data_out):
+                # Memory layout optimization only applies to numpy arrays
                 data_out = np.require(data_out, requirements=self.settings.order.upper()[0])
             msg_out = replace(
                 message,

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/tests/integration/ezmsg/test_add_system.py

@@ -76,7 +76,7 @@ def test_add_two_signals_system(
     messages: list[AxisArray] = [_ for _ in message_log(test_filename)]
     os.remove(test_filename)
 
-    assert len(messages) == n_messages
+    assert len(messages) >= n_messages
 
     # Verify each message has correct shape
     for msg in messages:
@@ -137,7 +137,7 @@ def test_const_add_system(
     messages: list[AxisArray] = [_ for _ in message_log(test_filename)]
     os.remove(test_filename)
 
-    assert len(messages) == n_messages
+    assert len(messages) >= n_messages
 
     # Verify the constant was added
     data = np.concatenate([_.data for _ in messages]).squeeze()

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/tests/integration/ezmsg/test_difference_system.py

@@ -74,7 +74,7 @@ def test_difference_two_signals_system(
     messages: list[AxisArray] = [_ for _ in message_log(test_filename)]
     os.remove(test_filename)
 
-    assert len(messages) == n_messages
+    assert len(messages) >= n_messages
 
     # Verify each message has correct shape
     for msg in messages:
@@ -135,7 +135,7 @@ def test_const_difference_system(
     messages: list[AxisArray] = [_ for _ in message_log(test_filename)]
     os.remove(test_filename)
 
-    assert len(messages) == n_messages
+    assert len(messages) >= n_messages
 
     # Verify the constant was subtracted
     data = np.concatenate([_.data for _ in messages]).squeeze()
@@ -191,7 +191,7 @@ def test_const_difference_subtrahend_false_system(
     messages: list[AxisArray] = [_ for _ in message_log(test_filename)]
     os.remove(test_filename)
 
-    assert len(messages) == n_messages
+    assert len(messages) >= n_messages
 
     # Verify: value - input
     data = np.concatenate([_.data for _ in messages]).squeeze()

{ezmsg_sigproc-2.8.0 → ezmsg_sigproc-2.9.0}/tests/integration/ezmsg/test_sampler_system.py

@@ -89,7 +89,7 @@ def test_sampler_system(test_name: str | None = None):
     messages: list[SampleTriggerMessage] = [_ for _ in message_log(test_filename)]
     os.remove(test_filename)
     ez.logger.info(f"Analyzing recording of {len(messages)} messages...")
-    assert len(messages) == n_msgs
+    assert len(messages) >= n_msgs
     assert all([_.sample.data.shape == (int(freq * sample_dur), 1) for _ in messages])
     # Test the sample window slice vs the trigger timestamps
     latencies = [_.sample.axes["time"].offset - (_.trigger.timestamp + _.trigger.period[0]) for _ in messages]