AffineTransform: handle empty MLX input in block-diagonal path

src/ezmsg/sigproc/affinetransform.py

@@ -305,7 +305,7 @@ def set_weights(self, weights, *, recalc_clusters=False) -> None:
                 all_in = np.concatenate([np.asarray(group) for group in self.settings.channel_clusters])
                 if np.any((all_in < 0) | (all_in >= n_in)):
                     raise ValueError(
-                        "channel_clusters contains out-of-range input indices " f"(valid range: 0..{n_in - 1})"
+                        f"channel_clusters contains out-of-range input indices (valid range: 0..{n_in - 1})"
                     )
 
                 # Derive output indices from non-zero weights for each input cluster
@@ -371,9 +371,13 @@ def _block_diagonal_matmul(self, xp, data, axis_idx):
         out_shape = data.shape[:-1] + (self._state.n_out,)
         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)
-            result[..., out_idx] = xp.matmul(chunk, sub_weights)
+        # Empty input: every cluster write would be zero-size, and MLX's scatter
+        # (indexed assignment) rejects zero-size updates. The zeros allocation is
+        # already the correct result, so skip the assignments.
+        if 0 not in data.shape:
+            for in_idx, out_idx, sub_weights in self._state.clusters:
+                chunk = xp.take(data, in_idx, axis=data.ndim - 1)
+                result[..., out_idx] = xp.matmul(chunk, sub_weights)
 
         if needs_permute:
             inv_dim_perm = list(range(result.ndim))

tests/unit/test_affine_transform.py

@@ -15,7 +15,7 @@
     _merge_small_clusters,
 )
 from tests.helpers.empty_time import N_CH, check_empty_result, check_state_not_corrupted, make_empty_msg, make_msg
-from tests.helpers.util import assert_messages_equal
+from tests.helpers.util import assert_messages_equal, requires_mlx
 
 
 def test_affine_transform():
@@ -1107,3 +1107,48 @@ def test_common_rereference_empty_first():
     result = proc(empty)
     check_empty_result(result)
     check_state_not_corrupted(proc, normal)
+
+
+@requires_mlx
+def test_affine_empty_block_diagonal_mlx():
+    """MLX scatter (indexed assignment) rejects zero-size updates; an empty
+    message through the block-diagonal path must short-circuit the cluster
+    writes and return a correctly shaped empty result."""
+    import mlx.core as mx
+
+    rng = np.random.default_rng(42)
+    n_ch = 64
+    weights = _make_block_diagonal_weights([32, 32], rng=rng)
+
+    def _mlx_msg(n_time: int) -> AxisArray:
+        return AxisArray(
+            data=mx.array(rng.standard_normal((n_time, n_ch)).astype(np.float32)),
+            dims=["time", "ch"],
+            axes={
+                "time": AxisArray.TimeAxis(fs=100.0),
+                "ch": AxisArray.CoordinateAxis(data=np.arange(n_ch).astype(str), dims=["ch"]),
+            },
+            key="test_affine_empty_block_diagonal_mlx",
+        )
+
+    proc = AffineTransformTransformer(AffineTransformSettings(weights=weights, axis="ch"))
+
+    # Empty startup message arrives first; state initializes from it.
+    empty_result = proc(_mlx_msg(0))
+    assert proc._state.clusters is not None and len(proc._state.clusters) == 2
+    assert isinstance(empty_result.data, mx.array)
+    assert empty_result.data.shape == (0, n_ch)
+    assert empty_result.data.dtype == mx.float32
+    assert empty_result.dims == ["time", "ch"]
+    assert np.array_equal(np.asarray(empty_result.axes["ch"].data), np.arange(n_ch).astype(str))
+    check_empty_result(empty_result)
+
+    # Non-empty messages through the same processor still compute correctly.
+    normal = _mlx_msg(4)
+    out = proc(normal)
+    assert isinstance(out.data, mx.array)
+    expected = np.asarray(normal.data) @ weights
+    assert np.allclose(np.asarray(out.data), expected, atol=1e-4)
+
+    # And another empty message after real data also passes through.
+    check_empty_result(proc(_mlx_msg(0)))