diff --git a/src/ezmsg/learn/process/ssr.py b/src/ezmsg/learn/process/ssr.py index 18b7c43..b696c85 100644 --- a/src/ezmsg/learn/process/ssr.py +++ b/src/ezmsg/learn/process/ssr.py @@ -28,8 +28,9 @@ :meth:`~SelfSupervisedRegressionTransformer.partial_fit` calls. **Solving.** Within each group the weight matrix *W* is obtained from -the inverse of the (ridge-regularised) group covariance -``C_inv = (C_group + lambda * I)^{-1}`` using the block-inverse identity:: +the inverse of the ridge-regularised group covariance, or its pseudoinverse +when an unregularised MLX covariance may be singular, using the block-inverse +identity:: W[:, c] = -C_inv[:, c] / C_inv[c, c], diag(W) = 0 @@ -116,6 +117,11 @@ class SelfSupervisedRegressionSettings(ez.Settings): """When ``True``, accumulate ``X^T X`` across :meth:`partial_fit` calls. When ``False``, each call replaces the previous statistics.""" + check_finite: bool = True + """Skip batches containing non-finite values. Disable for trusted finite + streams to avoid the device-to-host synchronization required to make this + decision on lazy backends such as MLX.""" + @processor_state class SelfSupervisedRegressionState: @@ -230,7 +236,7 @@ def _validate_groups(self, n_channels: int) -> None: # -- weight solving ------------------------------------------------------ def _solve_weights(self, cxx): - """Solve all per-channel ridge regressions via matrix inverse. + """Solve all per-channel ridge regressions via an inverse or pseudoinverse. Uses the block-inverse identity: for target channel *c* with references *r*, ``w_c = -C_inv[r, c] / C_inv[c, c]`` where @@ -253,15 +259,24 @@ def _solve_weights(self, cxx): if groups is None: groups = [np.arange(n, dtype=np.intp)] - W = xp_create(xp.zeros, (n, n), dtype=cxx.dtype, device=dev) - eye_n = xp_create(xp.eye, n, dtype=cxx.dtype, device=dev) + # When every group is already consecutive, collect row blocks and build + # the full matrix once. Otherwise retain the generic selection-matrix + # scatter for arbitrary index orderings. + normalized_groups = [np.asarray(group, dtype=np.intp).reshape(-1) for group in groups] + contiguous_assembly = all( + idx.size == 0 or np.array_equal(idx, np.arange(idx[0], idx[0] + idx.size, dtype=np.intp)) + for idx in normalized_groups + ) + W = None if contiguous_assembly else xp_create(xp.zeros, (n, n), dtype=cxx.dtype, device=dev) + group_blocks = [] + eye_n = None # MLX linalg ops are CPU-only; with unified memory the explicit CPU # stream is a scheduling hint, not a host copy, and results stay mlx. - inv_kwargs = {"stream": xp.cpu} if xp.__name__ == "mlx.core" else {} + is_mlx = xp.__name__ == "mlx.core" + inv_kwargs = {"stream": xp.cpu} if is_mlx else {} - for group in groups: - idx = np.asarray(group, dtype=np.intp).reshape(-1) + for idx in normalized_groups: k = idx.size if k < MIN_REREF_GROUP_SIZE: # Too few channels to rereference against -- leave these channels @@ -279,31 +294,65 @@ def _solve_weights(self, cxx): if self.settings.ridge_lambda > 0: sub = sub + self.settings.ridge_lambda * eye_k - # One inverse per group - try: - sub_inv = xp.linalg.inv(sub, **inv_kwargs) - except Exception: - sub_inv = xp.linalg.pinv(sub, **inv_kwargs) + # One inverse per group. MLX is lazy, so an exception raised while + # evaluating inv() cannot be caught here. With positive ridge the + # covariance is nonsingular; without it, use pinv() directly. + if is_mlx: + solve = xp.linalg.inv if self.settings.ridge_lambda > 0 else xp.linalg.pinv + sub_inv = solve(sub, **inv_kwargs) + else: + try: + sub_inv = xp.linalg.inv(sub) + except Exception: + sub_inv = xp.linalg.pinv(sub) - # Diagonal via element-wise product with identity - diag_vals = xp.sum(sub_inv * eye_k, axis=0) + diag_vals = xp.diag(sub_inv) # w_c = -C_inv[:, c] / C_inv[c, c], vectorised over all c - W_group = -(sub_inv / xp.reshape(diag_vals, (1, k))) + # A completely silent channel has a zero pseudoinverse diagonal; + # leave its prediction weights at zero instead of producing NaNs. + valid_diag = diag_vals != 0 + safe_diag = xp.where(valid_diag, diag_vals, xp.ones_like(diag_vals)) + W_group = -(sub_inv / xp.reshape(safe_diag, (1, k))) + W_group = W_group * xp.reshape(valid_diag, (1, k)) # Zero the diagonal W_group = W_group * (1.0 - eye_k) - # Scatter into full W. The no-op shortcut needs the group to be every - # channel *in order* -- a callable spec may return all n permuted, and - # then the sub-block still has to be scattered back. - if k == n and np.array_equal(idx, np.arange(n, dtype=np.intp)): + # Collect consecutive blocks for one-shot assembly; arbitrary index + # orders still need the generic scatter. + if contiguous_assembly: + start, stop = int(idx[0]), int(idx[-1]) + 1 + group_blocks.append((start, stop, W_group)) + elif k == n and np.array_equal(idx, np.arange(n, dtype=np.intp)): W = W + W_group else: # Selection matrix: columns of eye(n) at group indices + if eye_n is None: + eye_n = xp_create(xp.eye, n, dtype=cxx.dtype, device=dev) S = xp.take(eye_n, idx_xp, axis=1) # (n, k) W = W + xp.matmul(S, xp.matmul(W_group, xp.permute_dims(S, (1, 0)))) + if contiguous_assembly: + row_parts = [] + cursor = 0 + for start, stop, W_group in sorted(group_blocks, key=lambda block: block[0]): + if start > cursor: + row_parts.append(xp_create(xp.zeros, (start - cursor, n), dtype=cxx.dtype, device=dev)) + + col_parts = [] + if start: + col_parts.append(xp_create(xp.zeros, (stop - start, start), dtype=cxx.dtype, device=dev)) + col_parts.append(W_group) + if stop < n: + col_parts.append(xp_create(xp.zeros, (stop - start, n - stop), dtype=cxx.dtype, device=dev)) + row_parts.append(xp.concat(col_parts, axis=1)) + cursor = stop + + if cursor < n: + row_parts.append(xp_create(xp.zeros, (n - cursor, n), dtype=cxx.dtype, device=dev)) + W = xp.concat(row_parts, axis=0) if len(row_parts) > 1 else row_parts[0] + return W # -- partial_fit (self-supervised, accepts AxisArray) -------------------- @@ -311,7 +360,9 @@ def _solve_weights(self, cxx): def partial_fit(self, message: AxisArray) -> None: # type: ignore[override] xp = get_namespace(message.data) - if xp.any(xp.isnan(message.data)): + # This branch necessarily synchronizes lazy device backends. Trusted + # real-time streams can disable it with check_finite=False. + if self.settings.check_finite and not xp.all(xp.isfinite(message.data)): return # Hash check / state reset @@ -346,6 +397,12 @@ def partial_fit(self, message: AxisArray) -> None: # type: ignore[override] self._state.cxx = cxx_new self._state.n_samples += int(X.shape[0]) + # partial_fit has no output to materialize this sufficient statistic. + # Bound MLX's lazy dependency chain without synchronizing the caller; + # intermediate weight solves remain lazy and can be superseded. + if xp.__name__ == "mlx.core": + xp.async_eval(self._state.cxx) + self._state.weights = self._solve_weights(self._state.cxx) self._on_weights_updated() diff --git a/tests/benchmark/bench_lrr.py b/tests/benchmark/bench_lrr.py index ca02614..7cc7a78 100644 --- a/tests/benchmark/bench_lrr.py +++ b/tests/benchmark/bench_lrr.py @@ -8,6 +8,8 @@ 2. partial_fit (training) — numpy, varying chunk sizes 3. _process — torch MPS (Apple Silicon GPU) 4. partial_fit — torch MPS (Apple Silicon GPU) + 5. _process — MLX + 6. partial_fit — MLX, explicitly evaluated """ import time @@ -223,9 +225,6 @@ def bench_process_mlx() -> None: proc = LRRTransformer(LRRSettings(channel_groups=GROUPS)) proc.partial_fit(_make_msg(fit_data)) - def sync(): - mx.eval() - for chunk in CHUNK_SIZES: data_mlx = mx.random.normal(shape=(chunk, N_CH)) msg = _make_msg(data_mlx) @@ -252,48 +251,29 @@ def bench_partial_fit_mlx() -> None: _print_header("partial_fit (training) — MLX") print(f" {N_CH} channels, {N_GROUPS}x{GROUP_SIZE} groups, {WARMUP_ITERS} warmup, {BENCH_ITERS} iters") - # MLX linalg.inv doesn't support GPU yet; run inv on CPU stream - print(" NOTE: linalg.inv runs on mx.cpu stream (GPU not supported)") + print(" NOTE: CPU-stream linalg; check_finite=False; cxx/weights/effective explicitly evaluated") print() - import mlx.core as mx - - _ = np.random.default_rng(1) - proc = LRRTransformer(LRRSettings(channel_groups=GROUPS)) - - # Monkey-patch _solve_weights to use mx.cpu stream for inv - original_solve = proc._solve_weights + for ridge_lambda in (0.0, 1e-3): + solver = "pinv" if ridge_lambda == 0 else "inv" + print(f" ridge_lambda={ridge_lambda:g} ({solver})") + proc = LRRTransformer(LRRSettings(channel_groups=GROUPS, check_finite=False, ridge_lambda=ridge_lambda)) - def _solve_weights_cpu_inv(cxx): - from array_api_compat import get_namespace - - xp = get_namespace(cxx) - # If this is MLX, we need to override linalg.inv - if xp.__name__ == "mlx.core": - orig_inv = mx.linalg.inv - mx.linalg.inv = lambda a: orig_inv(a, stream=mx.cpu) - try: - return original_solve(cxx) - finally: - mx.linalg.inv = orig_inv - return original_solve(cxx) - - proc._solve_weights = _solve_weights_cpu_inv - - for chunk in CHUNK_SIZES: - data_mlx = mx.random.normal(shape=(chunk, N_CH)) - msg = _make_msg(data_mlx) - # Prime - proc.partial_fit(msg) - - def run(): + for chunk in CHUNK_SIZES: + data_mlx = mx.random.normal(shape=(chunk, N_CH)) + msg = _make_msg(data_mlx) + # Prime proc.partial_fit(msg) - mx.eval() - times = _bench_loop(run, WARMUP_ITERS, BENCH_ITERS) - median_us = np.median(times) * 1e6 - throughput = chunk / np.median(times) - _print_row(chunk, median_us, throughput / 1e3) + def run(): + proc.partial_fit(msg) + mx.eval(proc.state.cxx, proc.state.weights, proc.state.effective) + + times = _bench_loop(run, WARMUP_ITERS, BENCH_ITERS) + median_us = np.median(times) * 1e6 + throughput = chunk / np.median(times) + _print_row(chunk, median_us, throughput / 1e3) + print() # --------------------------------------------------------------------------- diff --git a/tests/unit/test_ssr.py b/tests/unit/test_ssr.py index 83ffe05..07ab04a 100644 --- a/tests/unit/test_ssr.py +++ b/tests/unit/test_ssr.py @@ -306,20 +306,75 @@ def test_ridge_handles_collinearity(self): assert out.data.shape == X.shape assert np.all(np.isfinite(out.data)) + def test_mlx_zero_ridge_handles_silent_and_collinear_channels(self): + """MLX must choose pinv before lazy evaluation, where inv failures can + no longer be caught. A completely silent target remains unchanged.""" + mx = pytest.importorskip("mlx.core") + rng = np.random.default_rng(60) + base = rng.standard_normal((200, 1)).astype(np.float32) + X = np.hstack([base, base, np.zeros_like(base), rng.standard_normal((200, 1)).astype(np.float32)]) + msg = _make_axisarray(mx.array(X)) + + proc = LRRTransformer(LRRSettings(ridge_lambda=0.0, check_finite=False)) + proc.partial_fit(msg) + mx.eval(proc.state.cxx, proc.state.weights, proc.state.effective) + + weights = np.asarray(proc.state.weights) + assert np.all(np.isfinite(weights)) + np.testing.assert_array_equal(weights[:, 2], 0.0) + + +class TestGroupAssembly: + def test_contiguous_fast_path_matches_generic_scatter(self): + """Changing only each group's index order forces the generic selection + matrix path without changing the regression represented by the groups.""" + rng = np.random.default_rng(61) + X = _random_data(n_times=300, rng=rng) + contiguous = [[0, 1, 2, 3], [4, 5, 6, 7]] + permuted = [group[::-1] for group in contiguous] + + proc_fast = LRRTransformer(LRRSettings(channel_groups=contiguous)) + proc_fast.partial_fit(_make_axisarray(X)) + proc_generic = LRRTransformer(LRRSettings(channel_groups=permuted)) + proc_generic.partial_fit(_make_axisarray(X)) + + np.testing.assert_allclose(proc_fast.state.weights, proc_generic.state.weights, atol=1e-10) + + +class TestMlxMaterialization: + def test_partial_fit_async_evaluates_covariance(self, monkeypatch): + """Training has no output to bound MLX's lazy covariance graph, so the + sufficient statistic is explicitly queued for asynchronous evaluation.""" + mx = pytest.importorskip("mlx.core") + original_async_eval = mx.async_eval + evaluated = [] + + def record_async_eval(*arrays): + evaluated.extend(arrays) + return original_async_eval(*arrays) + + monkeypatch.setattr(mx, "async_eval", record_async_eval) + X = mx.array(_random_data(n_times=100).astype(np.float32)) + proc = LRRTransformer(LRRSettings(check_finite=False, ridge_lambda=1e-3)) + proc.partial_fit(_make_axisarray(X)) + + assert evaluated == [proc.state.cxx] + -class TestNanDataSkipped: - def test_nan_data_skipped(self): - """partial_fit with NaN data is a no-op.""" +class TestNonfiniteDataSkipped: + @pytest.mark.parametrize("bad_value", [np.nan, np.inf]) + def test_nonfinite_data_skipped(self, bad_value): + """partial_fit with non-finite data is a no-op.""" rng = np.random.default_rng(7) X_good = _random_data(rng=rng) - X_nan = _random_data(rng=rng) - X_nan[0, 0] = np.nan + X_bad = _random_data(rng=rng) + X_bad[0, 0] = bad_value proc = LRRTransformer(LRRSettings()) proc.partial_fit(_make_axisarray(X_good)) W_before = proc.state.weights.copy() - proc.partial_fit(_make_axisarray(X_nan)) + proc.partial_fit(_make_axisarray(X_bad)) np.testing.assert_array_equal(proc.state.weights, W_before)