diff --git a/src/ezmsg/sigproc/affinetransform.py b/src/ezmsg/sigproc/affinetransform.py index d1c8bfff..23064d0e 100644 --- a/src/ezmsg/sigproc/affinetransform.py +++ b/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)) diff --git a/tests/unit/test_affine_transform.py b/tests/unit/test_affine_transform.py index a4ad2f31..4405efdd 100644 --- a/tests/unit/test_affine_transform.py +++ b/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)))