diff --git a/pyproject.toml b/pyproject.toml index a3779fa..d07c19d 100644 --- a/pyproject.toml +++ b/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", diff --git a/src/ezmsg/learn/collection/sample_adapt_regressor.py b/src/ezmsg/learn/collection/sample_adapt_regressor.py index b6eb4bf..3db795b 100644 --- a/src/ezmsg/learn/collection/sample_adapt_regressor.py +++ b/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`` diff --git a/src/ezmsg/learn/dim_reduce/adaptive_decomp.py b/src/ezmsg/learn/dim_reduce/adaptive_decomp.py index 7c97787..74edfb8 100644 --- a/src/ezmsg/learn/dim_reduce/adaptive_decomp.py +++ b/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,) diff --git a/src/ezmsg/learn/process/adaptive_linear_regressor.py b/src/ezmsg/learn/process/adaptive_linear_regressor.py index 3605458..3607d6b 100644 --- a/src/ezmsg/learn/process/adaptive_linear_regressor.py +++ b/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. diff --git a/src/ezmsg/learn/process/flatten.py b/src/ezmsg/learn/process/flatten.py index 02cbdc8..128d0ed 100644 --- a/src/ezmsg/learn/process/flatten.py +++ b/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] diff --git a/src/ezmsg/learn/process/mlp_old.py b/src/ezmsg/learn/process/mlp_old.py index 78d1360..fc0e012 100644 --- a/src/ezmsg/learn/process/mlp_old.py +++ b/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: diff --git a/src/ezmsg/learn/process/sgd.py b/src/ezmsg/learn/process/sgd.py index 44d67c5..3d5ab67 100644 --- a/src/ezmsg/learn/process/sgd.py +++ b/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() diff --git a/src/ezmsg/learn/process/sklearn.py b/src/ezmsg/learn/process/sklearn.py index 9d72270..dac9ccb 100644 --- a/src/ezmsg/learn/process/sklearn.py +++ b/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, diff --git a/src/ezmsg/learn/process/slda.py b/src/ezmsg/learn/process/slda.py index 8db090e..fb346f6 100644 --- a/src/ezmsg/learn/process/slda.py +++ b/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, diff --git a/src/ezmsg/learn/process/ssr.py b/src/ezmsg/learn/process/ssr.py index b696c85..b51d9f9 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/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) diff --git a/src/ezmsg/learn/process/torch.py b/src/ezmsg/learn/process/torch.py index 0efeb3f..c4fdd86 100644 --- a/src/ezmsg/learn/process/torch.py +++ b/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() } diff --git a/src/ezmsg/learn/util.py b/src/ezmsg/learn/util.py index f0f9b53..5a113ed 100644 --- a/src/ezmsg/learn/util.py +++ b/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" diff --git a/tests/unit/test_axis_fingerprint_priming.py b/tests/unit/test_axis_fingerprint_priming.py new file mode 100644 index 0000000..ad4a5de --- /dev/null +++ b/tests/unit/test_axis_fingerprint_priming.py @@ -0,0 +1,148 @@ +"""Coordinate axes built here are handed downstream ready to use. + +``CoordinateAxis.fingerprint`` is what every stateful consumer keys its cached +state on. It is computed on first access and cached on the instance, and the +cache pickles with the axis, so whoever touches it first pays and everyone +after gets it free. + +In-process that first toucher is usually the next stateful node, which primes +the axis as a side effect of hashing it. The gap is the process boundary: +unpickling builds a *new* axis object per message, so an axis that left its +producing process cold is re-checksummed by the first consumer in every +receiving process, on every message. Priming at construction closes that, and +these tests are what would notice it stopping. +""" + +import pickle + +import numpy as np +import pytest +from ezmsg.util.messages.axisarray import AxisArray, CoordinateAxis + +from ezmsg.learn.util import with_fingerprint + + +def signal(labels, n_time=32, fs=100.0, key="dev"): + return AxisArray( + np.random.default_rng(0).standard_normal((n_time, len(labels))), + dims=["time", "ch"], + axes={ + "time": AxisArray.TimeAxis(fs=fs), + "ch": CoordinateAxis(data=np.array(labels), dims=["ch"]), + }, + key=key, + chunk_dim="time", + ) + + +def created_axes(source: AxisArray, result: AxisArray) -> dict: + """Coordinate axes on *result* that are not objects *source* handed in. + + Identity, not equality: an axis that merely passed through was primed by + whoever hashed it, which would mask a producer that primes nothing. + """ + incoming = {id(a) for a in source.axes.values()} + return {d: a for d, a in result.axes.items() if isinstance(a, CoordinateAxis) and id(a) not in incoming} + + +class TestTheHelper: + def test_it_returns_the_same_axis(self): + axis = CoordinateAxis(data=np.array(["a", "b"]), dims=["ch"]) + assert with_fingerprint(axis) is axis + + def test_it_is_idempotent(self): + axis = with_fingerprint(CoordinateAxis(data=np.array(["a", "b"]), dims=["ch"])) + first = axis.__dict__["_fingerprint"] + assert with_fingerprint(axis).__dict__["_fingerprint"] is first + + @pytest.mark.parametrize("dtype", ["U8", "f8", "i4"]) + def test_priming_survives_the_transport(self, dtype): + """The whole point: the far side gets the answer without recomputing.""" + axis = with_fingerprint(CoordinateAxis(data=np.arange(8).astype(dtype), dims=["ch"])) + landed = pickle.loads(pickle.dumps(axis)) + assert "_fingerprint" in landed.__dict__ + assert landed.__dict__["_fingerprint"] == axis.fingerprint + + +class TestCreatedAxesArePrimed: + """Messages here carry no ``chunk_dim``: released ezmsg-sigproc does not set + it, so that is what these transformers actually receive today. It is why + ``FlattenTransformer.STREAMING_DIMS`` names ``win`` -- the base class's + ``("time",)`` fallback would exclude the lag dimension, which is the one + thing the lag axis is sized by.""" + + def test_the_lag_axis_flatten_builds(self): + """The lag axis is what *this* package builds. The merged output axis is + built by the inner ezmsg-sigproc transformer and primed there, not here.""" + from ezmsg.learn.process.flatten import FlattenSettings, FlattenTransformer + + proc = FlattenTransformer(FlattenSettings(preserve_axis="win", sample_axis="time", feature_axis="ch")) + proc( + AxisArray( + np.arange(24).reshape(2, 3, 4).astype(float), + dims=["win", "time", "ch"], + axes={ + "win": AxisArray.TimeAxis(fs=50.0), + "time": AxisArray.TimeAxis(fs=50.0), + "ch": CoordinateAxis(data=np.array(["a", "b", "c", "d"]), dims=["ch"]), + }, + key="dev", + ) + ) + lag_axis = proc._state.lag_axis + assert lag_axis is not None, "expected the lag case to be detected" + assert "_fingerprint" in lag_axis.__dict__ + + def test_the_component_axis_incremental_pca_builds(self): + from ezmsg.learn.dim_reduce.adaptive_decomp import ( + IncrementalPCASettings, + IncrementalPCATransformer, + ) + + proc = IncrementalPCATransformer(IncrementalPCASettings(n_components=2)) + msg = signal(["c0", "c1", "c2"], n_time=64) + proc.partial_fit(msg) + out = proc(msg) + assert out is not None + cold = [d for d, a in out.axes.items() if isinstance(a, CoordinateAxis) and "_fingerprint" not in a.__dict__] + assert not cold, f"handed downstream cold: {cold}" + + +class TestTheFlattenFallbackIsRight: + """``STREAMING_DIMS`` decides which dimension is excluded when the producer + is silent, and getting it wrong is not a small error.""" + + @staticmethod + def _msg(n_win, n_lag): + return AxisArray( + np.zeros((n_win, n_lag, 4)), + dims=["win", "time", "ch"], + axes={ + "win": AxisArray.TimeAxis(fs=50.0), + "time": AxisArray.TimeAxis(fs=50.0), + "ch": CoordinateAxis(data=np.array(["a", "b", "c", "d"]), dims=["ch"]), + }, + key="dev", + ) + + @staticmethod + def _proc(): + from ezmsg.learn.process.flatten import FlattenSettings, FlattenTransformer + + return FlattenTransformer(FlattenSettings(preserve_axis="win", sample_axis="time", feature_axis="ch")) + + def test_a_window_length_change_rebuilds(self): + """The lag axis is sized by it.""" + proc = self._proc() + proc(self._msg(n_win=2, n_lag=3)) + inner = proc._state.inner + proc(self._msg(n_win=2, n_lag=5)) + assert proc._state.inner is not inner + + def test_a_window_count_change_does_not(self): + """That is just how many windows arrived.""" + proc = self._proc() + proc(self._msg(n_win=2, n_lag=3)) + inner = proc._state.inner + proc(self._msg(n_win=7, n_lag=3)) + assert proc._state.inner is inner diff --git a/tests/unit/test_ssr.py b/tests/unit/test_ssr.py index 07ab04a..da45d5d 100644 --- a/tests/unit/test_ssr.py +++ b/tests/unit/test_ssr.py @@ -228,12 +228,17 @@ def test_missing_field_falls_back_to_block_size(self): np.testing.assert_array_equal(proc_field.state.weights, proc_block.state.weights) - def test_bank_field_value_change_is_not_detected(self): - """Intentional concession (mirrors the ezmsg-sigproc CAR fix): a live bank - remap at fixed key + channel count is NOT re-derived. ``_hash_message`` - folds only an O(1) "bank field present" boolean, not the field's bytes, so - the per-message hash does not scale with channel count. A genuine remap on - real hardware arrives with a new key or channel count (escape hatch below).""" + def test_bank_field_value_change_is_detected(self): + """A live bank remap at fixed key and channel count is now re-derived. + + This asserted the opposite until ezmsg-baseproc 1.12.0. The old hash + folded an O(1) "bank field present" boolean rather than the field's + bytes, so a remap that kept the key and the channel count looked + identical and the cached groups were reused -- silently rereferencing + each channel against the wrong bank. The concession existed because + digesting the field per message was thought to scale with channel + count; ``CoordinateAxis.fingerprint`` removes that cost by computing + the digest once per axis object rather than once per consumer.""" rng = np.random.default_rng(11) X = _random_data(n_ch=4, rng=rng) proc = LRRTransformer(LRRSettings(axis="ch", channel_groups="bank")) @@ -243,15 +248,14 @@ def test_bank_field_value_change_is_not_detected(self): assert [g.tolist() for g in proc.state.resolved_groups] == [[0, 1], [2, 3]] np.testing.assert_array_equal(proc.state.weights[np.ix_([0, 1], [2, 3])], 0.0) - # Same key + channel count, different banks -> hash unchanged, so the - # cached groups are (deliberately) NOT re-derived. + # Same key, same channel count, different banks -> re-derived. proc.partial_fit(_banked_axisarray(X, ["A", "B", "A", "B"], key="x")) - assert [g.tolist() for g in proc.state.resolved_groups] == [[0, 1], [2, 3]] - - # Escape hatch: a new key (as a real remap would carry) forces re-derivation. - proc.partial_fit(_banked_axisarray(X, ["A", "B", "A", "B"], key="y")) assert [g.tolist() for g in proc.state.resolved_groups] == [[0, 2], [1, 3]] + # A new key still forces re-derivation, as it always did. + proc.partial_fit(_banked_axisarray(X, ["B", "B", "A", "A"], key="y")) + assert [g.tolist() for g in proc.state.resolved_groups] == [[0, 1], [2, 3]] + class TestIncrementalAccumulates: def test_incremental_accumulates(self):