ezmsg-sigproc 2.13.1__py3-none-any.whl → 2.14.0__py3-none-any.whl

ezmsg/sigproc/__version__.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.13.1'
-__version_tuple__ = version_tuple = (2, 13, 1)
+__version__ = version = '2.14.0'
+__version_tuple__ = version_tuple = (2, 14, 0)
 
 __commit_id__ = commit_id = None

ezmsg/sigproc/affinetransform.py

@@ -23,6 +23,8 @@ from ezmsg.baseproc import (
 from ezmsg.util.messages.axisarray import AxisArray, AxisBase
 from ezmsg.util.messages.util import replace
 
+from ezmsg.sigproc.util.array import array_device, is_float_dtype, xp_asarray, xp_create
+
 
 def _find_block_diagonal_clusters(weights: np.ndarray) -> list[tuple[np.ndarray, np.ndarray]] | None:
     """Detect block-diagonal structure in a weight matrix.
@@ -245,13 +247,25 @@ class AffineTransformTransformer(
 
             self._state.new_axis = replace(message.axes[axis], data=np.array(new_labels))
 
-        # Convert to match message.data namespace for efficient operations in _process
+        # Convert to match message.data namespace and device for _process.
+        # Weights are stored as numpy float64 after cluster detection; some
+        # devices (e.g. MPS) don't support float64, so we downcast weight
+        # arrays to the message's dtype when the message is floating-point.
         xp = get_namespace(message.data)
+        dev = array_device(message.data)
+        msg_dt = message.data.dtype
+        # Downcast weights dtype only for float message data (avoids casting
+        # float weights to integer when message data happens to be int).
+        w_dt = msg_dt if is_float_dtype(xp, msg_dt) else None
         if self._state.weights is not None:
-            self._state.weights = xp.asarray(self._state.weights)
+            self._state.weights = xp_asarray(xp, self._state.weights, dtype=w_dt, device=dev)
         if self._state.clusters is not None:
             self._state.clusters = [
-                (xp.asarray(in_idx), xp.asarray(out_idx), xp.asarray(sub_w))
+                (
+                    xp_asarray(xp, in_idx, device=dev),
+                    xp_asarray(xp, out_idx, device=dev),
+                    xp_asarray(xp, sub_w, dtype=w_dt, device=dev),
+                )
                 for in_idx, out_idx, sub_w in self._state.clusters
             ]
 
@@ -345,7 +359,7 @@ class AffineTransformTransformer(
 
         # Pre-allocate output (omitted channels stay zero)
         out_shape = data.shape[:-1] + (self._state.n_out,)
-        result = xp.zeros(out_shape, dtype=data.dtype)
+        result = xp_create(xp.zeros, out_shape, dtype=data.dtype, device=array_device(data))
 
         for in_idx, out_idx, sub_weights in self._state.clusters:
             chunk = xp.take(data, in_idx, axis=data.ndim - 1)
@@ -371,7 +385,10 @@ class AffineTransformTransformer(
                 # The weights are stacked A|B where A is the transform and B is a single row
                 #  in the equation y = Ax + B. This supports NeuroKey's weights matrices.
                 sample_shape = data.shape[:axis_idx] + (1,) + data.shape[axis_idx + 1 :]
-                data = xp.concat((data, xp.ones(sample_shape, dtype=data.dtype)), axis=axis_idx)
+                data = xp.concat(
+                    (data, xp_create(xp.ones, sample_shape, dtype=data.dtype, device=array_device(data))),
+                    axis=axis_idx,
+                )
 
             if axis_idx in [-1, len(message.dims) - 1]:
                 data = xp.matmul(data, self._state.weights)

ezmsg/sigproc/util/array.py

@@ -0,0 +1,59 @@
+"""Portable helpers for Array API interoperability.
+
+These utilities smooth over differences between Array API libraries
+(NumPy, PyTorch, MLX, CuPy, etc.) — in particular around ``device``
+placement and ``dtype`` introspection, which are not uniformly supported.
+"""
+
+import numpy as np
+
+
+def array_device(x):
+    """Return the device of an array, or ``None`` for device-less libraries."""
+    try:
+        from array_api_compat import device
+
+        return device(x)
+    except (AttributeError, TypeError):
+        return None
+
+
+def xp_asarray(xp, obj, *, dtype=None, device=None):
+    """Portable ``xp.asarray`` that omits unsupported kwargs.
+
+    Some Array API libraries (e.g. MLX) don't accept a ``device`` keyword.
+    This helper builds the kwargs dict dynamically so that only supported
+    arguments are forwarded.
+    """
+    kwargs = {}
+    if dtype is not None:
+        kwargs["dtype"] = dtype
+    if device is not None:
+        kwargs["device"] = device
+    return xp.asarray(obj, **kwargs)
+
+
+def xp_create(fn, *args, dtype=None, device=None, **extra):
+    """Call a creation function (``zeros``, ``ones``, ``eye``) portably.
+
+    Omits ``device`` if it is ``None`` (for libraries that don't support it).
+    """
+    kwargs = dict(extra)
+    if dtype is not None:
+        kwargs["dtype"] = dtype
+    if device is not None:
+        kwargs["device"] = device
+    return fn(*args, **kwargs)
+
+
+def is_float_dtype(xp, dtype) -> bool:
+    """Check whether *dtype* is a real floating-point type, portably."""
+    try:
+        return xp.isdtype(dtype, "real floating")
+    except AttributeError:
+        pass
+    # Fallback for libraries without isdtype (e.g. MLX).
+    try:
+        return xp.issubdtype(dtype, xp.floating)
+    except (AttributeError, TypeError):
+        return np.issubdtype(np.dtype(dtype), np.floating)

{ezmsg_sigproc-2.13.1.dist-info → ezmsg_sigproc-2.14.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ezmsg-sigproc
-Version: 2.13.1
+Version: 2.14.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.13.1.dist-info → ezmsg_sigproc-2.14.0.dist-info}/RECORD

@@ -1,8 +1,8 @@
 ezmsg/sigproc/__init__.py,sha256=8K4IcOA3-pfzadoM6s2Sfg5460KlJUocGgyTJTJl96U,52
-ezmsg/sigproc/__version__.py,sha256=LrOgVsOxSmuj7RbXPLo3yIvC77lH9VlW4tk7Ihs46rY,706
+ezmsg/sigproc/__version__.py,sha256=ho7Izv0FLYm0b-otHGFFngHF9OeMd4aMhL-uSzYcJ5U,706
 ezmsg/sigproc/activation.py,sha256=83vnTa3ZcC4Q3VSWcGfaqhCEqYRNySUOyVpMHZXfz-c,2755
 ezmsg/sigproc/adaptive_lattice_notch.py,sha256=ThUR48mbSHuThkimtD0j4IXNMrOVcpZgGhE7PCYfXhU,8818
-ezmsg/sigproc/affinetransform.py,sha256=mjA21DRVYm0kS2tK7dNR_mU5XAxDbJGuuhnuzz0gtw4,21679
+ezmsg/sigproc/affinetransform.py,sha256=rnbQ2QtLq15eoFrqL-ij12yXSoLLq5ncOwRRzQgB1lA,22574
 ezmsg/sigproc/aggregate.py,sha256=7Hdz1m-S6Cl9h0oRQHeS_UTGBemhOB4XdFyX6cGcdHo,9362
 ezmsg/sigproc/bandpower.py,sha256=dAhH56sUrXNhcRFymTTwjdM_KcU5OxFzrR_sxIPAxyw,2264
 ezmsg/sigproc/base.py,sha256=SJvKEb8gw6mUMwlV5sH0iPG0bXrgS8tvkPwhI-j89MQ,3672
@@ -55,6 +55,7 @@ ezmsg/sigproc/math/log.py,sha256=JhjSqLnQnvx_3F4txRYHuUPSJ12Yj2HvRTsCMNvlxpo,202
 ezmsg/sigproc/math/pow.py,sha256=0sdlXFUEBXmpEV_i75oshGRjMguv8L13nLt7hlvdX3E,1284
 ezmsg/sigproc/math/scale.py,sha256=4_xHcHNuf13E1fxIF5vbkPfkN4En6zkfPIKID7lCERk,1133
 ezmsg/sigproc/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+ezmsg/sigproc/util/array.py,sha256=S6wTGRbtJpU5I9u2t0MTUAMSKkm1PJref145s87xHLg,1831
 ezmsg/sigproc/util/asio.py,sha256=aAj0e7OoBvkRy28k05HL2s9YPCTxOddc05xMN-qd4lQ,577
 ezmsg/sigproc/util/axisarray_buffer.py,sha256=TGDeC6CXmmp7OUuiGd6xYQijRGYDE4QGdWxjK5Vs3nE,14057
 ezmsg/sigproc/util/buffer.py,sha256=83Gm0IuowmcMlXgLFB_rz8_ZPhkwG4DNNejyWJDKJl8,19658
@@ -62,7 +63,7 @@ ezmsg/sigproc/util/message.py,sha256=ppN3IYtIAwrxWG9JOvgWFn1wDdIumkEzYFfqpH9VQkY
 ezmsg/sigproc/util/profile.py,sha256=eVOo9pXgusrnH1yfRdd2RsM7Dbe2UpyC0LJ9MfGpB08,416
 ezmsg/sigproc/util/sparse.py,sha256=NjbJitCtO0B6CENTlyd9c-lHEJwoCan-T3DIgPyeShw,4834
 ezmsg/sigproc/util/typeresolution.py,sha256=fMFzLi63dqCIclGFLcMdM870OYxJnkeWw6aWKNMk718,362
-ezmsg_sigproc-2.13.1.dist-info/METADATA,sha256=mjyCiq1zy3JbgayLzq3BQLB-i1LOouUeGHe0tn3v7MY,1909
-ezmsg_sigproc-2.13.1.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-ezmsg_sigproc-2.13.1.dist-info/licenses/LICENSE,sha256=seu0tKhhAMPCUgc1XpXGGaCxY1YaYvFJwqFuQZAl2go,1100
-ezmsg_sigproc-2.13.1.dist-info/RECORD,,
+ezmsg_sigproc-2.14.0.dist-info/METADATA,sha256=Ml-3BXd3vRBRe3H1j1pmWdlXmvnvphbYQ7SH_n4pw9Y,1909
+ezmsg_sigproc-2.14.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+ezmsg_sigproc-2.14.0.dist-info/licenses/LICENSE,sha256=seu0tKhhAMPCUgc1XpXGGaCxY1YaYvFJwqFuQZAl2go,1100
+ezmsg_sigproc-2.14.0.dist-info/RECORD,,