----
This project has now been incorporated into [GPJax](https://github.com/JaxGaussianProcesses/GPJax).
----
<p align="center">
<img width="700" height="300" src="https://raw.githubusercontent.com/JaxGaussianProcesses/JaxKern/main/docs/_static/logo/logo.png" alt="JaxKern's logo">
</p>
<h2 align='center'>Kernels in Jax.</h2>
[](https://codecov.io/gh/JaxGaussianProcesses/JaxKern)
[](https://dl.circleci.com/status-badge/redirect/gh/JaxGaussianProcesses/JaxKern/tree/main)
[](https://gpjax.readthedocs.io/en/latest/?badge=latest)
[](https://badge.fury.io/py/jaxkern)
[](https://pepy.tech/project/jaxkern)
[](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw)
## 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](https://people.eecs.berkeley.edu/~brecht/papers/07.rah.rec.nips.pdf) approach.
## Example
The following code snippet demonstrates how the first order Matérn kernel can be computed and, subsequently, approximated using random Fourier features.
```python
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](https://github.com/JaxGaussianProcesses/JaxLinOp).
Within [GPJax](https://github.com/JaxGaussianProcesses/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](https://gpjax.readthedocs.io/en/latest/).
Raw data
{
"_id": null,
"home_page": null,
"name": "jaxkern-nightly",
"maintainer": null,
"docs_url": null,
"requires_python": null,
"maintainer_email": null,
"keywords": "gaussian-processes jax machine-learning bayesian",
"author": "Daniel Dodd and Thomas Pinder",
"author_email": "tompinder@live.co.uk",
"download_url": "https://files.pythonhosted.org/packages/7c/2f/80d27e77f6b937a34d239436ad460cd50c230642ba88cc1fb84e7ff3c74a/jaxkern-nightly-0.0.5.dev20240504.tar.gz",
"platform": null,
"description": "----\nThis project has now been incorporated into [GPJax](https://github.com/JaxGaussianProcesses/GPJax).\n----\n\n<p align=\"center\">\n<img width=\"700\" height=\"300\" src=\"https://raw.githubusercontent.com/JaxGaussianProcesses/JaxKern/main/docs/_static/logo/logo.png\" alt=\"JaxKern's logo\">\n</p>\n<h2 align='center'>Kernels in Jax.</h2>\n\n[](https://codecov.io/gh/JaxGaussianProcesses/JaxKern)\n[](https://dl.circleci.com/status-badge/redirect/gh/JaxGaussianProcesses/JaxKern/tree/main)\n[](https://gpjax.readthedocs.io/en/latest/?badge=latest)\n[](https://badge.fury.io/py/jaxkern)\n[](https://pepy.tech/project/jaxkern)\n[](https://join.slack.com/t/gpjax/shared_invite/zt-1da57pmjn-rdBCVg9kApirEEn2E5Q2Zw)\n\n## Introduction\n\nJaxKern is Python library for working with kernel functions in JAX. We currently support the following kernels:\n* Stationary\n * Radial basis function (Squared exponential)\n * Mat\u00e9rn\n * Powered exponential\n * Rational quadratic\n * White noise\n * Periodic\n* Non-stationary\n * Linear \n * Polynomial\n* Non-Euclidean\n * Graph kernels\n\nIn addition to this, we implement kernel approximations using the [Random Fourier feature](https://people.eecs.berkeley.edu/~brecht/papers/07.rah.rec.nips.pdf) approach.\n\n## Example\n\nThe following code snippet demonstrates how the first order Mat\u00e9rn kernel can be computed and, subsequently, approximated using random Fourier features.\n```python\nimport jaxkern as jk\nimport jax.numpy as jnp\nimport jax.random as jr\nkey = jr.PRNGKey(123)\n\n# Define the points on which we'll evaluate the kernel\nX = jr.uniform(key, shape = (10, 1), minval=-3., maxval=3.)\nY = jr.uniform(key, shape = (20, 1), minval=-3., maxval=3.)\n\n# Instantiate the kernel and its parameters\nkernel = jk.Matern32()\nparams = kernel.init_params(key)\n\n# Compute the 10x10 Gram matrix\nKxx = kernel.gram(params, X)\n\n# Compute the 10x20 cross-covariance matrix\nKxy = kernel.cross_covariance(params, X, Y)\n\n# Build a RFF approximation\napprox = RFF(kernel, num_basis_fns = 5)\nrff_params = approx.init_params(key)\n\n# Build an approximation to the Gram matrix\nQff = approx.gram(rff_params, X)\n```\n\n## Code Structure\n\nAll 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](https://github.com/JaxGaussianProcesses/JaxLinOp). \n\nWithin [GPJax](https://github.com/JaxGaussianProcesses/GPJax), all kernel computations are handled using JaxKern.\n\n## Documentation\n\nA full set of documentation is a work in progress. However, many of the details in JaxKern can be found in the [GPJax documentation](https://gpjax.readthedocs.io/en/latest/).\n\n\n",
"bugtrack_url": null,
"license": "LICENSE",
"summary": "Kernels in Jax.",
"version": "0.0.5.dev20240504",
"project_urls": {
"Documentation": "https://JaxKern.readthedocs.io/en/latest/",
"Source": "https://github.com/JaxGaussianProcesses/JaxKern"
},
"split_keywords": [
"gaussian-processes",
"jax",
"machine-learning",
"bayesian"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "8828f31ff3a1bce9a46eb07e636c7a0f64979904c157ac6b0abe7349d390d1ec",
"md5": "7e7d1960d1d1a9b38f09277c91afa8ba",
"sha256": "55ec5250556717fca8c3c1afb5096aa50c59e21dcae4dc321144edcd9d3ffbb8"
},
"downloads": -1,
"filename": "jaxkern_nightly-0.0.5.dev20240504-py3-none-any.whl",
"has_sig": false,
"md5_digest": "7e7d1960d1d1a9b38f09277c91afa8ba",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": null,
"size": 35845,
"upload_time": "2024-05-04T00:04:52",
"upload_time_iso_8601": "2024-05-04T00:04:52.498394Z",
"url": "https://files.pythonhosted.org/packages/88/28/f31ff3a1bce9a46eb07e636c7a0f64979904c157ac6b0abe7349d390d1ec/jaxkern_nightly-0.0.5.dev20240504-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "7c2f80d27e77f6b937a34d239436ad460cd50c230642ba88cc1fb84e7ff3c74a",
"md5": "d4d4046557ebe500173d9f18388e23c2",
"sha256": "d85a22121221825ecfec93fee98235a40dc9f1abde00f2acf592796c7e8f7e11"
},
"downloads": -1,
"filename": "jaxkern-nightly-0.0.5.dev20240504.tar.gz",
"has_sig": false,
"md5_digest": "d4d4046557ebe500173d9f18388e23c2",
"packagetype": "sdist",
"python_version": "source",
"requires_python": null,
"size": 33746,
"upload_time": "2024-05-04T00:04:56",
"upload_time_iso_8601": "2024-05-04T00:04:56.392596Z",
"url": "https://files.pythonhosted.org/packages/7c/2f/80d27e77f6b937a34d239436ad460cd50c230642ba88cc1fb84e7ff3c74a/jaxkern-nightly-0.0.5.dev20240504.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-05-04 00:04:56",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "JaxGaussianProcesses",
"github_project": "JaxKern",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"circle": true,
"lcname": "jaxkern-nightly"
}
JSON
