JaxKern

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.

Metadata

Release files for jaxkern-nightly 0.0.5.dev20230222

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for jaxkern-nightly 0.0.5.dev20230222
File	Size	Uploaded
jaxkern-nightly-0.0.5.dev20230222.tar.gz	33.5 kB	Feb 22, 2023	Details

Built distribution (wheel)

Table of built distributions (wheels) for jaxkern-nightly 0.0.5.dev20230222
File	Interpreter	ABI	Platform	Reset
jaxkern_nightly-0.0.5.dev20230222-py3-none-any.whl 35.7 kB Feb 22, 2023	Python 3	none	any	Details

Total release size: 69.3 kB

Release files / jaxkern-nightly-0.0.5.dev20230222.tar.gz

Download URL	jaxkern-nightly-0.0.5.dev20230222.tar.gz
Size	33.5 kB
Tags	Source
SHA-256 checksum How to use checksums	`3935fec497dbfa61d03eb2c108c1cb404c9e7f4d4a1da8dda1f8eed769fea43f`
BLAKE2b-256 checksum How to use checksums	`b468210b0ce03a46f1adbc21a699c14a7355edaf88696d95ab596abf9fb9305e`
Upload date	Feb 22, 2023
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/4.0.2 CPython/3.9.0`

Release files / jaxkern_nightly-0.0.5.dev20230222-py3-none-any.whl

Download URL	jaxkern_nightly-0.0.5.dev20230222-py3-none-any.whl
Size	35.7 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`466180602f7e16dd1a6e1ecd91f1f419c8c5acce388472250f515c7baab2b954`
BLAKE2b-256 checksum How to use checksums	`6fed7423b41ec47945376b7573a4ff1c872a1e1c8303a46d5ba8e721bbf38b19`
Upload date	Feb 22, 2023
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/4.0.2 CPython/3.9.0`