Optimize MLX LRR training

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,39 +294,75 @@ 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) --------------------
 
     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()
 

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()
 
 
 # ---------------------------------------------------------------------------

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)