Kernels in Jax.
Project description
This project has now been incorporated into GPJax.
Kernels in Jax.
Introduction
JaxKern is Python library for working with kernel functions in JAX. We currently support the following kernels:
- Stationary
- Radial basis function (Squared exponential)
- Matérn
- Powered exponential
- Rational quadratic
- White noise
- Periodic
- Non-stationary
- Linear
- Polynomial
- Non-Euclidean
- Graph kernels
In addition to this, we implement kernel approximations using the Random Fourier feature approach.
Example
The following code snippet demonstrates how the first order Matérn kernel can be computed and, subsequently, approximated using random Fourier features.
import jaxkern as jk
import jax.numpy as jnp
import jax.random as jr
key = jr.PRNGKey(123)
# Define the points on which we'll evaluate the kernel
X = jr.uniform(key, shape = (10, 1), minval=-3., maxval=3.)
Y = jr.uniform(key, shape = (20, 1), minval=-3., maxval=3.)
# Instantiate the kernel and its parameters
kernel = jk.Matern32()
params = kernel.init_params(key)
# Compute the 10x10 Gram matrix
Kxx = kernel.gram(params, X)
# Compute the 10x20 cross-covariance matrix
Kxy = kernel.cross_covariance(params, X, Y)
# Build a RFF approximation
approx = RFF(kernel, num_basis_fns = 5)
rff_params = approx.init_params(key)
# Build an approximation to the Gram matrix
Qff = approx.gram(rff_params, X)
Code Structure
All kernels are supplied with a gram
and cross_covariance
method. When computing a Gram matrix, there is often some structure in the data (e.g., Markov) that can be exploited to yield a sparse matrix. To instruct JAX how to operate on this, the return type of gram
is a Linear Operator from JaxLinOp.
Within GPJax, all kernel computations are handled using JaxKern.
Documentation
A full set of documentation is a work in progress. However, many of the details in JaxKern can be found in the GPJax documentation.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Hashes for jaxkern-nightly-0.0.5.dev20240819.tar.gz
Algorithm | Hash digest | |
---|---|---|
SHA256 | c3f6d5e590610e5ff7a4fed061debe8cd966983ac7df38b5c0aa3a0bed14a66b |
|
MD5 | cdfa7d76d4425a80b032841e82c565fb |
|
BLAKE2b-256 | 560512262e74041251f8ef4ec0c4c42728844eabe4f2eab75f7d287af0250a66 |
Hashes for jaxkern_nightly-0.0.5.dev20240819-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | 396aa34fd3ed14d4a38ec19c70e382b3b95c89086435403dbd5aac632a7f3c56 |
|
MD5 | 3d47e031da86fe750468fc555e139823 |
|
BLAKE2b-256 | 743d73f2b66e82a6a4428c1d5f3173a6958595edbcc38ae256faa35a8ec9cb01 |