Skip to main content

Kernels in Jax.

Project description


This project has now been incorporated into GPJax.

JaxKern's logo

Kernels in Jax.

codecov CircleCI Documentation Status PyPI version Downloads Slack Invite

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

jaxkern-nightly-0.0.5.dev20240930.tar.gz (33.8 kB view details)

Uploaded Source

Built Distribution

File details

Details for the file jaxkern-nightly-0.0.5.dev20240930.tar.gz.

File metadata

File hashes

Hashes for jaxkern-nightly-0.0.5.dev20240930.tar.gz
Algorithm Hash digest
SHA256 91bf380c68e2c8290b43fc9b6ed996e6b5b9c92c73e6526db4bf1c2899f32a24
MD5 440f57db6f06f5a91f88e3e9ba31f51f
BLAKE2b-256 4edc1d59927d2b5a2c44a87808ae17ab4945ac01636b6b40fa7d4be9134a0c0d

See more details on using hashes here.

File details

Details for the file jaxkern_nightly-0.0.5.dev20240930-py3-none-any.whl.

File metadata

File hashes

Hashes for jaxkern_nightly-0.0.5.dev20240930-py3-none-any.whl
Algorithm Hash digest
SHA256 49625744a6e6635278f6d3d3a3df238ecb938f452aa538179c32d6f8f781a5db
MD5 a048b6de07f48ce50b6611eeec8f4ecc
BLAKE2b-256 7c6ecaa21262fd99b74e50c1da78b77cc6e72a36da0cbfd4bb5fc1957eb3d879

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page