ezmsg-learn 1.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

ezmsg/learn/process/sklearn.py

@@ -5,12 +5,11 @@ import typing
 import ezmsg.core as ez
 import numpy as np
 import pandas as pd
-from ezmsg.sigproc.base import (
+from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
     processor_state,
 )
-from ezmsg.sigproc.sampler import SampleMessage
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
@@ -45,11 +44,7 @@ class SklearnModelState:
     chan_ax: AxisArray.CoordinateAxis | None = None
 
 
-class SklearnModelProcessor(
-    BaseAdaptiveTransformer[
-        SklearnModelSettings, AxisArray, AxisArray, SklearnModelState
-    ]
-):
+class SklearnModelProcessor(BaseAdaptiveTransformer[SklearnModelSettings, AxisArray, AxisArray, SklearnModelState]):
     """
     Processor that wraps a scikit-learn, River, or HMMLearn model for use in the ezmsg framework.
 
@@ -115,40 +110,30 @@ class SklearnModelProcessor(
             if hasattr(self._state.model, "n_features_in_"):
                 expected = self._state.model.n_features_in_
                 if expected != n_input:
-                    raise ValueError(
-                        f"Model expects {expected} features, but got {n_input}"
-                    )
+                    raise ValueError(f"Model expects {expected} features, but got {n_input}")
         else:
             # No checkpoint, initialize from scratch
             self._init_model()
 
-    def partial_fit(self, message: SampleMessage) -> None:
-        X = message.sample.data
-        y = message.trigger.value
+    def partial_fit(self, message: AxisArray) -> None:
+        X = message.data
+        y = message.attrs["trigger"].value
         if self._state.model is None:
-            self._reset_state(message.sample)
+            self._reset_state(message)
         if hasattr(self._state.model, "partial_fit"):
             kwargs = {}
             if self.settings.partial_fit_classes is not None:
                 kwargs["classes"] = self.settings.partial_fit_classes
             self._state.model.partial_fit(X, y, **kwargs)
         elif hasattr(self._state.model, "learn_many"):
-            df_X = pd.DataFrame(
-                {
-                    k: v
-                    for k, v in zip(
-                        message.sample.axes["ch"].data, message.sample.data.T
-                    )
-                }
-            )
+            df_X = pd.DataFrame({k: v for k, v in zip(message.axes["ch"].data, message.data.T)})
             name = (
-                message.trigger.value.axes["ch"].data[0]
-                if hasattr(message.trigger.value, "axes")
-                and "ch" in message.trigger.value.axes
+                message.attrs["trigger"].value.axes["ch"].data[0]
+                if hasattr(message.attrs["trigger"].value, "axes") and "ch" in message.attrs["trigger"].value.axes
                 else "target"
             )
             ser_y = pd.Series(
-                data=np.asarray(message.trigger.value.data).flatten(),
+                data=np.asarray(message.attrs["trigger"].value.data).flatten(),
                 name=name,
             )
             self._state.model.learn_many(df_X, ser_y)
@@ -158,9 +143,7 @@ class SklearnModelProcessor(
                 features = {f"f{i}": xi[i] for i in range(len(xi))}
                 self._state.model.learn_one(features, yi)
         else:
-            raise NotImplementedError(
-                "Model does not support partial_fit or learn_many"
-            )
+            raise NotImplementedError("Model does not support partial_fit or learn_many")
 
     def fit(self, X: np.ndarray, y: np.ndarray) -> None:
         if self._state.model is None:
@@ -192,9 +175,7 @@ class SklearnModelProcessor(
 
     def _process(self, message: AxisArray) -> AxisArray:
         if self._state.model is None:
-            raise RuntimeError(
-                "Model has not been fit yet. Call `fit()` or `partial_fit()` before processing."
-            )
+            raise RuntimeError("Model has not been fit yet. Call `fit()` or `partial_fit()` before processing.")
         X = message.data
         original_shape = X.shape
         n_input = X.shape[message.get_axis_idx("ch")]
@@ -204,9 +185,7 @@ class SklearnModelProcessor(
         if hasattr(self._state.model, "n_features_in_"):
             expected = self._state.model.n_features_in_
             if expected != n_input:
-                raise ValueError(
-                    f"Model expects {expected} features, but got {n_input}"
-                )
+                raise ValueError(f"Model expects {expected} features, but got {n_input}")
 
         if hasattr(self._state.model, "predict"):
             y_pred = self._state.model.predict(X)
@@ -216,14 +195,7 @@ class SklearnModelProcessor(
             y_pred = np.array(list(y_pred))
         elif hasattr(self._state.model, "predict_one"):
             # river's random forest does not support predict_many
-            y_pred = np.array(
-                [
-                    self._state.model.predict_one(
-                        {f"f{i}": xi[i] for i in range(len(xi))}
-                    )
-                    for xi in X
-                ]
-            )
+            y_pred = np.array([self._state.model.predict_one({f"f{i}": xi[i] for i in range(len(xi))}) for xi in X])
         else:
             raise NotImplementedError("Model does not support predict or predict_many")
 
@@ -235,9 +207,7 @@ class SklearnModelProcessor(
         y_pred = y_pred.reshape(output_shape)
 
         if self._state.chan_ax is None:
-            self._state.chan_ax = AxisArray.CoordinateAxis(
-                data=np.arange(output_shape[1]), dims=["ch"]
-            )
+            self._state.chan_ax = AxisArray.CoordinateAxis(data=np.arange(output_shape[1]), dims=["ch"])
 
         return replace(
             message,
@@ -246,11 +216,7 @@ class SklearnModelProcessor(
         )
 
 
-class SklearnModelUnit(
-    BaseAdaptiveTransformerUnit[
-        SklearnModelSettings, AxisArray, AxisArray, SklearnModelProcessor
-    ]
-):
+class SklearnModelUnit(BaseAdaptiveTransformerUnit[SklearnModelSettings, AxisArray, AxisArray, SklearnModelProcessor]):
     """
     Unit wrapper for the `SklearnModelProcessor`.
 

