nltools 0.6.0.dev3__tar.gz → 0.6.0.dev4__tar.gz

{nltools-0.6.0.dev3 → nltools-0.6.0.dev4}/.gitignore

@@ -96,3 +96,6 @@ docs/tutorials/workflows/[0-9]*_*.md
 # page outputs docs_show replays (pages/). Deliberately not under .cache/, which
 # is zensical's and which `zensical build -c` wipes whole.
 /.tutorial-cache/
+
+# Local planning artifacts (specs, plans, SDD workspaces)
+.superpowers/

{nltools-0.6.0.dev3 → nltools-0.6.0.dev4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: nltools
-Version: 0.6.0.dev3
+Version: 0.6.0.dev4
 Summary: A Python package to analyze neuroimaging data
 Project-URL: Homepage, https://nltools.org
 Author-email: "Luke J. Chang" <luke.j.chang@dartmouth.edu>, Eshin Jolly <eshin.jolly@gmail.com>

{nltools-0.6.0.dev3 → nltools-0.6.0.dev4}/nltools/models/ridge.py

@@ -12,7 +12,7 @@ import numbers
 from collections.abc import Mapping, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
-from typing import Any
+from typing import Any, Literal
 
 import numpy as np
 
@@ -71,6 +71,31 @@ def _scoped_himalaya_backend(name: str):
         set_backend(previous, on_error="raise")
 
 
+def _solver_form(
+    n_samples: int, n_features: int, n_spaces: int = 1
+) -> Literal["primal", "kernel"]:
+    """Return `'kernel'` when the design is wide, otherwise `'primal'`.
+
+    Himalaya's flowchart routes designs with more features than samples to its
+    kernel-form solvers, whose cross-validation working set scales with the
+    sample count instead of the feature count. Both forms solve the same
+    problem; only the cost differs. A banded fit holds one `(n_samples,
+    n_samples)` kernel per space, so it is only cheaper in the kernel form when
+    the average space is wider than the sample count, not when the total is. A
+    tie stays primal, which for one space is also the threshold of Himalaya's
+    own "slower than kernel ridge" warning.
+
+    Args:
+        n_samples (int): Rows of the design.
+        n_features (int): Total columns across every feature space.
+        n_spaces (int): Number of feature spaces. Defaults to 1.
+
+    Returns:
+        str: `'kernel'` or `'primal'`.
+    """
+    return "kernel" if n_spaces * n_samples < n_features else "primal"
+
+
 def _himalaya_backend_name(backend) -> str:
     """Map a resolved nltools `_Backend` to its Himalaya backend name.
 
@@ -100,26 +125,62 @@ def _batch_sizes(
     n_targets: int,
     n_alphas: int,
     itemsize: int,
+    *,
+    solver_form: str = "primal",
+    n_spaces: int = 1,
 ) -> dict[str, int]:
     """Size the batches of a whole cross-validated or banded fit.
 
     Only the per-item working-set estimates live here; the budget itself and
     the batch arithmetic come from `nltools.algorithms.backends`. The estimates
-    follow Himalaya's dominant allocations: decomposition matrices of
-    `(n_alphas_batch, n_features, n_samples)`, cross-validated predictions of
-    `(n_alphas_batch, n_samples, n_targets_batch)`, and refit weights of
-    `(n_features, n_targets_batch_refit)`.
+    follow Himalaya's dominant allocations. In the primal form they are
+    decomposition matrices of `(n_alphas_batch, n_features, n_samples)`,
+    cross-validated predictions of `(n_alphas_batch, n_samples,
+    n_targets_batch)`, and refit weights of `(n_alphas_batch, n_features,
+    n_targets_batch_refit)`. In the kernel form every `n_features` above
+    becomes `n_samples`: the decomposition is of the `(n_samples, n_samples)`
+    kernel and the refit weights are dual.
+
+    The arrays that stay on the device for the whole fit are charged first,
+    and the batches are sized from what remains: the design and targets in
+    both forms, plus in the kernel form one `(n_samples, n_samples)` kernel per
+    space and the same-sized kernel Himalaya sums or slices per fold.
+
+    Args:
+        backend (Backend): Backend returned by `_resolve_backend`.
+        memory_budget_gb (float | None): Explicit budget, or None to measure.
+        n_samples (int): Rows of the design.
+        n_features (int): Total columns across every feature space.
+        n_targets (int): Columns of `y`.
+        n_alphas (int): Candidate alphas.
+        itemsize (int): Bytes per element of the working dtype.
+        solver_form (str): `'primal'` or `'kernel'`, from `_solver_form`.
+        n_spaces (int): Number of feature spaces. Defaults to 1.
 
     Returns:
         dict[str, int]: `n_targets_batch`, `n_targets_batch_refit`, and
             `n_alphas_batch`.
+
+    Raises:
+        ValueError: If the resident arrays alone exceed the budget, or one
+            batch item does.
     """
     budget_gb = _device_memory_budget(
         backend, max_gpu_memory_gb=memory_budget_gb, cap_for_batching=True
     )
+    resident_bytes = (n_samples * n_features + n_samples * n_targets) * itemsize
+    if solver_form == "kernel":
+        resident_bytes += (n_spaces + 1) * n_samples * n_samples * itemsize
+    budget_gb -= resident_bytes / 1e9
+    if budget_gb <= 0:
+        raise ValueError(
+            f"the resident design, targets and kernels need "
+            f"{resident_bytes / 1e9:.6g} GB, exceeding the budget"
+        )
+    weight_rows = n_samples if solver_form == "kernel" else n_features
     n_alphas_batch, _ = _auto_batch_size(
         n_alphas,
-        n_features * n_samples * itemsize,
+        weight_rows * n_samples * itemsize,
         budget_gb=budget_gb,
         overhead=_WORKING_SET_OVERHEAD,
     )
@@ -131,7 +192,7 @@ def _batch_sizes(
     )
     n_targets_batch_refit, _ = _auto_batch_size(
         n_targets,
-        n_alphas_batch * n_features * itemsize,
+        n_alphas_batch * weight_rows * itemsize,
         budget_gb=budget_gb,
         overhead=_WORKING_SET_OVERHEAD,
     )
@@ -156,8 +217,11 @@ def _refit_targets_batch(
 
     With one shared alpha Himalaya reuses a single shrinkage operator, so a
     target costs only its own columns of `Y` and of the weights. With a
-    per-target alpha it instead holds an `(n_targets_batch, n_samples,
-    n_samples)` block, which dominates everything else.
+    per-target alpha it instead holds one square matrix per target, which
+    dominates everything else: `(n_samples, n_samples)` in the kernel form a
+    wide design runs, and `(n_features, n_samples)` with `n_features <=
+    n_samples` in the primal form a tall one runs, so the sample-count square
+    bounds both.
 
     Returns:
         int: Target batch size in `[1, n_targets]`.
@@ -401,6 +465,12 @@ def _refit_fixed_hyperparameters(
     `sqrt(gamma[k])` before the solve and the resulting coefficients are scaled
     back, which is equivalent to the per-space penalty `alpha / gamma[k]`.
 
+    The scaled design is one ordinary ridge system, so it follows Himalaya's
+    flowchart with one space: a tall design runs `solve_ridge_svd`, a wide one
+    forms its linear kernel, runs `solve_kernel_ridge_eigenvalues`, and turns
+    the dual weights back into coefficients with Himalaya's own host-side
+    `primal_weights_kernel_ridge`.
+
     Targets that selected the same weight vector share one decomposition. The
     grouping is an implementation detail and does not change the result.
 
@@ -427,6 +497,11 @@ def _refit_fixed_hyperparameters(
         np.ndarray: Coefficients of shape `(n_features, n_targets)` in the
             original, unscaled feature coordinates, as CPU NumPy.
     """
+    from himalaya.kernel_ridge import (
+        linear_kernel,
+        primal_weights_kernel_ridge,
+        solve_kernel_ridge_eigenvalues,
+    )
     from himalaya.ridge import solve_ridge_svd
 
     if isinstance(feature_spaces, _ResidentDesign):
@@ -444,6 +519,7 @@ def _refit_fixed_hyperparameters(
 
     alphas = np.broadcast_to(np.asarray(alpha, dtype=np.float64), (n_targets,))
     coef = np.zeros((n_features, n_targets), dtype=np.float64)
+    solver_form = _solver_form(n_samples, n_features)
 
     with _scoped_himalaya_backend(resident.backend_name):
         stacked = _take_rows(resident.design, row_indices)
@@ -459,8 +535,8 @@ def _refit_fixed_hyperparameters(
                 ).astype(dtype)
                 design = stacked * _on_active_backend(scale)[0]
             # A shared alpha needs one shrinkage vector; a per-target alpha
-            # makes Himalaya hold an (n_targets_batch, n_samples, n_samples)
-            # block instead, so the two paths get different batch estimates.
+            # makes Himalaya hold a block of one square matrix per target
+            # instead, so the two paths get different batch estimates.
             group_alphas = alphas[columns]
             shared_alpha = bool(np.all(group_alphas == group_alphas[0]))
             batch = n_targets_batch
@@ -479,16 +555,29 @@ def _refit_fixed_hyperparameters(
                 if len(columns) == n_targets
                 else _take_columns(all_targets, columns)
             )
-            solved = solve_ridge_svd(
-                design,
-                targets,
-                alpha=dtype.type(group_alphas[0])
+            group_alpha = (
+                dtype.type(group_alphas[0])
                 if shared_alpha
-                else group_alphas.astype(dtype),
-                fit_intercept=False,
-                n_targets_batch=batch,
-                warn=False,
+                else group_alphas.astype(dtype)
             )
+            if solver_form == "kernel":
+                dual = solve_kernel_ridge_eigenvalues(
+                    linear_kernel(design),
+                    targets,
+                    alpha=group_alpha,
+                    fit_intercept=False,
+                    n_targets_batch=batch,
+                )
+                solved = primal_weights_kernel_ridge(dual, design)
+            else:
+                solved = solve_ridge_svd(
+                    design,
+                    targets,
+                    alpha=group_alpha,
+                    fit_intercept=False,
+                    n_targets_batch=batch,
+                    warn=False,
+                )
             solved = np.asarray(_to_cpu_numpy(solved), dtype=np.float64)
             if scale is not None:
                 solved = solved * scale[:, None]
@@ -516,6 +605,49 @@ def _on_active_backend(*arrays):
     return [backend.asarray(array) for array in arrays]
 
 
+def _to_backend_arrays(spaces, targets, alphas, dtype):
+    """Move a fit's inputs onto Himalaya's active backend in the working dtype.
+
+    The designs come back as a list even when there is one, because Himalaya's
+    banded solvers take a list of spaces: their `check_arrays` converts a list
+    element by element, which is what lets spaces of differing widths through,
+    whereas a tuple would be handed to `asarray` whole and fail as ragged.
+
+    Args:
+        spaces (Sequence[np.ndarray]): Feature matrices in coefficient order,
+            already contiguous in `dtype`.
+        targets (np.ndarray): `(n_samples, n_targets)` targets in `dtype`.
+        alphas (np.ndarray): Candidate alphas, float64.
+        dtype (np.dtype): Working dtype for the alpha grid.
+
+    Returns:
+        tuple: `(designs, targets, alphas)` on the active backend, with
+            `designs` a list holding one entry per space.
+    """
+    converted = _on_active_backend(*spaces, targets, alphas.astype(dtype))
+    return converted[:-2], converted[-2], converted[-1]
+
+
+def _linear_kernels(spaces):
+    """Stack the linear kernels of several feature spaces for the banded solver.
+
+    Each kernel is Himalaya's own `linear_kernel`, called inside the same
+    `_scoped_himalaya_backend` block as the solver that consumes it, so it is
+    built on the device and in the dtype the solve runs in.
+
+    Args:
+        spaces (Sequence): Feature matrices already on the active backend,
+            each `(n_samples, n_features_k)`.
+
+    Returns:
+        Array: Shape `(n_spaces, n_samples, n_samples)` on the active backend.
+    """
+    from himalaya.backend import get_backend
+    from himalaya.kernel_ridge import linear_kernel
+
+    return get_backend().stack([linear_kernel(space) for space in spaces])
+
+
 def _to_cpu_numpy(array) -> np.ndarray:
     """Return `array` as CPU NumPy, whatever backend produced it.
 
@@ -663,7 +795,10 @@ class _Ridge:
         cv_scores_ (float | np.ndarray | None): None for a fixed-alpha fit. For
             ordinary Ridge, the fold-averaged negative-MSE score at the selected
             alpha. For banded Ridge, `(search_iterations,)` or
-            `(search_iterations, n_targets)` fold-averaged scores.
+            `(search_iterations, n_targets)` fold-averaged scores. Stored as
+            Himalaya reports them: in the kernel form a candidate alpha below
+            the float32 rounding floor of the linear kernel scores `-1e5`, and
+            a target whose every candidate scored that way keeps the first.
         feature_space_weights_ (np.ndarray | None): None for ordinary Ridge.
             Strictly positive weights whose columns sum to one, shaped
             `(n_spaces,)` or `(n_spaces, n_targets)`.
@@ -672,6 +807,8 @@ class _Ridge:
         feature_space_sizes_ (tuple[int, ...] | None): Feature counts aligned
             with `feature_space_names_`; None for ordinary Ridge.
         backend_ (Backend): The resolved execution backend.
+        solver_form_ (str): `'primal'` or `'kernel'`, the Himalaya solver family
+            the fit ran. Chosen from the design shape by `_solver_form`.
         n_samples_ (int): Fitted sample count.
         n_features_in_ (int): Total fitted feature count across spaces.
         is_fitted_ (bool): True after a successful fit.
@@ -1010,12 +1147,15 @@ class _Ridge:
         n_targets = targets.shape[1]
         alphas = np.atleast_1d(np.asarray(alpha, dtype=np.float64))
 
-        self.backend_ = backend
+        # Every fitted attribute is assigned only once the solve has succeeded,
+        # so a fit that raises leaves the previous fitted state intact.
         if scalar_alpha:
             # The fixed-alpha refit sizes its own batch from the same budget; it
             # never runs the cross-validation or alpha loops the others measure.
-            self._fit_fixed_alpha(spaces, targets, float(alpha))
+            solver_form = _solver_form(n_samples, n_features)
+            self._fit_fixed_alpha(backend, spaces, targets, float(alpha))
         else:
+            solver_form = _solver_form(n_samples, n_features, len(spaces))
             try:
                 batches = _batch_sizes(
                     backend,
@@ -1025,6 +1165,8 @@ class _Ridge:
                     n_targets=n_targets,
                     n_alphas=alphas.size,
                     itemsize=dtype.itemsize,
+                    solver_form=solver_form,
+                    n_spaces=len(spaces),
                 )
             except ValueError as error:
                 raise ValueError(
@@ -1032,11 +1174,21 @@ class _Ridge:
                     f"fit with n_samples={n_samples}, n_features={n_features}, "
                     f"n_targets={n_targets}, n_alphas={alphas.size} ({error})"
                 ) from error
-            if is_banded:
-                self._fit_banded(spaces, targets, alphas, dtype, batches)
+            if is_banded and solver_form == "kernel":
+                self._fit_banded_kernel(
+                    backend, spaces, targets, alphas, dtype, batches
+                )
+            elif is_banded:
+                self._fit_banded(backend, spaces, targets, alphas, dtype, batches)
+            elif solver_form == "kernel":
+                self._fit_ordinary_cv_kernel(
+                    backend, spaces, targets, alphas, dtype, batches
+                )
             else:
-                self._fit_ordinary_cv(spaces, targets, alphas, dtype, batches)
+                self._fit_ordinary_cv(backend, spaces, targets, alphas, dtype, batches)
 
+        self.backend_ = backend
+        self.solver_form_ = solver_form
         self.feature_space_names_ = names
         self.feature_space_sizes_ = sizes if is_banded else None
         self.n_samples_ = int(n_samples)
@@ -1046,10 +1198,11 @@ class _Ridge:
         self.is_fitted_ = True
         return self
 
-    def _fit_fixed_alpha(self, spaces, targets, alpha) -> None:
+    def _fit_fixed_alpha(self, backend, spaces, targets, alpha) -> None:
         """Solve a fixed-alpha ordinary Ridge and store the fitted state.
 
         Args:
+            backend (Backend): Resolved execution backend.
             spaces (list[np.ndarray]): One feature matrix.
             targets (np.ndarray): `(n_samples, n_targets)` targets.
             alpha (float): The fixed regularization strength.
@@ -1058,17 +1211,41 @@ class _Ridge:
             spaces,
             targets,
             alpha,
-            backend=self.backend_,
+            backend=backend,
             memory_budget_gb=self.memory_budget_gb,
         )
         self.alpha_ = float(alpha)
         self.cv_scores_ = None
         self.feature_space_weights_ = None
 
-    def _fit_ordinary_cv(self, spaces, targets, alphas, dtype, batches) -> None:
+    def _store_ordinary_state(self, coef, best_alphas, cv_scores, alphas) -> None:
+        """Normalize an ordinary cross-validated solve into the fitted state.
+
+        Shared by the primal and kernel forms, whose solvers differ only in how
+        `coef` is obtained. The selected alphas are snapped back onto the
+        candidate grid, and every array is stored as CPU float64.
+
+        Args:
+            coef: `(n_features, n_targets)` coefficients, on any backend.
+            best_alphas: `(n_targets,)` selected alphas, on any backend.
+            cv_scores: Fold-averaged scores at the selected alpha, on any backend.
+            alphas (np.ndarray): Candidate alphas.
+        """
+        self.coef_ = np.asarray(_to_cpu_numpy(coef), dtype=np.float64)
+        selected = _snap_to_grid(_to_cpu_numpy(best_alphas), alphas)
+        self.alpha_ = selected if self.per_target_alpha else float(selected[0])
+        self.cv_scores_ = np.asarray(
+            _to_cpu_numpy(cv_scores), dtype=np.float64
+        ).reshape(-1)
+        self.feature_space_weights_ = None
+
+    def _fit_ordinary_cv(
+        self, backend, spaces, targets, alphas, dtype, batches
+    ) -> None:
         """Select an alpha by cross-validation and store the fitted state.
 
         Args:
+            backend (Backend): Resolved execution backend.
             spaces (list[np.ndarray]): One feature matrix.
             targets (np.ndarray): `(n_samples, n_targets)` targets.
             alphas (np.ndarray): Candidate alphas.
@@ -1078,12 +1255,12 @@ class _Ridge:
         from himalaya.ridge import solve_ridge_cv_svd
         from himalaya.scoring import l2_neg_loss
 
-        with _scoped_himalaya_backend(_himalaya_backend_name(self.backend_)):
-            design, y_device, alpha_device = _on_active_backend(
-                spaces[0], targets, alphas.astype(dtype)
+        with _scoped_himalaya_backend(_himalaya_backend_name(backend)):
+            designs, y_device, alpha_device = _to_backend_arrays(
+                spaces, targets, alphas, dtype
             )
             best_alphas, coefs, cv_scores = solve_ridge_cv_svd(
-                design,
+                designs[0],
                 y_device,
                 alphas=alpha_device,
                 fit_intercept=False,
@@ -1095,49 +1272,131 @@ class _Ridge:
                 **batches,
             )
 
-        self.coef_ = np.asarray(_to_cpu_numpy(coefs), dtype=np.float64)
-        selected = _snap_to_grid(_to_cpu_numpy(best_alphas), alphas)
-        self.alpha_ = float(selected[0]) if not self.per_target_alpha else selected
-        self.cv_scores_ = np.asarray(
-            _to_cpu_numpy(cv_scores), dtype=np.float64
-        ).reshape(-1)
-        self.feature_space_weights_ = None
+        self._store_ordinary_state(coefs, best_alphas, cv_scores, alphas)
 
-    def _fit_banded(self, spaces, targets, alphas, dtype, batches) -> None:
-        """Run the banded random search and store the fitted state.
+    def _fit_ordinary_cv_kernel(
+        self, backend, spaces, targets, alphas, dtype, batches
+    ) -> None:
+        """Select an alpha by cross-validation in the kernel form and store the state.
+
+        The wide-design counterpart of `_fit_ordinary_cv`. Himalaya solves on the
+        `(n_samples, n_samples)` linear kernel and returns dual weights on the
+        host; its own `primal_weights_kernel_ridge` turns them back into `coef_`
+        on the host, which is where Himalaya keeps primal weights because they
+        can be large. The device copy of the design is released once the kernel
+        exists, so only the kernel stays resident through the solve.
 
         Args:
-            spaces (list[np.ndarray]): Feature matrices in coefficient order.
+            backend (Backend): Resolved execution backend.
+            spaces (list[np.ndarray]): One feature matrix.
             targets (np.ndarray): `(n_samples, n_targets)` targets.
             alphas (np.ndarray): Candidate alphas.
             dtype (np.dtype): Working dtype.
             batches (dict[str, int]): Himalaya batch sizes.
         """
-        from himalaya.kernel_ridge import generate_dirichlet_samples
-        from himalaya.ridge import solve_group_ridge_random_search
+        from himalaya.kernel_ridge import (
+            linear_kernel,
+            primal_weights_kernel_ridge,
+            solve_kernel_ridge_cv_eigenvalues,
+        )
         from himalaya.scoring import l2_neg_loss
 
-        # Himalaya's sampler ends with `get_backend().asarray(gammas)`, so the
-        # candidates would otherwise take the dtype and device of whatever
-        # backend happened to be globally active. They are validated and clamped
-        # on the host, so draw them under an explicit numpy scope.
+        with _scoped_himalaya_backend(_himalaya_backend_name(backend)):
+            designs, y_device, alpha_device = _to_backend_arrays(
+                spaces, targets, alphas, dtype
+            )
+            kernel = linear_kernel(designs[0])
+            del designs
+            best_alphas, dual_weights, cv_scores = solve_kernel_ridge_cv_eigenvalues(
+                kernel,
+                y_device,
+                alphas=alpha_device,
+                fit_intercept=False,
+                score_func=l2_neg_loss,
+                cv=self._resolved_cv(),
+                local_alpha=self.per_target_alpha,
+                conservative=self.prefer_conservative_alpha,
+                **batches,
+            )
+
+        dual = np.asarray(_to_cpu_numpy(dual_weights), dtype=dtype)
+        with _scoped_himalaya_backend("numpy"):
+            coef = primal_weights_kernel_ridge(dual, spaces[0])
+        self._store_ordinary_state(coef, best_alphas, cv_scores, alphas)
+
+    def _draw_feature_space_candidates(self, n_spaces: int, dtype) -> np.ndarray:
+        """Draw and condition the Dirichlet candidates for a banded search.
+
+        Himalaya's sampler ends with `get_backend().asarray(gammas)`, so the
+        candidates would otherwise take the dtype and device of whatever
+        backend happened to be globally active. They are validated and clamped
+        on the host, so they are drawn under an explicit numpy scope.
+
+        Args:
+            n_spaces (int): Number of feature spaces.
+            dtype (np.dtype): Working dtype, which sets the underflow floor.
+
+        Returns:
+            np.ndarray: `(search_iterations, n_spaces)` weights, each row on
+                the simplex and every entry at least `finfo(dtype).tiny`.
+        """
+        from himalaya.kernel_ridge import generate_dirichlet_samples
+
         with _scoped_himalaya_backend("numpy"):
             candidates = _to_cpu_numpy(
                 generate_dirichlet_samples(
                     n_samples=self.search_iterations,
-                    n_kernels=len(spaces),
+                    n_kernels=n_spaces,
                     concentration=self._concentration_for_himalaya(),
                     random_state=self.random_state,
                 )
             )
-        candidates = _prepare_feature_space_weights(candidates, dtype)
-
-        with _scoped_himalaya_backend(_himalaya_backend_name(self.backend_)):
-            converted = _on_active_backend(*spaces, targets, alphas.astype(dtype))
-            designs, y_device, alpha_device = (
-                converted[:-2],
-                converted[-2],
-                converted[-1],
+        return _prepare_feature_space_weights(candidates, dtype)
+
+    def _store_banded_state(self, deltas, refit_weights, cv_scores, alphas) -> None:
+        """Normalize a banded random search into the fitted state.
+
+        Shared by the primal and kernel forms. Himalaya reports the banded
+        solution as `deltas = log(gamma / alpha)` with each gamma column summing
+        to one, so the simplex weights and the selected alpha both fall out of a
+        log-sum-exp over the spaces. The recovered alpha is snapped back onto
+        the candidate grid, and every array is stored as CPU float64.
+
+        Args:
+            deltas: `(n_spaces, n_targets)` deltas, on any backend.
+            refit_weights: `(n_features, n_targets)` coefficients in original
+                feature coordinates, on any backend.
+            cv_scores: `(search_iterations, n_targets)` scores, on any backend.
+            alphas (np.ndarray): Candidate alphas.
+        """
+        from scipy.special import logsumexp, softmax
+
+        deltas = np.asarray(_to_cpu_numpy(deltas), dtype=np.float64)
+        self.coef_ = np.asarray(_to_cpu_numpy(refit_weights), dtype=np.float64)
+        self.cv_scores_ = np.asarray(_to_cpu_numpy(cv_scores), dtype=np.float64)
+        self.feature_space_weights_ = softmax(deltas, axis=0)
+        selected = _snap_to_grid(np.exp(-logsumexp(deltas, axis=0)), alphas)
+        self.alpha_ = selected if self.per_target_alpha else float(selected[0])
+
+    def _fit_banded(self, backend, spaces, targets, alphas, dtype, batches) -> None:
+        """Run the banded random search in the primal form and store the fitted state.
+
+        Args:
+            backend (Backend): Resolved execution backend.
+            spaces (list[np.ndarray]): Feature matrices in coefficient order.
+            targets (np.ndarray): `(n_samples, n_targets)` targets.
+            alphas (np.ndarray): Candidate alphas.
+            dtype (np.dtype): Working dtype.
+            batches (dict[str, int]): Himalaya batch sizes.
+        """
+        from himalaya.ridge import solve_group_ridge_random_search
+        from himalaya.scoring import l2_neg_loss
+
+        candidates = self._draw_feature_space_candidates(len(spaces), dtype)
+
+        with _scoped_himalaya_backend(_himalaya_backend_name(backend)):
+            designs, y_device, alpha_device = _to_backend_arrays(
+                spaces, targets, alphas, dtype
             )
             deltas, refit_weights, cv_scores = solve_group_ridge_random_search(
                 designs,
@@ -1156,19 +1415,68 @@ class _Ridge:
                 **batches,
             )
 
-        deltas = np.asarray(_to_cpu_numpy(deltas), dtype=np.float64)
-        self.coef_ = np.asarray(_to_cpu_numpy(refit_weights), dtype=np.float64)
-        self.cv_scores_ = np.asarray(_to_cpu_numpy(cv_scores), dtype=np.float64)
+        self._store_banded_state(deltas, refit_weights, cv_scores, alphas)
 
-        # deltas = log(gamma / alpha) with each gamma column summing to one, so
-        # the simplex weights and the selected alpha both fall out of a
-        # log-sum-exp over the spaces.
-        shifted = deltas - deltas.max(axis=0, keepdims=True)
-        weights = np.exp(shifted)
-        self.feature_space_weights_ = weights / weights.sum(axis=0, keepdims=True)
-        log_total = deltas.max(axis=0) + np.log(np.exp(shifted).sum(axis=0))
-        selected = _snap_to_grid(np.exp(-log_total), alphas)
-        self.alpha_ = selected if self.per_target_alpha else float(selected[0])
+    def _fit_banded_kernel(
+        self, backend, spaces, targets, alphas, dtype, batches
+    ) -> None:
+        """Run the banded random search in the kernel form and store the fitted state.
+
+        The wide-design counterpart of `_fit_banded`. Himalaya searches over one
+        linear kernel per feature space and returns dual weights, as its own
+        `MultipleKernelRidgeCV` asks it to; its `primal_weights_weighted_kernel_ridge`
+        then recovers `coef_` once on the host from the raw spaces and the
+        deltas. Asking the solver for primal weights instead would rebuild the
+        gamma-scaled design on the device for every improving candidate. The
+        deltas carry the same meaning as in the primal search and go through
+        the same recovery.
+
+        Args:
+            backend (Backend): Resolved execution backend.
+            spaces (list[np.ndarray]): Feature matrices in coefficient order.
+            targets (np.ndarray): `(n_samples, n_targets)` targets.
+            alphas (np.ndarray): Candidate alphas.
+            dtype (np.dtype): Working dtype.
+            batches (dict[str, int]): Himalaya batch sizes.
+        """
+        from himalaya.kernel_ridge import (
+            primal_weights_weighted_kernel_ridge,
+            solve_multiple_kernel_ridge_random_search,
+        )
+        from himalaya.scoring import l2_neg_loss
+
+        candidates = self._draw_feature_space_candidates(len(spaces), dtype)
+
+        with _scoped_himalaya_backend(_himalaya_backend_name(backend)):
+            designs, y_device, alpha_device = _to_backend_arrays(
+                spaces, targets, alphas, dtype
+            )
+            kernels = _linear_kernels(designs)
+            deltas, refit_weights, cv_scores = (
+                solve_multiple_kernel_ridge_random_search(
+                    kernels,
+                    y_device,
+                    n_iter=candidates,
+                    alphas=alpha_device,
+                    fit_intercept=False,
+                    score_func=l2_neg_loss,
+                    cv=self._resolved_cv(),
+                    return_weights="dual",
+                    local_alpha=self.per_target_alpha,
+                    random_state=self.random_state,
+                    progress_bar=self.progress_bar,
+                    conservative=self.prefer_conservative_alpha,
+                    **batches,
+                )
+            )
+
+        deltas = np.asarray(_to_cpu_numpy(deltas), dtype=np.float64)
+        dual = np.asarray(_to_cpu_numpy(refit_weights), dtype=dtype)
+        with _scoped_himalaya_backend("numpy"):
+            per_space = primal_weights_weighted_kernel_ridge(dual, deltas, spaces)
+        self._store_banded_state(
+            deltas, np.concatenate(per_space, axis=0), cv_scores, alphas
+        )
 
     def _concentration_for_himalaya(self):
         """Return `dirichlet_concentration` in the form Himalaya's sampler takes.