solver

Classes

ParallelStateSpaceSolver

A solver that uses jax.lax.associative_scan to implement

Module Contents

class smolgp.solvers.parallel.solver.ParallelStateSpaceSolver(kernel: smolgp.kernels.base.StateSpaceModel, X: tinygp.helpers.JAXArray, noise: tinygp.helpers.JAXArray)

Bases: smolgp.solvers.solver.StateSpaceSolver

A solver that uses jax.lax.associative_scan to implement parallel Kalman filtering and RTS smoothing.

Inherits from StateSpaceSolver and overrides the Kalman and RTS methods to use the parallel implementations. Methods which do not benefit from associative scans are inherited from StateSpaceSolver.

Kalman(y, return_v_S=True) → Any

Wrapper for Kalman filter used with this solver

log_probability(y) → tinygp.helpers.JAXArray

The marginal log likelihood, reduced from this solver’s own filter.

Overrides StateSpaceSolver.log_probability() deliberately, as that is an optimized sequential scan. The generic path to reuse the Kalman filtered v and S is better here, as those are determined via associative scan, hence the likelihood stays log-depth.

RTS(kalman_results) → Any

Wrapper for RTS smoother used with this solver

condition(y, return_v_S=False) → tinygp.helpers.JAXArray

Compute the Kalman predicted, filtered, and RTS smoothed means and covariances at each of the input coordinates