Keep MLX dense event-rate state on device

src/ezmsg/event/kernel_activation.py

@@ -6,13 +6,13 @@
 uses a state-based approach that is O(n_events + n_bins) instead of
 O(n_samples).
 
-Input may be either ``sparse.COO`` (the typical output of
-:class:`ezmsg.event.peak.ThresholdCrossingTransformer` in default mode) or a
-dense array (from the same transformer with ``output_format=DENSE``). When the
-input is dense and the configuration is COUNT + SUM (the rate-computation
-case), the binning runs on the input's array namespace and stays on device
-(e.g., MLX, CuPy). Other configurations with dense input fall back to
-event-extraction and use the same code path as sparse input.
+Input may be either ``sparse.COO`` (the default output of
+:class:`ezmsg.event.peak.ThresholdCrossingTransformer`) or a dense array from the
+same transformer with ``output_format=DENSE``. When the input is dense and the
+configuration is COUNT + SUM (the rate-computation case), the binning runs on the
+input's array namespace and stays on device (e.g., MLX, CuPy). Other configurations
+with dense input fall back to event extraction and use the same code path as sparse
+input.
 """
 
 from enum import Enum
@@ -97,6 +97,9 @@ class BinnedKernelActivationState:
     # Current activation level per channel (for exponential/alpha)
     activation: npt.NDArray[np.float64] | None = None
 
+    dense_carry: object | None = None
+    """Partial-bin COUNT+SUM state in the dense input's array namespace."""
+
     # For alpha kernel: auxiliary state variable
     alpha_aux: npt.NDArray[np.float64] | None = None
 
@@ -145,13 +148,15 @@ def _hash_message(self, message: AxisArray) -> int:
         if "time" not in message.axes or not hasattr(message.axes["time"], "gain"):
             raise ValueError("Could not determine sample rate from input message")
         # str(dtype) works for numpy ('bool', 'float32', ...) and mlx (which doesn't expose dtype.kind).
-        return hash((message.data.ndim, str(message.data.dtype), n_channels, message.axes["time"].gain))
+        backend = "sparse" if isinstance(message.data, sparse.SparseArray) else get_namespace(message.data).__name__
+        return hash((message.data.ndim, str(message.data.dtype), n_channels, message.axes["time"].gain, backend))
 
     def _reset_state(self, message: AxisArray) -> None:
         """Initialize state for new input stream."""
         n_channels = message.data.shape[message.get_axis_idx("ch")] if "ch" in message.dims else 1
 
         self._state.activation = np.zeros(n_channels, dtype=np.float64)
+        self._state.dense_carry = None
         self._state.samples_since_update = np.zeros(n_channels, dtype=np.int64)
 
         # For alpha kernel, we need auxiliary state
@@ -420,13 +425,17 @@ def _process_dense_count_sum(self, message: AxisArray) -> AxisArray:
         else:
             contrib = (data != 0).astype(xp.float32)
 
-        # Pull state into the input namespace for on-device math.
-        overflow_xp = xp.asarray(self._state.activation.reshape(feature_shape)).astype(xp.float32)
+        # Keep dense partial-bin state in the input namespace. In particular, do
+        # not round-trip an MLX carry through np.asarray here: that synchronizes
+        # the device on every source chunk, including chunks that close no bin.
+        overflow_xp = (
+            xp.zeros(feature_shape, dtype=xp.float32) if self._state.dense_carry is None else self._state.dense_carry
+        )
 
         if n_bins == 0:
             # No complete bins this chunk — accumulate everything into the carry-over.
             new_overflow = overflow_xp + (xp.sum(contrib, axis=0) if n_samples > 0 else overflow_xp * 0)
