dymad.modules.krr

Classes

KRRBase(kernel[, ridge_init, jitter, device])

Base class for Kernel Ridge Regression, in particular:

KRRMultiOutputIndep(kernel[, ridge_init, ...])

Scalar KRR for multiple outputs, and one kernel per output

KRRMultiOutputShared(kernel[, ridge_init, ...])

Scalar KRR for multiple outputs but one single kernel

KRROperatorValued(kernel[, ridge_init, ...])

Operator-valued kernel K(X,Z) -> (N,M,Dy,Dy).

KRRTangent(kernel[, ridge_init, jitter, device])

KRR for vector fields on a manifold, using a specialized tangent kernel.
class dymad.modules.krr.KRRBase(kernel, ridge_init=0, jitter=1e-10, device=None)

Bases: Module

Base class for Kernel Ridge Regression, in particular:

  • Multi-output single scalar kernel (the most common case)

  • Multi-output multiple scalar kernel (i.e., one kernel per output)

  • True operator-valued kernel (i.e., matrix-valued)

Subclasses must implement:

  • _ensure_solved(self)

  • _predict_from_solution(self, Xnew) -> (M, Dy)

fit()

Precompute the linear solve, which can be backprop’d.

Return type:

Tensor

forward(Xnew)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

property ridge: Tensor
set_train_data(X, Y)
Return type:

None

class dymad.modules.krr.KRRMultiOutputIndep(kernel, ridge_init=0, jitter=1e-10, device=None)

Bases: KRRBase

Scalar KRR for multiple outputs, and one kernel per output

  • A ModuleList of Dy scalar kernels (one per output).

  • Dy independent NxN Choleskys; Dy ridges (vector).

class dymad.modules.krr.KRRMultiOutputShared(kernel, ridge_init=0, jitter=1e-10, device=None)

Bases: KRRBase

Scalar KRR for multiple outputs but one single kernel

  • One NxN Cholesky; solve Dy outputs together. One lambda (scalar) by default.

class dymad.modules.krr.KRROperatorValued(kernel, ridge_init=0, jitter=1e-10, device=None)

Bases: KRRBase

Operator-valued kernel K(X,Z) -> (N,M,Dy,Dy).

Solves (Kxx + lambda I) vec(alpha) = vec(Y), using a single (N*Dy)x(N*Dy) Cholesky.

class dymad.modules.krr.KRRTangent(kernel, ridge_init=0, jitter=1e-10, device=None)

Bases: KRRBase

KRR for vector fields on a manifold, using a specialized tangent kernel.

The formulation is based on Geometrically constraint Multivariate KRR (GMKRR) from

Huang, He, Harlim, Li, ICLR2025

Solves (Kxx + lambda I) vec(alpha) = vec(Y), but Kxx is given in a factorized form, so effectively we solve a smaller system in intrinsic dimension d << Dy.

(kxx + lambda I) vec(alpha) = vec(T^T * Y)

set_manifold(manifold)
Return type:

None