Let the base class decide when these transformers reset

pyproject.toml

@@ -9,8 +9,12 @@ license = "MIT"
 requires-python = ">=3.10"
 dynamic = ["version"]
 dependencies = [
-    "ezmsg>=3.9.0",
-    "ezmsg-baseproc>=1.7.0",
+    # 3.10.0b2 for CoordinateAxis.fingerprint and AxisArray.chunk_dim, which
+    # ezmsg-baseproc's default state hash reads. Pinned directly rather than
+    # left to the transitive requirement: uv only enables pre-releases for a
+    # package named with a pre-release marker in *this* file.
+    "ezmsg>=3.10.0b2",
+    "ezmsg-baseproc>=1.12.0",
     "ezmsg-sigproc>=3.0.0",
     "numpy",
     "scipy",

src/ezmsg/learn/collection/sample_adapt_regressor.py

@@ -31,6 +31,8 @@
 from ezmsg.learn.process.seqseqsampler import SeqSeqSamplerSettings, SeqSeqSamplerUnit
 from ezmsg.learn.util import AdaptiveLinearRegressor
 
+from ..util import with_fingerprint
+
 #: Default torch model class used when ``model_type == "mlp"``.
 DEFAULT_TORCH_MODEL_CLASS = "ezmsg.learn.model.mlp.MLP"
 
@@ -87,7 +89,9 @@ class DecodeOutputAdapterProcessor(
 
     def _reset_state(self, message: AxisArray) -> None:
         if self.settings.output_labels is not None:
-            self.state.ch_axis = AxisArray.CoordinateAxis(data=np.asarray(self.settings.output_labels), dims=["ch"])
+            self.state.ch_axis = with_fingerprint(
+                AxisArray.CoordinateAxis(data=np.asarray(self.settings.output_labels), dims=["ch"])
+            )
 
     def _process(self, message: AxisArray) -> AxisArray | None:
         data = np.asarray(message.data, dtype=float)
@@ -100,7 +104,9 @@ def _process(self, message: AxisArray) -> AxisArray | None:
             ch_axis = self.state.ch_axis
         else:
             data = data.reshape((data.shape[0], -1)) if data.ndim > 1 else data.reshape((1, -1))
-            ch_axis = AxisArray.CoordinateAxis(data=np.asarray([f"ch{i}" for i in range(data.shape[-1])]), dims=["ch"])
+            ch_axis = with_fingerprint(
+                AxisArray.CoordinateAxis(data=np.asarray([f"ch{i}" for i in range(data.shape[-1])]), dims=["ch"])
+            )
 
         # The decoder engines carry a ``time`` axis through (kalman keeps the
         # input's; the torch path inherits the windower's renamed ``win``->``time``

src/ezmsg/learn/dim_reduce/adaptive_decomp.py

@@ -21,6 +21,7 @@
 from ezmsg.util.messages.axisarray import AxisArray, replace
 
 from .._optional import missing_extra
+from ..util import with_fingerprint
 
 try:
     from sklearn.decomposition import IncrementalPCA, MiniBatchNMF
@@ -97,16 +98,6 @@ def _calculate_axis_groups(self, message: AxisArray):
             ]
         self._state.axis_groups = iter_axis, targ_axes, off_targ_axes
 
-    def _hash_message(self, message: AxisArray) -> int:
-        iter_axis = (
-            self.settings.axis[1:]
-            if self.settings.axis.startswith("!")
-            else ("win" if "win" in message.dims else "time")
-        )
-        ax_idx = message.get_axis_idx(iter_axis)
-        sample_shape = message.data.shape[:ax_idx] + message.data.shape[ax_idx + 1 :]
-        return hash((sample_shape, message.key))
-
     def _reset_state(self, message: AxisArray) -> None:
         """Reset state"""
         self._calculate_axis_groups(message)
@@ -123,10 +114,12 @@ def _reset_state(self, message: AxisArray) -> None:
         else:
             targ_ax_name = "components"
         out_dims += [targ_ax_name]
-        out_axes[targ_ax_name] = AxisArray.CoordinateAxis(
-            data=np.arange(self.settings.n_components).astype(str),
-            dims=[targ_ax_name],
-            unit="component",
+        out_axes[targ_ax_name] = with_fingerprint(
+            AxisArray.CoordinateAxis(
+                data=np.arange(self.settings.n_components).astype(str),
+                dims=[targ_ax_name],
+                unit="component",
+            )
         )
         out_shape = [message.data.shape[message.get_axis_idx(_)] for _ in off_targ_axes]
         out_shape = (0,) + tuple(out_shape) + (self.settings.n_components,)

src/ezmsg/learn/process/adaptive_linear_regressor.py

@@ -22,7 +22,7 @@
 from ezmsg.util.messages.axisarray import AxisArray, replace
 
 from .._optional import missing_extra
-from ..util import AdaptiveLinearRegressor, RegressorType, get_regressor
+from ..util import AdaptiveLinearRegressor, RegressorType, get_regressor, with_fingerprint
 
 try:
     import pandas as pd
@@ -86,7 +86,7 @@ def _prediction_template_from_signal(message: AxisArray, output_labels: list[typ
         dims=["time", "ch"],
         axes={
             "time": replace(message.axes["time"], offset=message.axes["time"].offset),
-            "ch": AxisArray.CoordinateAxis(data=np.asarray(output_labels), dims=["ch"]),
+            "ch": with_fingerprint(AxisArray.CoordinateAxis(data=np.asarray(output_labels), dims=["ch"])),
         },
         key=message.key + "_pred",
     )
@@ -148,8 +148,12 @@ def __init__(self, *args, **kwargs):
             self.state.model = self._regressor_klass(**self.settings.model_kwargs)
 
     def _hash_message(self, message: AxisArray) -> int:
-        # So far, nothing to reset so hash can be constant.
-        return -1
+        # Nothing to reset -- `.model` is built in __init__ and `.template` is
+        # updated in partial_fit -- so a constant is both correct and the
+        # cheapest possible hash. Zero rather than -1: -1 is the sentinel
+        # `_hash` starts at and that `_request_reset()` writes, so returning it
+        # made an explicitly requested reset compare equal and be swallowed.
+        return 0
 
     def _reset_state(self, message: AxisArray) -> None:
         # So far, there is nothing to reset.

src/ezmsg/learn/process/flatten.py

@@ -35,6 +35,8 @@
 )
 from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis, replace
 
+from ..util import with_fingerprint
+
 
 class FlattenSettings(ez.Settings):
     """Settings for the learn-side :obj:`Flatten`.
@@ -104,7 +106,7 @@ def _build_lag_axis(sample_dim: str, sample_size: int) -> CoordinateAxis:
     data = np.empty(sample_size, dtype=dtype)
     data["lag"] = lags
     data["label"] = label_strs
-    return CoordinateAxis(data=data, dims=[sample_dim])
+    return with_fingerprint(CoordinateAxis(data=data, dims=[sample_dim]))
 
 
 class FlattenTransformer(BaseStatefulTransformer[FlattenSettings, AxisArray, AxisArray, _LagFlattenState]):
@@ -118,8 +120,19 @@ class FlattenTransformer(BaseStatefulTransformer[FlattenSettings, AxisArray, Axi
     sigproc-composed ``"label"`` (``"t-2/c0"`` style).
     """
 
-    def _hash_message(self, message: AxisArray) -> int:
-        return hash((tuple(message.dims), tuple(message.data.shape)))
+    STREAMING_DIMS = ("win",)
+    """Fallback chunk dimension when the producer does not declare one.
+
+    The base class defaults to ``("time",)``, which is exactly wrong here: the
+    canonical input is ``(win, time, ch[, feature])``, where ``win`` is what
+    grows per message and ``time`` is the *lag* dimension inside each window.
+    The lag count sizes the lag axis built below, so excluding ``time`` would
+    stop this noticing a window-length change, while including ``win`` would
+    rebuild the inner transformer every time the window count jittered.
+
+    Consulted only when :attr:`AxisArray.chunk_dim` is absent; a producer that
+    declares it -- ezmsg-sigproc's ``Window`` does -- overrides this.
+    """
 
     def _reset_state(self, message: AxisArray) -> None:
         preserve_axis = self.settings.preserve_axis or message.dims[0]

src/ezmsg/learn/process/mlp_old.py

@@ -12,6 +12,7 @@
 
 from .._optional import missing_extra
 from ..model.mlp_old import MLP
+from ..util import with_fingerprint
 
 try:
     import torch
@@ -65,12 +66,6 @@ class MLPState:
 
 
 class MLPProcessor(BaseAdaptiveTransformer[MLPSettings, AxisArray, AxisArray, MLPState]):
-    def _hash_message(self, message: AxisArray) -> int:
-        hash_items = (message.key,)
-        if "ch" in message.dims:
-            hash_items += (message.data.shape[message.get_axis_idx("ch")],)
-        return hash(hash_items)
-
     def _reset_state(self, message: AxisArray) -> None:
         # Create the model
         self._state.model = MLP(
@@ -118,8 +113,8 @@ def _reset_state(self, message: AxisArray) -> None:
 
         # Create the output channel axis for reuse in each output.
         n_output_channels = self.settings.hidden_channels[-1]
-        self._state.chan_ax = AxisArray.CoordinateAxis(
-            data=np.array([f"ch{_}" for _ in range(n_output_channels)]), dims=["ch"]
+        self._state.chan_ax = with_fingerprint(
+            AxisArray.CoordinateAxis(data=np.array([f"ch{_}" for _ in range(n_output_channels)]), dims=["ch"])
         )
 
     def save_checkpoint(self, path: str) -> None:

src/ezmsg/learn/process/sgd.py

@@ -64,6 +64,23 @@ def _refreshed_model(self):
             )
         return model
 
+    def _hash_message(self, message: AxisArray) -> int:
+        """Constant: inference must never rebuild the model.
+
+        The model's lifecycle belongs to `partial_fit`, which sets `_hash` to 0
+        itself once it has trained. Training samples arrive as
+        `(time, ch, freq)` and inference windows as `(win, time, ch, freq)`, so
+        any hash that reads the layout differs between the two and makes every
+        alternation throw the fitted model away -- which is what
+        `_refreshed_model()` below does.
+
+        This was previously inherited from ezmsg-baseproc's old default, which
+        returned a constant for everything. Now that the default keys on the
+        message layout, the assumption has to be stated here rather than
+        depended upon.
+        """
+        return 0
+
     def _reset_state(self, message: AxisArray) -> None:
         self._state.model = self._refreshed_model()
 

src/ezmsg/learn/process/sklearn.py

@@ -14,6 +14,7 @@
 from ezmsg.util.messages.util import replace
 
 from .._optional import missing_extra
+from ..util import with_fingerprint
 
 try:
     import pandas as pd
@@ -235,7 +236,7 @@ def _process(self, message: AxisArray) -> AxisArray | None:
                 chan_labels = np.asarray(self._state.model.classes_)
             else:
                 chan_labels = np.arange(output_shape[1])
-            self._state.chan_ax = AxisArray.CoordinateAxis(data=chan_labels, dims=["ch"])
+            self._state.chan_ax = with_fingerprint(AxisArray.CoordinateAxis(data=chan_labels, dims=["ch"]))
 
         return replace(
             message,

src/ezmsg/learn/process/slda.py

@@ -21,7 +21,7 @@
 from ezmsg.util.messages.util import replace
 
 from .._optional import missing_extra
-from ..util import ClassifierMessage
+from ..util import ClassifierMessage, with_fingerprint
 
 try:
     from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
@@ -80,7 +80,7 @@ def _reset_state(self, message: AxisArray) -> None:
             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": with_fingerprint(AxisArray.CoordinateAxis(data=np.array(out_labels), dims=["classes"])),
             },
             labels=out_labels,
             key=message.key,

src/ezmsg/learn/process/ssr.py

@@ -63,7 +63,6 @@
 from ezmsg.sigproc.util.channels import (
     ChannelGroupSpec,
     group_spec_fields,
-    group_spec_fingerprint,
     resolve_channel_groups,
     validate_channel_groups,
 )
@@ -149,18 +148,6 @@ class SelfSupervisedRegressionTransformer(
 
     # -- 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)
-        # group_spec_fingerprint contributes an O(1) "can this spec resolve?"
-        # boolean rather than the field's bytes, so the per-message hash does not
-        # grow with channel count. See its docstring for what that deliberately
-        # does not detect. Mirrors the ezmsg-sigproc transformers' hash.
-        return hash(
-            (message.key, message.data.shape[axis_idx])
-            + group_spec_fingerprint(message, axis, self.settings.channel_groups)
-        )
-
     def _reset_state(self, message: AxisArray) -> None:
         axis = self.settings.axis or message.dims[-1]
         axis_idx = message.get_axis_idx(axis)

src/ezmsg/learn/process/torch.py

@@ -15,6 +15,7 @@
 from ezmsg.util.messages.util import replace
 
 from .._optional import missing_extra
+from ..util import with_fingerprint
 from .base import ModelInitMixin
 
 try:
@@ -243,9 +244,11 @@ def _common_reset_state(self: P, message: AxisArray, model_kwargs: dict) -> None
 
         output_sizes = self._infer_output_sizes(self._state.model, n_input)
         self._state.chan_ax = {
-            head: AxisArray.CoordinateAxis(
-                data=np.array([f"{head}_ch{_}" for _ in range(size)]),
-                dims=["ch"],
+            head: with_fingerprint(
+                AxisArray.CoordinateAxis(
+                    data=np.array([f"{head}_ch{_}" for _ in range(size)]),
+                    dims=["ch"],
+                )
             )
             for head, size in output_sizes.items()
         }

src/ezmsg/learn/util.py

@@ -9,6 +9,28 @@
 # from sklearn.neural_network import MLPClassifier
 
 
+def with_fingerprint(axis: AxisArray.CoordinateAxis) -> AxisArray.CoordinateAxis:
+    """Compute *axis*'s fingerprint now, and return the axis.
+
+    Every stateful consumer reads the fingerprint of the coordinate axes that
+    describe a stream's configuration, and the value is cached on the instance
+    and pickled with it. Computing it where the axis is built pays the checksum
+    once, for everybody:
+
+    * In this process the axis object is reused for the life of the stream, so
+      one call covers every message and every consumer downstream of it.
+    * Across a process boundary it is better than that. Unpickling hands out a
+      *new* axis object per message, so a cold axis is re-checksummed by the
+      first consumer in every receiving process, on every message.
+
+    Apply it to axes that describe the stream -- channel labels, class labels,
+    lag labels -- not to per-message coordinates along the chunk dimension,
+    whose fingerprint no consumer reads.
+    """
+    axis.fingerprint
+    return axis
+
+
 class RegressorType(str, Enum):
     ADAPTIVE = "adaptive"
     STATIC = "static"