aterenin/sparsegpax

An ultra-lightweight JAX implementation of sparse Gaussian processes via pathwise sampling.

SparseGPAX

┌───────────────────┐
│⢰⢦⣀⣀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡆⠀⠀⢀⣴⣶⢄⠀⠀⠀⠀⠀⠀⠀⠀│
│⢨⣿⣷⣻⣵⡀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⣠⣾⣿⣿⣿⡆⠀⠀⠀⠀⠀⠀⠀│
│⠨⡿⣿⣿⣿⣧⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⣼⣿⣿⠿⣿⣾⣿⡀⠀⠀⠀⠀⠀⠀│
│⠸⠋⠑⢿⣿⣿⡆⠀⠀⠀⠀⠀⠀⠀⠀⣿⣿⣿⡟⠉⠺⣻⣿⣧⠀⠀⠀⠀⠀⠀│
│⠀⠀⠀⠀⠹⣿⣿⡀⠀⠀⠀⠀⠀⠀⣼⣿⣿⠝⠀⠀⠀⠈⢿⣿⡆⠀⠀⠀⠀⠀│
│⠉⠉⠉⠉⠉⢻⣿⣯⠉⠉⠉⠉⠉⣹⣿⣿⠏⠉⠉⠉⠉⠉⠉⣿⣿⣍⡉⠉⠉⠉│
│⠀⠀⠀⠀⠀⠀⢿⣿⡆⠀⠀⢀⣮⣿⣿⡏⠀⠀⠀⠀⠀⠀⠀⠸⣿⣏⢟⡄⡜⠀│
│⠀⠀⠀⠀⠀⠀⠈⢿⣿⣦⣶⣫⣿⣿⡏⡇⠀⠀⠀⠀⠀⠀⠀⠀⠹⣿⣯⣫⣦⠂│
│⠀⠀⠀⠀⠀⠀⠀⠘⢿⣿⣿⣿⡿⡟⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠹⣿⣿⣿⡄│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠙⠾⠿⠏⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠹⢿⠿⠀│
│⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⡇⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠱⠀│
└───────────────────┘

An ultra-lightweight JAX implementation of sparse Gaussian processes.

It supports models of the form

(f | u)(.) = f(.) + K_{(.)z} (K_{zz} + \Lambda)^{-1} (\mu - f(z) - \epsilon)

trained using doubly stochastic sparse variational inference via pathwise sampling.

sparsegpax's Issues

Use of tf2jax Operators?

It not really an issue, its more a question. I notice in the code that you make wide use of the following tf2jax operations in order to utilise tensorflow operators e.g.

tf2jax.linalg.LinearOperatorLowerTriangular(cholesky).solvevec(tf2jax.linalg.LinearOperatorLowerTriangular(cholesky).solvevec(residual), adjoint=True)

Is there a particular reason e.g speed, why one would want to do this over just using the native jax equivalent operators?