-            self._state.activation = np.asarray(new_overflow).reshape(self._state.activation.shape)
+            self._state.dense_carry = new_overflow
             return replace(
                 message,
                 data=xp.zeros((0,) + feature_shape, dtype=xp.float32),
@@ -467,7 +476,7 @@ def _process_dense_count_sum(self, message: AxisArray) -> AxisArray:
             new_overflow = xp.sum(contrib[last_bin_end:], axis=0)
         else:
             new_overflow = xp.zeros(feature_shape, dtype=cumsum.dtype)
-        self._state.activation = np.asarray(new_overflow).reshape(self._state.activation.shape)
+        self._state.dense_carry = new_overflow
 
         if self.settings.rate_normalize:
             output = output / step.output_gain

src/ezmsg/event/peak.py

@@ -154,7 +154,7 @@ def _reset_state(self, message: AxisArray) -> None:
             self._state.elapsed = np.full((n_features,), self._state.refrac_width + 1, dtype=np.int32)
 
     def _can_use_mlx_metal(self, xp, data) -> bool:
-        """The Metal kernel runs automatically for MLX + DENSE + a config it supports.
+        """The Metal kernel runs automatically for MLX + dense + a supported config.
 
         It cannot recover peak values, align on peaks, enforce a min peak width, or
         auto-scale, so any of those settings disable the path.
@@ -253,6 +253,12 @@ def _process(self, message: AxisArray) -> AxisArray:
         # (the cpu path's prepended buffer alone yields no new crossings, and the metal
         # kernel has nothing to scan). Short-circuit before backend dispatch.
         if message.data.shape[0] == 0:
+            # _hash_message ignores the sample count, so an empty first chunk still triggers
+            # _reset_state, which seeds _state.data from data[:1] -> an empty time axis. Force
+            # a re-reset so the first non-empty chunk seeds the prev-sample reference (the
+            # metal path indexes _state.data[0], and the cpu path prepends it).
+            if self._state.data is not None and self._state.data.shape[0] == 0:
+                self._request_reset()
             return self._empty_output(message, xp)
 
         # MLX-on-device fast path: bypass the numpy event-detection logic and run

tests/test_peak.py

@@ -183,6 +183,33 @@ def _require_mlx():
     return mx
 
 
+def test_threshold_crossing_defaults_to_sparse():
+    proc = ThresholdCrossingTransformer(ThresholdSettings(threshold=-1.0, refrac_dur=0.0))
+    msg = AxisArray(
+        data=np.zeros((10, 2), dtype=np.float32),
+        dims=["time", "ch"],
+        axes={"time": AxisArray.TimeAxis(fs=1000.0)},
+    )
+
+    result = proc(msg)
+
+    assert isinstance(result.data, sparse.SparseArray)
+
+
+def test_threshold_crossing_mlx_defaults_to_sparse():
+    mx = _require_mlx()
+    proc = ThresholdCrossingTransformer(ThresholdSettings(threshold=-1.0, refrac_dur=0.0))
+    msg = AxisArray(
+        data=mx.zeros((10, 2), dtype=mx.float32),
+        dims=["time", "ch"],
+        axes={"time": AxisArray.TimeAxis(fs=1000.0)},
+    )
+
+    result = proc(msg)
+
+    assert isinstance(result.data, sparse.SparseArray)
+
+
 @pytest.mark.parametrize(
     ("threshold", "refrac_dur", "stride"),
     [
@@ -292,3 +319,33 @@ def test_threshold_crossing_empty_time_first(return_peak_val: bool, auto_scale_t
     out_normal = proc(msg_normal)
     assert isinstance(out_normal.data, sparse.SparseArray)
     assert out_normal.data.shape[1] == N_CH
+
+
+def test_threshold_crossing_empty_time_first_mlx_metal():
+    """Empty → normal for MLX inputs: the empty first chunk must not leave the Metal path with
+    an empty prev-sample reference (regression for `[take] ... from an empty axis`)."""
+    mx = _require_mlx()
+    fs = 1000.0
+
+    proc = ThresholdCrossingTransformer(
+        ThresholdSettings(
+            threshold=-1.0,
+            refrac_dur=0.001,
+            output_format=OutputFormat.DENSE,
+        )
+    )
+
+    def mlx_msg(n_time: int) -> AxisArray:
+        return AxisArray(
+            data=mx.array(np.random.randn(n_time, N_CH).astype(np.float32)),
+            dims=["time", "ch"],
+            axes={"time": AxisArray.TimeAxis(fs=fs)},
+        )
+
+    out_empty = proc(mlx_msg(0))
+    assert out_empty.data.shape[0] == 0
+
+    out_normal = proc(mlx_msg(64))
+    mx.eval(out_normal.data)
+    assert not isinstance(out_normal.data, sparse.SparseArray)
+    assert out_normal.data.shape == (64, N_CH)

tests/test_rate.py

@@ -157,6 +157,36 @@ def test_rate_dense_input_matches_sparse(fs: float, bin_duration: float, refrac_
         np.testing.assert_allclose(sp_msg.data, ds_msg.data)
 
 
+def test_rate_dense_threshold_preserves_mlx_backend():
+    mx = pytest.importorskip("mlx.core")
+    try:
+        mx.eval(mx.array([1.0], dtype=mx.float32))
+    except RuntimeError as exc:
+        pytest.skip(f"MLX device unavailable: {exc}")
+
+    fs = 30_000.0
+    data = np.zeros((600, 4), dtype=np.float32)
+    data[10] = -2.0
+    data[310] = -2.0
+    message = _make_msg(mx.array(data), fs, 0.0)
+
+    threshold = ThresholdCrossingTransformer(
+        threshold=-1.0,
+        refrac_dur=0.001,
+        output_format=OutputFormat.DENSE,
+    )
+    rate = Rate(EventRateSettings(bin_duration=0.02, fractional=False))
+
+    events = threshold(message)
+    result = rate(events)
+    mx.eval(events.data, result.data)
+
+    assert isinstance(events.data, mx.array)
+    assert isinstance(result.data, mx.array)
+    assert isinstance(rate.state.dense_carry, mx.array)
+    np.testing.assert_allclose(np.asarray(result.data), np.full((1, 4), 100.0))
+
+
 def test_rate_empty_time_first():
     """Empty → normal: empty first message triggers _reset_state on empty data."""
     proc = Rate(EventRateSettings(bin_duration=0.02))