# SPDX-License-Identifier: Apache-2.0
# Copyright 2026 Andrew Hendel. Part of mojolearn, https://doi.org/10.5281/zenodo.22068632
"""Kernel methods: kernel ridge regression and kernel feature approximations.
`KernelRidge`, `Nystroem` and `RBFSampler` follow scikit-learn's estimators of
the same names: every parameter means what scikit-learn's parameter of that
name means, and an invalid value raises an error that names the parameter.
Inputs are converted to float32; sparse input is densified.
`KernelRidge` and `Nystroem` support the kernels 'linear', 'poly' (or
'polynomial'), 'rbf', 'sigmoid', 'laplacian', 'cosine', 'chi2' and
'additive_chi2'. `KernelRidge` also accepts 'precomputed'; `Nystroem` refuses
it with `ValueError`. `RBFSampler` approximates the RBF kernel only.
`gamma=None` is scikit-learn's `1 / n_features` (1.0 for 'chi2').
A fitted model keeps every array it computed, so it can be inspected, saved
with `save` and loaded with `load`; `mojolearn.host_model(path)` predicts or
transforms from a saved model on a CPU without a GPU.
"""
# Maintainer notes (moved from the module docstring):
# The Python door of `kernel_methods/estimator.mojo` through
# `bindings/_mojolearn_kernel_methods.mojo` (workstream D, 2026-09-14).
# `kernel_methods/README.md` carries the parameter mapping to scikit-learn;
# every refusal is raised by name, here for the shapes and the spellings, on the
# Mojo host for the values (`km_validate_matrix`, `km_validate_kernel_params`,
# the alpha, gamma, degree and n_components refusals, DEVIATION 1686).
# THE MODEL IS ITS ARRAYS. Fit returns every field `kernel_methods/
# estimator.mojo`'s structs carry, so a fitted object can be read, saved and
# hashed, and predict/transform send the same arrays back down to rebuild the
# struct. `Nystroem` keeps its eigenvalues, eigenvectors and Jacobi sweep count,
# which scikit-learn discards, because the lane's identity argument is which of
# those moved (the estimator's header, point 3).
# The kernels are `svm/impl/svm_parameter.mojo`'s codes plus the lane's
# laplacian (and cosine, chi2, additive_chi2). `gamma=None` is scikit-learn's
# `1 / n_features`, one Python division, correctly rounded on every host.
# No speed claim: the lane has no published number and this door adds none.
from . import _backend, _serialize
from ._scale_gamma import scale_gamma, scale_gamma_x
from ._array import Array
from ._buffer import addr, addr_ro, all_finite, as_f32_c, as_f32_dense_c, as_f64_c, empty, full
from ._lazy_out import _empty_out, _mapped_out
from ._mode import NumericModeMixin
#: The saved-model formats (lane/inference-linear-svm, 2026-09-15):
#: `mojolearn.host_model(path)` predicts or transforms from each on a CPU.
_KERNEL_RIDGE_FORMAT = "mojolearn-kernel-ridge-1"
_NYSTROEM_FORMAT = "mojolearn-nystroem-1"
_RBF_SAMPLER_FORMAT = "mojolearn-rbf-sampler-1"
def _km_header(arrays, path, cls, meta_fields, hyper_fields):
"""The estimator, `meta` `<i8` and `hyper` `<f8` checks the three loads
share; returns `(meta, hyper)`."""
from .linear_model import _check_saved_by
_check_saved_by(arrays, path, cls)
meta = _serialize.exact(arrays, "meta", "<i8")
if meta.size != meta_fields:
raise ValueError(f"mojolearn: {path!r} meta holds {meta.size} fields, {meta_fields} are needed")
hyper = _serialize.exact(arrays, "hyper", "<f8")
if hyper.size != hyper_fields:
raise ValueError(f"mojolearn: {path!r} hyper holds {hyper.size} fields, {hyper_fields} are needed")
return meta, hyper
def _km_array(arrays, name, dtype, shape, path):
value = _serialize.exact(arrays, name, dtype)
if tuple(value.shape) != tuple(shape):
raise ValueError(f"mojolearn: {path!r} {name} shape {tuple(value.shape)} is not {tuple(shape)}")
return value
def _km_write(est, path, fmt, arrays):
from .linear_model import _saved_mode
arrays.update(format=fmt, estimator=type(est).__name__, numeric_mode=_saved_mode(est))
return _serialize.write_npz(path, arrays)
#: `checks/numerics.mojo` codes, duplicated on purpose (the GP's reason):
#: the read-back must not share a table with the thing it checks.
_MODE_CODE = {"fast": 0, "identical": 1, "deterministic": 2}
#: `kernel_methods/checks/kernel_matrix.mojo`'s KM_KERNEL_* codes:
#: `svm/impl/svm_parameter.mojo`'s enumeration (0 to 4) plus the lane's
#: laplacian at 5. A renumbering on either side is a WRONG KERNEL, not a
#: failure, which is why the surface test plants a fixture whose answer
#: moves if the codes move.
KERNEL_LINEAR = 0
KERNEL_POLYNOMIAL = 1
KERNEL_RBF = 2
KERNEL_SIGMOID = 3
KERNEL_PRECOMPUTED = 4
KERNEL_LAPLACIAN = 5
KERNEL_COSINE = 6
KERNEL_CHI2 = 7
KERNEL_ADDITIVE_CHI2 = 8
_KERNELS = {
"linear": KERNEL_LINEAR,
"poly": KERNEL_POLYNOMIAL,
"polynomial": KERNEL_POLYNOMIAL,
"rbf": KERNEL_RBF,
"sigmoid": KERNEL_SIGMOID,
"laplacian": KERNEL_LAPLACIAN,
"cosine": KERNEL_COSINE,
"chi2": KERNEL_CHI2,
"additive_chi2": KERNEL_ADDITIVE_CHI2,
}
#: Kernels that read gamma. Under 'linear' it is ignored on both sides,
#: exactly as scikit-learn's `pairwise_kernels` ignores it.
_NEEDS_GAMMA = frozenset({KERNEL_POLYNOMIAL, KERNEL_RBF, KERNEL_SIGMOID, KERNEL_LAPLACIAN})
def _kernel_code(kernel, where, precomputed=False):
if kernel == "precomputed" and precomputed:
return KERNEL_PRECOMPUTED
if kernel == "precomputed":
raise ValueError(
f"mojolearn {where}: kernel='precomputed' is refused by name; the "
"lane forms the kernel matrix on the device and takes no matrix "
"from the caller (kernel_methods/NOT_IMPLEMENTED.tsv)"
)
if isinstance(kernel, str):
if kernel not in _KERNELS:
raise ValueError(
f"mojolearn {where}: kernel must be one of "
f"{sorted(_KERNELS)}, got {kernel!r}"
)
return _KERNELS[kernel]
if isinstance(kernel, bool) or not isinstance(kernel, int):
raise TypeError(
f"mojolearn {where}: kernel must be a name or a KM_KERNEL code, "
f"got {type(kernel).__name__}"
)
return int(kernel)
def _real(v, name, where):
if isinstance(v, bool) or not isinstance(v, (int, float)):
raise TypeError(
f"mojolearn {where}: {name} must be a real number, got "
f"{type(v).__name__}"
)
return float(v)
def _gamma_for(gamma, n_features, where, kernel=None):
"""scikit-learn's `gamma=None` is `1 / n_features`, except for
`chi2_kernel`, whose own default is `gamma=1.0`."""
if gamma is None and kernel == KERNEL_CHI2:
return 1.0
if gamma is None:
return 1.0 / float(n_features)
return _real(gamma, "gamma", where)
def _check_chi2_input(x, kernel, where):
"""scikit-learn's chi2 kernels refuse negative input ("X contains
negative values"); so do these, before any device work."""
if kernel in (KERNEL_CHI2, KERNEL_ADDITIVE_CHI2) and x.size and float(x.min()) < 0.0:
raise ValueError(f"mojolearn {where}: X contains negative values (the chi2 kernels need x >= 0)")
def _sqrt_weights(sample_weight, n, where):
"""The float64 sample weights the binding takes the square roots of
(`np.sqrt(np.atleast_1d(sample_weight))`: each factor the correctly
rounded binary64 square root rounded once to float32, in Mojo). The
finiteness and sign checks run in the core helpers. None when
unweighted."""
if sample_weight is None:
return None
if isinstance(sample_weight, (int, float)) and not isinstance(sample_weight, bool):
w = full((n,), float(sample_weight), "<f8")
else:
w, _ = as_f64_c(sample_weight, ndim=1, name="sample_weight")
if w.shape[0] != n:
raise ValueError(
f"mojolearn {where}: sample_weight has {w.shape[0]} entries, X has {n} rows"
)
if n and (not all_finite(w) or w.min() < 0.0):
raise ValueError(
f"mojolearn {where}: sample_weight must be finite and >= 0"
)
return w
class _KernelMethodBase(NumericModeMixin):
_BINDING = "_mojolearn_kernel_methods"
_WHERE = "kernel method"
def _extension(self):
"""The binding for THIS estimator's tier, with the binary's own
compile-time answer cross-checked against it (the GP's pattern)."""
mod = self._bind()
want = getattr(self, "numeric_mode", None) or _backend.default_mode()
fn = getattr(mod, "kernel_methods_numeric_mode", None)
if fn is not None:
got = int(fn())
if got != _MODE_CODE.get(want):
raise RuntimeError(
f"mojolearn {self._WHERE}: numeric_mode={want!r} was requested "
f"but {mod.__name__} reports compile-time mode code {got}; "
"the binary and the directory it sits in disagree, rebuild "
"it with bash bindings/build_kernel_methods.sh"
)
return mod
# Maintainer notes: `alpha` IS the Cholesky ridge (DEVIATION 1660); a kernel
# matrix that does not factor is refused by name, never solved by a
# least-squares fallback (DEVIATION 1662). Negative or NaN alpha and
# non-positive gamma are refused on the Mojo host. The cosine and chi2 kernels
# are each cell one pinned chain. `info_` is LAPACK's info from the
# factorization.
[docs]
class KernelRidge(_KernelMethodBase):
"""Kernel ridge regression.
Follows `sklearn.kernel_ridge.KernelRidge` (and cuML's). Fitting forms
the kernel matrix `K` of the training rows, adds `alpha` to its
diagonal, factors it with a Cholesky factorization and solves for the
dual coefficients; prediction is `K(X, X_fit_) @ dual_coef_`. The
training matrix is kept on the model because prediction needs it.
If `K + alpha * I` cannot be factored at the given `alpha`, `fit`
raises an error rather than falling back to a least-squares solve;
increase `alpha` in that case. Inputs and outputs are float32. Supports
the `'identical'` (default) and `'fast'` numeric modes; `'fast'` is
available on Apple GPUs only.
Parameters
----------
alpha : float, default 1.0
Regularization strength, added to the kernel matrix diagonal. Must
be >= 0; zero is accepted, negative or NaN is refused.
kernel : str, default 'linear'
One of 'linear', 'poly', 'rbf', 'sigmoid', 'laplacian', 'cosine',
'chi2', 'additive_chi2' or 'precomputed'.
'cosine', 'chi2' and 'additive_chi2' are
scikit-learn's pairwise kernels; both chi2 kernels refuse negative
values in `X` with `ValueError`, as scikit-learn's do. With
'precomputed', `X` is the n x n kernel matrix at fit and the
q x n cross-kernel against the training rows at predict, used as
given (symmetry is not checked).
gamma : float or None, default None
Kernel coefficient for 'poly', 'rbf', 'sigmoid', 'laplacian' and
the chi2 kernels. None means `1 / n_features` (1.0 for 'chi2', its
scikit-learn default). Must be positive where the kernel reads it.
degree : int, default 3
Degree of the 'poly' kernel.
coef0 : float, default 1.0
Independent term of the 'poly' and 'sigmoid' kernels.
numeric_mode : {'identical', 'fast'} or None, default None
Numeric mode for this estimator; None uses the process default.
Attributes
----------
dual_coef_ : Array of shape (n_samples,) or (n_samples, n_targets), float32
Dual coefficients, 1-D when `y` was 1-D.
X_fit_ : Array of shape (n_samples, n_features), float32
The training data, used by `predict`.
info_ : int
Status of the factorization; always 0 on a fitted model.
n_features_in_ : int
Number of features seen during `fit`.
n_targets_ : int
Number of regression targets.
Examples
--------
>>> import numpy as np, mojolearn
>>> rng = np.random.default_rng(0)
>>> X = rng.random((1000, 8), dtype=np.float32)
>>> y = X.sum(axis=1)
>>> model = mojolearn.KernelRidge(alpha=1.0, kernel="rbf").fit(X, y)
>>> pred = model.predict(X[:5])
"""
#: scikit-learn's estimator kind: `cross_val_score` stratifies a
#: classifier's default folds, as scikit-learn's does.
_estimator_type = "regressor"
_WHERE = "KernelRidge"
[docs]
def __init__(self, alpha=1.0, kernel="linear", gamma=None, degree=3, coef0=1.0):
self.alpha = alpha
self.kernel = kernel
self.gamma = gamma
self.degree = degree
self.coef0 = coef0
def _kp(self, n_features):
k = _kernel_code(self.kernel, self._WHERE, precomputed=True)
if isinstance(self.degree, bool) or not isinstance(self.degree, int):
raise TypeError(
f"mojolearn {self._WHERE}: degree must be an int, got "
f"{type(self.degree).__name__}"
)
return k, int(self.degree), _gamma_for(self.gamma, n_features, self._WHERE, k), _real(self.coef0, "coef0", self._WHERE)
# Maintainer notes: `kernel_ridge_fit_host`. `sample_weight` is
# scikit-learn's `_solve_cholesky_kernel` weighted arm: `sw = sqrt(w)`
# (binary64, rounded once to float32), `y * sw`, `K * outer(sw, sw)`,
# then `dual_coef * sw` (float32 on both columns, DEVIATION 1688).
[docs]
def fit(self, X, y, sample_weight=None):
"""Fit the kernel ridge model.
Parameters
----------
X : array-like of shape (n_samples, n_features)
Training data, converted to float32. With
`kernel='precomputed'`, the square n x n kernel matrix.
y : array-like of shape (n_samples,) or (n_samples, n_targets)
Target values.
sample_weight : float or array-like of shape (n_samples,), default None
Per-sample weights, finite and >= 0, applied as scikit-learn
does. None weights every sample equally.
Returns
-------
self : KernelRidge
The fitted estimator.
"""
x, _ = as_f32_dense_c(X, ndim=2, name="X")
n, d = x.shape
yy, _ = as_f32_c(y, ndim=None, name="y")
if yy.ndim == 1:
t, squeeze = 1, True
elif yy.ndim == 2:
t, squeeze = yy.shape[1], False
else:
raise ValueError(f"mojolearn {self._WHERE}: y must be 1-D or 2-D, got {yy.ndim}-D")
if yy.shape[0] != n:
raise ValueError(
f"mojolearn {self._WHERE}: y has {yy.shape[0]} rows, X has {n}"
)
kernel, degree, gamma, coef0 = self._kp(d)
_check_chi2_input(x, kernel, self._WHERE)
if kernel == KERNEL_PRECOMPUTED and d != n:
raise ValueError(
f"mojolearn {self._WHERE}: kernel='precomputed' needs the square "
f"n x n kernel matrix as X, got {n} x {d}"
)
alpha = _real(self.alpha, "alpha", self._WHERE)
flat_y = yy.reshape((n * t,))
dual = empty((n * t,), "<f4")
scalars = empty((1,), "<f8")
sw = _sqrt_weights(sample_weight, n, self._WHERE)
tail = [] if sw is None else [addr_ro(sw, name="sample_weight")]
ext = self._extension()
# glue: the pointer entry when the binding registers it (IDENTICAL
# device builds, `KRR_IDN_PTR_IN`); the same addrs and params.
fit_fn = getattr(ext, "kernel_ridge_fit_ptr", None) or ext.kernel_ridge_fit
info = fit_fn(
# ORDER MATCHES bindings/_mojolearn_kernel_methods.mojo::kernel_ridge_fit_binding.
# x, y, dual_out, scalars_out[, sw]
[addr_ro(x, name="X"), addr_ro(flat_y, name="y"), addr(dual, name="dual_coef_"), addr(scalars, name="scalars")] + tail,
# n, d, t, kernel, degree, gamma, coef0, alpha
[n, d, t, kernel, degree, gamma, coef0, alpha],
)
self.X_fit_ = x
self.n_features_in_ = d
self.n_targets_ = t
self._squeeze = squeeze
self.dual_coef_ = dual if squeeze else dual.reshape((n, t))
self.info_ = int(info)
self._kernel_params = (kernel, degree, gamma, coef0, alpha)
return self
# Maintainer note: `K(X, X_fit_) . dual_coef_` through the identical GEMM
# (DEVIATION 1680).
[docs]
def predict(self, X):
"""Predict with the fitted model: `K(X, X_fit_) @ dual_coef_`.
Parameters
----------
X : array-like of shape (n_queries, n_features)
Query data. With `kernel='precomputed'`, the kernel between the
queries and the training rows, shape (n_queries, n_samples).
Returns
-------
Array of shape (n_queries,) or (n_queries, n_targets), float32
Predictions, 1-D when the fitted `y` was 1-D.
"""
if not hasattr(self, "dual_coef_"):
raise ValueError(f"mojolearn {self._WHERE}: call fit before predict")
xq, _ = as_f32_dense_c(X, ndim=2, name="X")
q, d = xq.shape
if d != self.n_features_in_:
raise ValueError(
f"mojolearn {self._WHERE}: X has {d} features, the fit had {self.n_features_in_}"
)
n, t = self.X_fit_.shape[0], self.n_targets_
kernel, degree, gamma, coef0, alpha = self._kernel_params
_check_chi2_input(xq, kernel, self._WHERE)
flat_dual = self.dual_coef_.reshape((n * t,))
out = empty((q * t,), "<f4")
ext = self._extension()
# glue: the pointer entry when the binding registers it (`KRR_IDN_PTR_IN`)
predict_fn = getattr(ext, "kernel_ridge_predict_ptr", None) or ext.kernel_ridge_predict
predict_fn(
# ORDER MATCHES bindings/_mojolearn_kernel_methods.mojo::kernel_ridge_predict_binding.
# x_fit, dual, x_new, out
[addr_ro(self.X_fit_, name="X_fit_"), addr_ro(flat_dual, name="dual_coef_"), addr_ro(xq, name="X"), addr(out, name="predict")],
# n, d, t, kernel, degree, gamma, coef0, alpha, info, q
[n, d, t, kernel, degree, gamma, coef0, alpha, self.info_, q],
)
if self._squeeze:
return out
return out.reshape((q, t))
# Maintainer note, file layout: `X_fit_` and `dual_coef_` (flat,
# `n * n_targets`) as fitted, `hyper` `<f8` [gamma, coef0, alpha] as the fit
# resolved them, `meta` `<i8` [n, n_features, n_targets, kernel code, degree,
# info, 1-D target] and the numeric mode.
[docs]
def save(self, path):
"""Save the fitted model to `path` as an `.npz` file.
The file holds the training data, the dual coefficients, the kernel
parameters as resolved at fit and the numeric mode. Load it with
`KernelRidge.load`, or predict from it on a CPU without a GPU with
`mojolearn.host_model(path)`.
Parameters
----------
path : str or path-like
Destination file.
Raises
------
RuntimeError
If the model has not been fitted.
"""
if not hasattr(self, "dual_coef_"):
raise RuntimeError("this estimator is not fitted yet")
kernel, degree, gamma, coef0, alpha = self._kernel_params
n, t = self.X_fit_.shape[0], self.n_targets_
return _km_write(self, path, _KERNEL_RIDGE_FORMAT, {
"x_fit": self.X_fit_,
"dual": self.dual_coef_.reshape((n * t,)),
"hyper": Array.from_list([float(gamma), float(coef0), float(alpha)], "<f8"),
"meta": Array.from_list([int(n), int(self.n_features_in_), int(t), int(kernel), int(degree),
int(self.info_), 1 if self._squeeze else 0], "<i8"),
})
[docs]
@classmethod
def load(cls, path):
"""Load a model written by `save`.
Parameters
----------
path : str or path-like
File written by `KernelRidge.save`.
Returns
-------
KernelRidge
A fitted model, ready to predict, with the saved numeric mode.
"""
from .linear_model import _restore_mode
arrays = _serialize.read_npz(path, _KERNEL_RIDGE_FORMAT)
meta, hyper = _km_header(arrays, path, cls, 7, 3)
n, d, t, kernel, degree, info, squeeze = (int(v) for v in meta.tolist()) # glue: seven saved meta fields
gamma, coef0, alpha = (float(v) for v in hyper.tolist()) # glue: three saved hyperparameters
obj = cls(alpha=alpha, kernel=kernel, gamma=gamma, degree=degree, coef0=coef0)
_restore_mode(obj, arrays)
obj.X_fit_ = _km_array(arrays, "x_fit", "<f4", (n, d), path)
dual = _km_array(arrays, "dual", "<f4", (n * t,), path)
obj.n_features_in_ = d
obj.n_targets_ = t
obj._squeeze = bool(squeeze)
obj.dual_coef_ = dual if squeeze else dual.reshape((n, t))
obj.info_ = info
obj._kernel_params = (kernel, degree, gamma, coef0, alpha)
return obj
# Maintainer notes: `random_state` seeds `km_basis_indices`, the device
# permutation. `normalization_` is NOT bitwise symmetric (DEVIATION 1674);
# transform uses its transpose, scikit-learn's arm. `eigenvalues_` are clipped
# at 1e-12 (DEVIATION 1670). n_components > n_samples is refused on the Mojo
# host. The eigenvalues, eigenvectors and sweep count are kept (scikit-learn
# discards them) because they are part of the identity argument.
[docs]
class Nystroem(_KernelMethodBase):
"""Approximate a kernel map with the Nystroem method.
Follows `sklearn.kernel_approximation.Nystroem`. `fit` picks
`n_components` training rows at random as the basis, forms their
kernel matrix and computes its inverse square root with a Jacobi
eigensolver on the device; `transform` maps rows to
`K(X, components_) @ normalization_.T`, a feature space whose inner
products approximate the kernel. Inputs and outputs are float32.
Supports the `'identical'` (default) and `'fast'` numeric modes;
`'fast'` is available on Apple GPUs only.
Parameters
----------
kernel : str, default 'rbf'
One of 'linear', 'poly', 'rbf', 'sigmoid', 'laplacian', 'cosine',
'chi2' or 'additive_chi2'. 'precomputed' is refused with
`ValueError`. Both chi2 kernels refuse negative values in `X`.
gamma : float or None, default None
Kernel coefficient. None means `1 / n_features` (1.0 for 'chi2').
degree : int, default 3
Degree of the 'poly' kernel.
coef0 : float, default 1.0
Independent term of the 'poly' and 'sigmoid' kernels.
n_components : int, default 100
Number of basis rows. Must be positive and at most `n_samples`.
random_state : int, default 0
Seed of the basis selection.
numeric_mode : {'identical', 'fast'} or None, default None
Numeric mode for this estimator; None uses the process default.
Attributes
----------
components_ : Array of shape (n_components, n_features), float32
The basis rows.
component_indices_ : Array of shape (n_components,), int32
Indices of the basis rows in the training data.
normalization_ : Array of shape (n_components, n_components), float32
Inverse square root of the basis kernel matrix.
eigenvalues_ : Array of shape (n_components,), float32
Absolute eigenvalues of the basis kernel matrix (its singular
values), descending. A negative eigenvalue keeps its sign in the
right factor of `normalization_`, as scikit-learn's SVD does.
eigenvectors_ : Array of shape (n_components, n_components), float32
Eigenvectors of the basis kernel matrix.
sweeps_ : int
Number of Jacobi sweeps the eigensolver ran.
n_features_in_ : int
Number of features seen during `fit`.
Examples
--------
>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> features = mojolearn.Nystroem(n_components=64).fit_transform(X)
"""
_WHERE = "Nystroem"
[docs]
def __init__(self, kernel="rbf", gamma=None, degree=3, coef0=1.0, n_components=100, random_state=0):
self.kernel = kernel
self.gamma = gamma
self.degree = degree
self.coef0 = coef0
self.n_components = n_components
self.random_state = random_state
[docs]
def fit(self, X, y=None):
"""Choose the basis rows and compute the normalization.
Parameters
----------
X : array-like of shape (n_samples, n_features)
Training data, converted to float32.
y : ignored
Present for scikit-learn API compatibility.
Returns
-------
self : Nystroem
The fitted transformer.
"""
x, _ = as_f32_dense_c(X, ndim=2, name="X")
n, d = x.shape
kernel = _kernel_code(self.kernel, self._WHERE)
if isinstance(self.degree, bool) or not isinstance(self.degree, int):
raise TypeError(f"mojolearn {self._WHERE}: degree must be an int")
if isinstance(self.n_components, bool) or not isinstance(self.n_components, int):
raise TypeError(f"mojolearn {self._WHERE}: n_components must be an int")
q = int(self.n_components)
if q < 1:
raise ValueError(f"mojolearn {self._WHERE}: n_components must be positive, got {q}")
gamma = _gamma_for(self.gamma, d, self._WHERE, kernel)
_check_chi2_input(x, kernel, self._WHERE)
coef0 = _real(self.coef0, "coef0", self._WHERE)
seed = int(self.random_state)
components = empty((q * d,), "<f4")
indices = empty((q,), "<i4")
normalization = empty((q * q,), "<f4")
eigenvalues = empty((q,), "<f4")
eigenvectors = empty((q * q,), "<f4")
scalars = empty((1,), "<f8")
ext = self._extension()
# glue: the pointer entry when the binding registers it (`NYS_IDN_FIT_PTR_IN`)
fit_fn = getattr(ext, "nystroem_fit_ptr", None) or ext.nystroem_fit
sweeps = fit_fn(
# ORDER MATCHES bindings/_mojolearn_kernel_methods.mojo::nystroem_fit_binding.
# x, components_out, indices_out, normalization_out, eigenvalues_out, eigenvectors_out, scalars_out
[addr_ro(x, name="X"), addr(components, name="components_"), addr(indices, name="component_indices_"),
addr(normalization, name="normalization_"), addr(eigenvalues, name="eigenvalues_"),
addr(eigenvectors, name="eigenvectors_"), addr(scalars, name="scalars")],
# n, d, kernel, degree, gamma, coef0, q, seed
[n, d, kernel, int(self.degree), gamma, coef0, q, seed],
)
self.n_features_in_ = d
self.components_ = components.reshape((q, d))
self.component_indices_ = indices
self.normalization_ = normalization.reshape((q, q))
self.eigenvalues_ = eigenvalues
self.eigenvectors_ = eigenvectors.reshape((q, q))
self.sweeps_ = int(sweeps)
self._kernel_params = (kernel, int(self.degree), gamma, coef0, seed)
return self
# Maintainer note, file layout: every model array as fitted, `hyper` `<f8`
# [gamma, coef0] as the fit resolved them, `meta` `<i8` [n_components,
# n_features, kernel code, degree, seed, sweeps] and the numeric mode.
[docs]
def save(self, path):
"""Save the fitted transformer to `path` as an `.npz` file.
Load it with `Nystroem.load`, or transform from it on a CPU without
a GPU with `mojolearn.host_model(path)`.
Parameters
----------
path : str or path-like
Destination file.
Raises
------
RuntimeError
If the transformer has not been fitted.
"""
if not hasattr(self, "components_"):
raise RuntimeError("this estimator is not fitted yet")
kernel, degree, gamma, coef0, seed = self._kernel_params
q, d = self.components_.shape
return _km_write(self, path, _NYSTROEM_FORMAT, {
"components": self.components_,
"component_indices": self.component_indices_,
"normalization": self.normalization_,
"eigenvalues": self.eigenvalues_,
"eigenvectors": self.eigenvectors_,
"hyper": Array.from_list([float(gamma), float(coef0)], "<f8"),
"meta": Array.from_list([int(q), int(d), int(kernel), int(degree), int(seed), int(self.sweeps_)], "<i8"),
})
[docs]
@classmethod
def load(cls, path):
"""Load a transformer written by `save`.
Parameters
----------
path : str or path-like
File written by `Nystroem.save`.
Returns
-------
Nystroem
A fitted transformer, ready to transform, with the saved
numeric mode.
"""
from .linear_model import _restore_mode
arrays = _serialize.read_npz(path, _NYSTROEM_FORMAT)
meta, hyper = _km_header(arrays, path, cls, 6, 2)
q, d, kernel, degree, seed, sweeps = (int(v) for v in meta.tolist()) # glue: six saved meta fields
gamma, coef0 = (float(v) for v in hyper.tolist()) # glue: two saved hyperparameters
obj = cls(kernel=kernel, gamma=gamma, degree=degree, coef0=coef0, n_components=q, random_state=seed)
_restore_mode(obj, arrays)
obj.n_features_in_ = d
obj.components_ = _km_array(arrays, "components", "<f4", (q, d), path)
obj.component_indices_ = _km_array(arrays, "component_indices", "<i4", (q,), path)
obj.normalization_ = _km_array(arrays, "normalization", "<f4", (q, q), path)
obj.eigenvalues_ = _km_array(arrays, "eigenvalues", "<f4", (q,), path)
obj.eigenvectors_ = _km_array(arrays, "eigenvectors", "<f4", (q, q), path)
obj.sweeps_ = sweeps
obj._kernel_params = (kernel, degree, gamma, coef0, seed)
return obj
# Maintainer notes: gamma <= 0 is refused on the Mojo host. gamma='scale' uses
# the EXACT variance of X's float32 cells, the reciprocal rounded once
# (`_scale_gamma.scale_gamma`). `sigma_` and `scale_` are computed once on the
# host (DEVIATION 1678).
[docs]
class RBFSampler(_KernelMethodBase):
"""Approximate the RBF kernel with random Fourier features.
Follows `sklearn.kernel_approximation.RBFSampler`. `fit` draws random
weights and offsets from `random_state`; `transform` maps rows to
`scale_ * cos(X @ random_weights_ + random_offset_)`, a feature space
whose inner products approximate `exp(-gamma * ||x - y||^2)`. `fit`
reads only the number of features of `X` (and its variance when
`gamma='scale'`). Inputs and outputs are float32. Supports the
`'identical'` (default) and `'fast'` numeric modes; `'fast'` is
available on Apple GPUs only.
Parameters
----------
gamma : float or 'scale', default 1.0
RBF kernel coefficient; must be positive. 'scale' is scikit-learn's
`1 / (n_features * X.var())` (1.0 when the variance is zero),
computed from the exact variance of the float32 data.
n_components : int, default 100
Number of random features. Must be positive.
random_state : int, default 0
Seed of the random draws.
numeric_mode : {'identical', 'fast'} or None, default None
Numeric mode for this estimator; None uses the process default.
Attributes
----------
random_weights_ : Array of shape (n_features, n_components), float32
The random projection weights.
random_offset_ : Array of shape (n_components,), float32
The random phase offsets.
sigma_ : float
`sqrt(2 * gamma)`, the standard deviation of the weights.
scale_ : float
`sqrt(2 / n_components)`, the output scale.
n_features_in_ : int
Number of features seen during `fit`.
Examples
--------
>>> import numpy as np, mojolearn
>>> X = np.random.default_rng(0).random((1000, 8), dtype=np.float32)
>>> features = mojolearn.RBFSampler(gamma=0.5, n_components=256).fit_transform(X)
"""
_WHERE = "RBFSampler"
[docs]
def __init__(self, gamma=1.0, n_components=100, random_state=0):
self.gamma = gamma
self.n_components = n_components
self.random_state = random_state
[docs]
def fit(self, X, y=None):
"""Draw the random weights and offsets.
Reads only `X.shape[1]` (and, for `gamma='scale'`, the variance of
`X`), as scikit-learn does.
Parameters
----------
X : array-like of shape (n_samples, n_features)
Training data, converted to float32.
y : ignored
Present for scikit-learn API compatibility.
Returns
-------
self : RBFSampler
The fitted transformer.
"""
x, _ = as_f32_dense_c(X, ndim=2, name="X")
_, d = x.shape
gamma, q, seed = self._fit_params(x, d)
weights = empty((d * q,), "<f4")
offset = empty((q,), "<f4")
scalars = empty((2,), "<f8")
self._extension().rbf_sampler_fit(
# ORDER MATCHES bindings/_mojolearn_kernel_methods.mojo::rbf_sampler_fit_binding.
# weights_out, offset_out, scalars_out
[addr(weights, name="random_weights_"), addr(offset, name="random_offset_"), addr(scalars, name="scalars")],
# d, q, gamma, seed
[d, q, gamma, seed],
)
self._set_fitted(d, q, gamma, seed, weights, offset, scalars)
return self
def _set_fitted(self, d, q, gamma, seed, weights, offset, scalars):
self.n_features_in_ = d
self.random_weights_ = weights.reshape((d, q))
self.random_offset_ = offset
self.sigma_ = float(scalars[0])
self.scale_ = float(scalars[1])
self._params = (gamma, seed)
def _fit_params(self, x, d):
"""(gamma, n_components, seed) as `fit` resolves and refuses them."""
if isinstance(self.gamma, str) and self.gamma == "scale":
# scikit-learn's 1 / (n_features * X.var()), 1.0 at zero
# variance: the EXACT variance of the float32 cells, the
# reciprocal rounded once (`_scale_gamma.scale_gamma`), so every
# host draws from the same gamma bits.
gamma = scale_gamma_x(self._extension(), x, d)
else:
gamma = _real(self.gamma, "gamma", self._WHERE)
if isinstance(self.n_components, bool) or not isinstance(self.n_components, int):
raise TypeError(f"mojolearn {self._WHERE}: n_components must be an int")
q = int(self.n_components)
if q < 1:
# BY NAME, BEFORE ANY BUFFER. The first MI300X run of this surface
# (2026-09-14) saw n_components=0 refused as "null float32 buffer
# address": the zero-length weights buffer reached the binding's
# address check before anything had said which argument was wrong.
raise ValueError(f"mojolearn {self._WHERE}: n_components must be positive, got {q}")
seed = int(self.random_state)
return gamma, q, seed
# Maintainer note, file layout: the weights and offsets as drawn, `hyper`
# `<f8` [gamma, sigma_, scale_], `meta` `<i8` [n_features, n_components,
# seed] and the numeric mode.
[docs]
def save(self, path):
"""Save the fitted sampler to `path` as an `.npz` file.
Load it with `RBFSampler.load`, or transform from it on a CPU
without a GPU with `mojolearn.host_model(path)`.
Parameters
----------
path : str or path-like
Destination file.
Raises
------
RuntimeError
If the sampler has not been fitted.
"""
if not hasattr(self, "random_weights_"):
raise RuntimeError("this estimator is not fitted yet")
gamma, seed = self._params
d, q = self.random_weights_.shape
return _km_write(self, path, _RBF_SAMPLER_FORMAT, {
"random_weights": self.random_weights_,
"random_offset": self.random_offset_,
"hyper": Array.from_list([float(gamma), float(self.sigma_), float(self.scale_)], "<f8"),
"meta": Array.from_list([int(d), int(q), int(seed)], "<i8"),
})
[docs]
@classmethod
def load(cls, path):
"""Load a sampler written by `save`.
Parameters
----------
path : str or path-like
File written by `RBFSampler.save`.
Returns
-------
RBFSampler
A fitted sampler with the saved weights (nothing is drawn
again), ready to transform, with the saved numeric mode.
"""
from .linear_model import _restore_mode
arrays = _serialize.read_npz(path, _RBF_SAMPLER_FORMAT)
meta, hyper = _km_header(arrays, path, cls, 3, 3)
d, q, seed = (int(v) for v in meta.tolist()) # glue: three saved meta fields
gamma, sigma, scale = (float(v) for v in hyper.tolist()) # glue: three saved hyperparameters
obj = cls(gamma=gamma, n_components=q, random_state=seed)
_restore_mode(obj, arrays)
obj.n_features_in_ = d
obj.random_weights_ = _km_array(arrays, "random_weights", "<f4", (d, q), path)
obj.random_offset_ = _km_array(arrays, "random_offset", "<f4", (q,), path)
obj.sigma_ = sigma
obj.scale_ = scale
obj._params = (gamma, seed)
return obj
__all__ = ["KernelRidge", "Nystroem", "RBFSampler"]