ezmsg/learn/process/slda.py

@@ -2,10 +2,10 @@ import typing
 
 import ezmsg.core as ez
 import numpy as np
-from ezmsg.sigproc.base import (
+from ezmsg.baseproc import (
     BaseStatefulTransformer,
-    processor_state,
     BaseTransformerUnit,
+    processor_state,
 )
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
@@ -25,9 +25,7 @@ class SLDAState:
     out_template: typing.Optional[ClassifierMessage] = None
 
 
-class SLDATransformer(
-    BaseStatefulTransformer[SLDASettings, AxisArray, ClassifierMessage, SLDAState]
-):
+class SLDATransformer(BaseStatefulTransformer[SLDASettings, AxisArray, ClassifierMessage, SLDAState]):
     def _reset_state(self, message: AxisArray) -> None:
         if self.settings.settings_path[-4:] == ".mat":
             # Expects a very specific format from a specific project. Not for general use.
@@ -67,9 +65,7 @@ class SLDATransformer(
             dims=[self.settings.axis, "classes"],
             axes={
                 self.settings.axis: message.axes[self.settings.axis],
-                "classes": AxisArray.CoordinateAxis(
-                    data=np.array(out_labels), dims=["classes"]
-                ),
+                "classes": AxisArray.CoordinateAxis(data=np.array(out_labels), dims=["classes"]),
             },
             labels=out_labels,
             key=message.key,
@@ -80,10 +76,7 @@ class SLDATransformer(
         X = np.moveaxis(message.data, samp_ax_idx, 0)
 
         if X.shape[0]:
-            if (
-                isinstance(self.settings.settings_path, str)
-                and self.settings.settings_path[-4:] == ".mat"
-            ):
+            if isinstance(self.settings.settings_path, str) and self.settings.settings_path[-4:] == ".mat":
                 # Assumes F-contiguous weights
                 pred_probas = []
                 for samp in X:
@@ -113,7 +106,5 @@ class SLDATransformer(
             return self.state.out_template
 
 
-class SLDA(
-    BaseTransformerUnit[SLDASettings, AxisArray, ClassifierMessage, SLDATransformer]
-):
+class SLDA(BaseTransformerUnit[SLDASettings, AxisArray, ClassifierMessage, SLDATransformer]):
     SETTINGS = SLDASettings

ezmsg/learn/process/ssr.py

@@ -0,0 +1,374 @@
+"""Self-supervised regression framework and LRR implementation.
+
+This module provides a general framework for self-supervised channel
+regression via :class:`SelfSupervisedRegressionTransformer`, and a
+concrete implementation — Linear Regression Rereferencing (LRR) — via
+:class:`LRRTransformer`.
+
+**Framework.**  The base class accumulates the channel covariance
+``C = X^T X`` and solves per-cluster ridge regressions to obtain a weight
+matrix *W*.  Subclasses define what to *do* with *W* by implementing
+:meth:`~SelfSupervisedRegressionTransformer._on_weights_updated` and
+:meth:`~SelfSupervisedRegressionTransformer._process`.
+
+**LRR.**  For each channel *c*, predict it from the other channels in its
+cluster via ridge regression, then subtract the prediction::
+
+    y = X - X @ W = X @ (I - W)
+
+The effective weight matrix ``I - W`` is passed to
+:class:`~ezmsg.sigproc.affinetransform.AffineTransformTransformer`, which
+automatically exploits block-diagonal structure when ``channel_clusters``
+are provided.
+
+**Fitting.**  Given data matrix *X* of shape ``(samples, channels)``, the
+sufficient statistic is the channel covariance ``C = X^T X``.  When
+``incremental=True`` (default), *C* is accumulated across
+:meth:`~SelfSupervisedRegressionTransformer.partial_fit` calls.
+
+**Solving.**  Within each cluster the weight matrix *W* is obtained from
+the inverse of the (ridge-regularised) cluster covariance
+``C_inv = (C_cluster + lambda * I)^{-1}`` using the block-inverse identity::
+
+    W[:, c] = -C_inv[:, c] / C_inv[c, c],    diag(W) = 0
+
+This replaces the naive per-channel Cholesky loop with a single matrix
+inverse per cluster, keeping the linear algebra in the source array
+namespace so that GPU-backed arrays benefit from device-side computation.
+"""
+
+from __future__ import annotations
+
+import os
+import typing
+from abc import abstractmethod
+from pathlib import Path
+
+import ezmsg.core as ez
+import numpy as np
+from array_api_compat import get_namespace
+from ezmsg.baseproc import (
+    BaseAdaptiveTransformer,
+    BaseAdaptiveTransformerUnit,
+    processor_state,
+)
+from ezmsg.baseproc.protocols import SettingsType, StateType
+from ezmsg.sigproc.affinetransform import (
+    AffineTransformSettings,
+    AffineTransformTransformer,
+)
+from ezmsg.sigproc.util.array import array_device, xp_create
+from ezmsg.util.messages.axisarray import AxisArray
+
+# ---------------------------------------------------------------------------
+# Base: Self-supervised regression
+# ---------------------------------------------------------------------------
+
+
+class SelfSupervisedRegressionSettings(ez.Settings):
+    """Settings common to all self-supervised regression modes."""
+
+    weights: np.ndarray | str | Path | None = None
+    """Pre-calculated weight matrix *W* or path to a CSV file (``np.loadtxt``
+    compatible).  If provided, the transformer is ready immediately."""
+
+    axis: str | None = None
+    """Channel axis name. ``None`` defaults to the last dimension."""
+
+    channel_clusters: list[list[int]] | None = None
+    """Per-cluster regression.  ``None`` treats all channels as one cluster."""
+
+    ridge_lambda: float = 0.0
+    """Ridge (L2) regularisation parameter."""
+
+    incremental: bool = True
+    """When ``True``, accumulate ``X^T X`` across :meth:`partial_fit` calls.
+    When ``False``, each call replaces the previous statistics."""
+
+
+@processor_state
+class SelfSupervisedRegressionState:
+    cxx: object | None = None  # Array API; namespace matches source data.
+    n_samples: int = 0
+    weights: object | None = None  # Array API; namespace matches cxx.
+
+
+class SelfSupervisedRegressionTransformer(
+    BaseAdaptiveTransformer[SettingsType, AxisArray, AxisArray, StateType],
+    typing.Generic[SettingsType, StateType],
+):
+    """Abstract base for self-supervised regression transformers.
+
+    Subclasses must implement:
+
+    * :meth:`_on_weights_updated` — called whenever the weight matrix *W* is
+      (re)computed, so the subclass can build whatever internal transform it
+      needs (e.g. ``I - W`` for LRR).
+    * :meth:`_process` — the per-message transform step.
+    """
+
+    # -- message hash / state management ------------------------------------
+
+    def _hash_message(self, message: AxisArray) -> int:
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+        return hash((message.key, message.data.shape[axis_idx]))
+
+    def _reset_state(self, message: AxisArray) -> None:
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+        n_channels = message.data.shape[axis_idx]
+
+        self._validate_clusters(n_channels)
+        self._state.cxx = None
+        self._state.n_samples = 0
+        self._state.weights = None
+
+        # If pre-calculated weights are provided, load and go.
+        weights = self.settings.weights
+        if weights is not None:
+            if isinstance(weights, str):
+                weights = Path(os.path.abspath(os.path.expanduser(weights)))
+            if isinstance(weights, Path):
+                weights = np.loadtxt(weights, delimiter=",")
+            weights = np.asarray(weights, dtype=np.float64)
+            self._state.weights = weights
+            self._on_weights_updated()
+
+    # -- cluster validation --------------------------------------------------
+
+    def _validate_clusters(self, n_channels: int) -> None:
+        """Raise if any cluster index is out of range."""
+        clusters = self.settings.channel_clusters
+        if clusters is None:
+            return
+        all_indices = np.concatenate([np.asarray(g) for g in clusters])
+        if np.any((all_indices < 0) | (all_indices >= n_channels)):
+            raise ValueError(f"channel_clusters contains out-of-range indices (valid range: 0..{n_channels - 1})")
+
+    # -- weight solving ------------------------------------------------------
+
+    def _solve_weights(self, cxx):
+        """Solve all per-channel ridge regressions via matrix inverse.
+
+        Uses the block-inverse identity: for target channel *c* with
+        references *r*, ``w_c = -C_inv[r, c] / C_inv[c, c]`` where
+        ``C_inv = (C_cluster + λI)⁻¹``.  This replaces the per-channel
+        Cholesky loop with one matrix inverse per cluster.
+
+        All computation stays in the source array namespace so that
+        GPU-backed arrays benefit from device-side execution.  Cluster
+        results are scattered into the full matrix via a selection-matrix
+        multiply (``S @ W_cluster @ S^T``) to avoid numpy fancy indexing.
+
+        Returns weight matrix *W* in the same namespace as *cxx*, with
+        ``diag(W) == 0``.
+        """
+        xp = get_namespace(cxx)
+        dev = array_device(cxx)
+        n = cxx.shape[0]
+
+        clusters = self.settings.channel_clusters
+        if clusters is None:
+            clusters = [list(range(n))]
+
+        W = xp_create(xp.zeros, (n, n), dtype=cxx.dtype, device=dev)
+        eye_n = xp_create(xp.eye, n, dtype=cxx.dtype, device=dev)
+
+        for cluster in clusters:
+            k = len(cluster)
+            if k <= 1:
+                continue
+
+            idx_xp = xp.asarray(cluster) if dev is None else xp.asarray(cluster, device=dev)
+            eye_k = xp_create(xp.eye, k, dtype=cxx.dtype, device=dev)
+
+            # Extract cluster sub-covariance (stays on device)
+            sub = xp.take(xp.take(cxx, idx_xp, axis=0), idx_xp, axis=1)
+
+            if self.settings.ridge_lambda > 0:
+                sub = sub + self.settings.ridge_lambda * eye_k
+
+            # One inverse per cluster
+            try:
+                sub_inv = xp.linalg.inv(sub)
+            except Exception:
+                sub_inv = xp.linalg.pinv(sub)
+
+            # Diagonal via element-wise product with identity
+            diag_vals = xp.sum(sub_inv * eye_k, axis=0)
+
+            # w_c = -C_inv[:, c] / C_inv[c, c], vectorised over all c
+            W_cluster = -(sub_inv / xp.reshape(diag_vals, (1, k)))
+
+            # Zero the diagonal
+            W_cluster = W_cluster * (1.0 - eye_k)
+
+            # Scatter into full W
+            if k == n:
+                W = W + W_cluster
+            else:
+                # Selection matrix: columns of eye(n) at cluster indices
+                S = xp.take(eye_n, idx_xp, axis=1)  # (n, k)
+                W = W + xp.matmul(S, xp.matmul(W_cluster, xp.permute_dims(S, (1, 0))))
+
+        return W
+
+    # -- partial_fit (self-supervised, accepts AxisArray) --------------------
+
+    def partial_fit(self, message: AxisArray) -> None:  # type: ignore[override]
+        xp = get_namespace(message.data)
+
+        if xp.any(xp.isnan(message.data)):
+            return
+
+        # Hash check / state reset
+        msg_hash = self._hash_message(message)
+        if self._hash != msg_hash:
+            self._reset_state(message)
+            self._hash = msg_hash
+
+        axis = self.settings.axis or message.dims[-1]
+        axis_idx = message.get_axis_idx(axis)
+        data = message.data
+
+        # Move channel axis to last, flatten to 2-D
+        if axis_idx != data.ndim - 1:
+            perm = list(range(data.ndim))
+            perm.append(perm.pop(axis_idx))
+            data = xp.permute_dims(data, perm)
+
+        n_channels = data.shape[-1]
+        X = xp.reshape(data, (-1, n_channels))
+
+        # Covariance stays in the source namespace for accumulation.
+        cxx_new = xp.matmul(xp.permute_dims(X, (1, 0)), X)
+
+        if self.settings.incremental and self._state.cxx is not None:
+            self._state.cxx = self._state.cxx + cxx_new
+        else:
+            self._state.cxx = cxx_new
+        self._state.n_samples += int(X.shape[0])
+
+        self._state.weights = self._solve_weights(self._state.cxx)
+        self._on_weights_updated()
+
+    # -- convenience APIs ----------------------------------------------------
+
+    def fit(self, X: np.ndarray) -> None:
+        """Batch fit from a raw numpy array (samples x channels)."""
+        n_channels = X.shape[-1]
+        self._validate_clusters(n_channels)
+        X = np.asarray(X, dtype=np.float64).reshape(-1, n_channels)
+        self._state.cxx = X.T @ X
+        self._state.n_samples = X.shape[0]
+        self._state.weights = self._solve_weights(self._state.cxx)
+        self._on_weights_updated()
+
+    def fit_transform(self, message: AxisArray) -> AxisArray:
+        """Convenience: ``partial_fit`` then ``_process``."""
+        self.partial_fit(message)
+        return self._process(message)
+
+    # -- abstract hooks for subclasses ---------------------------------------
+
+    @abstractmethod
+    def _on_weights_updated(self) -> None:
+        """Called after ``self._state.weights`` has been set/updated.
+
+        Subclasses should build or refresh whatever internal transform
+        object they need for :meth:`_process`.
+        """
+        ...
+
+    @abstractmethod
+    def _process(self, message: AxisArray) -> AxisArray: ...
+
+
+# ---------------------------------------------------------------------------
+# Concrete: Linear Regression Rereferencing (LRR)
+# ---------------------------------------------------------------------------
+
+
+class LRRSettings(SelfSupervisedRegressionSettings):
+    """Settings for :class:`LRRTransformer`."""
+
+    min_cluster_size: int = 32
+    """Passed to :class:`AffineTransformTransformer` for the block-diagonal
+    merge threshold."""
+
+
+@processor_state
+class LRRState(SelfSupervisedRegressionState):
+    affine: AffineTransformTransformer | None = None
+
+
+class LRRTransformer(
+    SelfSupervisedRegressionTransformer[LRRSettings, LRRState],
+):
+    """Adaptive LRR transformer.
+
+    ``partial_fit`` accepts a plain :class:`AxisArray` (self-supervised),
+    and the transform step is delegated to an internal :class:`AffineTransformTransformer`.
+    """
+
+    # -- state management (clear own state, then delegate to base) ----------
+
+    def _reset_state(self, message: AxisArray) -> None:
+        self._state.affine = None
+        super()._reset_state(message)
+
+    # -- weights → affine transform -----------------------------------------
+
+    def _on_weights_updated(self) -> None:
+        xp = get_namespace(self._state.weights)
+        dev = array_device(self._state.weights)
+        n = self._state.weights.shape[0]
+        effective = xp_create(xp.eye, n, dtype=self._state.weights.dtype, device=dev) - self._state.weights
+
+        # Prefer in-place weight update when the affine transformer supports
+        # it (avoids a full _reset_state round-trip on every partial_fit).
+        if self._state.affine is not None:
+            self._state.affine.set_weights(effective)
+        else:
+            self._state.affine = AffineTransformTransformer(
+                AffineTransformSettings(
+                    weights=effective,
+                    axis=self.settings.axis,
+                    channel_clusters=self.settings.channel_clusters,
+                    min_cluster_size=self.settings.min_cluster_size,
+                )
+            )
+
+    # -- transform -----------------------------------------------------------
+
+    def _process(self, message: AxisArray) -> AxisArray:
+        if self._state.affine is None:
+            raise RuntimeError(
+                "LRRTransformer has not been fitted. Call partial_fit() or provide pre-calculated weights."
+            )
+        return self._state.affine(message)
+
+
+class LRRUnit(
+    BaseAdaptiveTransformerUnit[
+        LRRSettings,
+        AxisArray,
+        AxisArray,
+        LRRTransformer,
+    ],
+):
+    """ezmsg Unit wrapping :class:`LRRTransformer`.
+
+    Follows the :class:`BaseAdaptiveDecompUnit` pattern — accepts
+    :class:`AxisArray` (not :class:`SampleMessage`) for self-supervised
+    training via ``INPUT_SAMPLE``.
+    """
+
+    SETTINGS = LRRSettings
+
+    INPUT_SAMPLE = ez.InputStream(AxisArray)
+
+    @ez.subscriber(INPUT_SAMPLE)
+    async def on_sample(self, msg: AxisArray) -> None:
+        await self.processor.apartial_fit(msg)

ezmsg/learn/process/torch.py

@@ -4,15 +4,14 @@ import typing
 import ezmsg.core as ez
 import numpy as np
 import torch
-from ezmsg.sigproc.base import (
+from ezmsg.baseproc import (
     BaseAdaptiveTransformer,
     BaseAdaptiveTransformerUnit,
     BaseStatefulTransformer,
     BaseTransformerUnit,
     processor_state,
 )
-from ezmsg.sigproc.sampler import SampleMessage
-from ezmsg.sigproc.util.profile import profile_subpub
+from ezmsg.baseproc.util.profile import profile_subpub
 from ezmsg.util.messages.axisarray import AxisArray
 from ezmsg.util.messages.util import replace
 
@@ -113,9 +112,7 @@ class TorchProcessorMixin:
         module = importlib.import_module(module_path)
         return getattr(module, class_name)
 
-    def _infer_output_sizes(
-        self: P, model: torch.nn.Module, n_input: int
-    ) -> dict[str, int]:
+    def _infer_output_sizes(self: P, model: torch.nn.Module, n_input: int) -> dict[str, int]:
         """Simple inference to get output channel size. Override if needed."""
         dummy_input = torch.zeros(1, 1, n_input, device=self._state.device)
         with torch.no_grad():
@@ -133,9 +130,7 @@ class TorchProcessorMixin:
             weight_decay=self.settings.weight_decay,
         )
         self._state.scheduler = (
-            torch.optim.lr_scheduler.ExponentialLR(
-                self._state.optimizer, gamma=self.settings.scheduler_gamma
-            )
+            torch.optim.lr_scheduler.ExponentialLR(self._state.optimizer, gamma=self.settings.scheduler_gamma)
             if self.settings.scheduler_gamma > 0.0
             else None
         )
@@ -191,9 +186,7 @@ class TorchProcessorMixin:
             output_messages = [
                 replace(
                     message,
-                    data=value.cpu().numpy().squeeze(0)
-                    if added_batch_dim
-                    else value.cpu().numpy(),
+                    data=value.cpu().numpy().squeeze(0) if added_batch_dim else value.cpu().numpy(),
                     axes={
                         **message.axes,
                         "ch": self._state.chan_ax[key],
@@ -207,9 +200,7 @@ class TorchProcessorMixin:
         return [
             replace(
                 message,
-                data=output.cpu().numpy().squeeze(0)
-                if added_batch_dim
-                else output.cpu().numpy(),
+                data=output.cpu().numpy().squeeze(0) if added_batch_dim else output.cpu().numpy(),
                 axes={
                     **message.axes,
                     "ch": self._state.chan_ax["output"],
@@ -229,11 +220,7 @@ class TorchProcessorMixin:
         else:
             model_kwargs["input_size"] = n_input
 
-        device = (
-            "cuda"
-            if torch.cuda.is_available()
-            else ("mps" if torch.mps.is_available() else "cpu")
-        )
+        device = "cuda" if torch.cuda.is_available() else ("mps" if torch.mps.is_available() else "cpu")
         device = self.settings.device or device
         self._state.device = torch.device(device)
 
@@ -260,9 +247,7 @@ class TorchProcessorMixin:
 
 
 class TorchSimpleProcessor(
-    BaseStatefulTransformer[
-        TorchSimpleSettings, AxisArray, AxisArray, TorchSimpleState
-    ],
+    BaseStatefulTransformer[TorchSimpleSettings, AxisArray, AxisArray, TorchSimpleState],
     TorchProcessorMixin,
     ModelInitMixin,
 ):
@@ -308,13 +293,13 @@ class TorchModelProcessor(
     def _process(self, message: AxisArray) -> list[AxisArray]:
         return self._common_process(message)
 
-    def partial_fit(self, message: SampleMessage) -> None:
+    def partial_fit(self, message: AxisArray) -> None:
         self._state.model.train()
 
-        X = self._to_tensor(message.sample.data)
+        X = self._to_tensor(message.data)
         X, batched = self._ensure_batched(X)
 
-        y_targ = message.trigger.value
+        y_targ = message.attrs["trigger"].value
         if not isinstance(y_targ, dict):
             y_targ = {"output": y_targ}
         y_targ = {key: self._to_tensor(value) for key, value in y_targ.items()}
@@ -339,9 +324,7 @@ class TorchModelProcessor(
             for key in y_targ.keys():
                 loss_fn = loss_fns.get(key)
                 if loss_fn is None:
-                    raise ValueError(
-                        f"Loss function for key '{key}' is not defined in settings."
-                    )
+                    raise ValueError(f"Loss function for key '{key}' is not defined in settings.")
                 if isinstance(loss_fn, torch.nn.CrossEntropyLoss):
                     loss = loss_fn(y_pred[key].permute(0, 2, 1), y_targ[key].long())
                 else: