GPJax

https://docs.jaxgaussianprocesses.com/
Admin
GPJax is a didactic Gaussian process library in JAX, supporting GPU acceleration and just-in-time...
Tokens:244,420
Snippets:1,433
Trust Score:6.3
Update:1 month ago
Show doc for...
Tokens:
Raw
### Setup and imports for multi-output GP example

Source: https://docs.jaxgaussianprocesses.com/_examples/multioutput

Imports required libraries, enables 64-bit precision, and sets up plotting style and color cycle.

```python
from examples.utils import use_mpl_style
from jax import config
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import install_import_hook
import matplotlib as mpl
import matplotlib.pyplot as plt

config.update("jax_enable_x64", True)

with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx

key = jr.key(42)
use_mpl_style()
cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]

```

--------------------------------

### Setup and imports for OAK example

Source: https://docs.jaxgaussianprocesses.com/_examples/oak

Enables Float64 for stable matrix inversions, imports required libraries (JAX, GPJax, matplotlib, etc.), and initializes the random key and plotting style.

```python
# Enable Float64 for more stable matrix inversions.
from jax import config

config.update("jax_enable_x64", True)

from examples.utils import use_mpl_style
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import install_import_hook
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpy as np

with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx
    from gpjax.kernels.additive import (
        OrthogonalAdditiveKernel,
        predict_first_order,
        rank_first_order,
        sobol_indices,
    )

key = jr.key(123)
use_mpl_style()
colours = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
```

--------------------------------

### Import libraries and configure JAX for heteroscedastic GP

Source: https://docs.jaxgaussianprocesses.com/_examples/heteroscedastic_inference

Imports and setup for the heteroscedastic GP example. Enables Float64 for stable linear algebra, applies the example plotting style, and initializes a JAX random key and matplotlib color cycle.

```python
from jax import config
import jax.numpy as jnp
import jax.random as jr
import matplotlib as mpl
import matplotlib.pyplot as plt
import optax as ox

from examples.utils import use_mpl_style
import gpjax as gpx
from gpjax.likelihoods import (
    HeteroscedasticGaussian,
    LogNormalTransform,
    SoftplusTransform,
)
from gpjax.variational_families import (
    HeteroscedasticVariationalFamily,
    VariationalGaussianInit,
)

# Enable Float64 for stable linear algebra.
config.update("jax_enable_x64", True)


use_mpl_style()
key = jr.key(123)
cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
```

--------------------------------

### Import libraries and configure JAX for Numpyro

Source: https://docs.jaxgaussianprocesses.com/_examples/numpyro_integration

Imports and setup for the Numpyro integration example. Enables 64-bit JAX precision, applies the example matplotlib style, and creates a random key split into four subkeys for data generation and inference.

```python
from examples.utils import use_mpl_style
import gpjax as gpx
from jax import config
import jax.numpy as jnp
import jax.random as jr
import matplotlib as mpl
import matplotlib.pyplot as plt
import numpyro
import numpyro.distributions as dist
from numpyro.infer import (
    MCMC,
    NUTS,
    Predictive,
)

config.update("jax_enable_x64", True)

use_mpl_style()
cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]

key = jr.key(123)
keys = jr.split(key, 4)
```

--------------------------------

### Setup and imports for GPJax ocean modelling

Source: https://docs.jaxgaussianprocesses.com/_examples/oceanmodelling

Imports necessary libraries, enables 64-bit JAX, installs beartype hooks for GPJax, sets plotting style, and initializes a random key. This setup is required before running the data loading and modelling code.

```python
from examples.utils import use_mpl_style
from gpjax.kernels.computations import DenseKernelComputation
from jax import (
    config,
    hessian,
)
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import (
    Array,
    Float,
    install_import_hook,
)
from matplotlib import rcParams
import matplotlib.pyplot as plt
import numpyro.distributions as npd
import pandas as pd

config.update("jax_enable_x64", True)


with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx


# set the default style for plotting
use_mpl_style()

key = jr.key(42)

colors = rcParams["axes.prop_cycle"].by_key()["color"]

```

--------------------------------

### Setup and imports for OILMM notebook

Source: https://docs.jaxgaussianprocesses.com/_examples/oilmm

Imports necessary libraries, enables 64-bit precision, and sets up plotting style. This is the initial setup for the notebook.

```python
from pathlib import Path

from examples.utils import use_mpl_style
from jax import config
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import install_import_hook
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd

config.update("jax_enable_x64", True)

with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx

key = jr.key(123)
use_mpl_style()
cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
```

--------------------------------

### Install GPJax with uv

Source: https://docs.jaxgaussianprocesses.com/contributing

Installs project dependencies using uv, as specified in pyproject.toml.

```bash
$ uv sync

```

--------------------------------

### create_oilmm_with_kernels usage example

Source: https://docs.jaxgaussianprocesses.com/api/models/oilmm

Example usage of create_oilmm_with_kernels with RBF and Matern52 kernels for 6 outputs.

```python
import gpjax as gpx
import jax.random as jr
model = gpx.models.create_oilmm_with_kernels(
    latent_kernels=[gpx.kernels.RBF(), gpx.kernels.Matern52()],
    num_outputs=6,
    key=jr.key(42)
)
```

--------------------------------

### Install GPJax 0.14.0

Source: https://docs.jaxgaussianprocesses.com/migration

Install GPJax 0.14.0 using pip or uv. New dependencies equinox>=0.11 and paramax>=0.0.5 are pulled in automatically; Flax is no longer a runtime dependency.

```bash
pip install "gpjax==0.14.0"
# or
uv add "gpjax==0.14.0"

```

--------------------------------

### Install development version of GPJax from source

Source: https://docs.jaxgaussianprocesses.com/installation

Clones the GPJax repository and synchronizes dependencies with uv. The development version may be unstable and contain bugs.

```bash
git clone https://github.com/thomaspinder/GPJax.git
cd GPJax
uv sync
```

--------------------------------

### to_sde usage example

Source: https://docs.jaxgaussianprocesses.com/api/state_space/kernels

Example of converting a Matern32 kernel to its SDE representation using to_sde. The example shows that the resulting SDE class name is 'Matern32SDE'.

```python
>>> import gpjax as gpx
>>> from gpjax.state_space import to_sde
>>> sde = to_sde(gpx.kernels.Matern32(lengthscale=1.0, variance=1.0))
>>> sde.__class__.__name__
'Matern32SDE'
```

--------------------------------

### get_batch usage example

Source: https://docs.jaxgaussianprocesses.com/api/fit

Example usage of get_batch: creates a dataset from 100 points, then batches it with batch_size=16 using a random key.

```python
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr

X = jnp.linspace(0, 1, 100).reshape(-1, 1)
y = jnp.sin(X)
D = gpx.Dataset(X=X, y=y)

from gpjax.fit import get_batch
batch = get_batch(D, batch_size=16, key=jr.key(0))
```

--------------------------------

### Import dependencies and enable Float64

Source: https://docs.jaxgaussianprocesses.com/_examples/backend

Sets up the environment for GPJax, enabling Float64 for stable matrix inversions and installing a type-checking import hook.

```python
import typing as tp

import equinox as eqx
from examples.utils import use_mpl_style
from gpjax.mean_functions import (
    AbstractMeanFunction,
    Constant,
)
from gpjax.parameters import (
    PositiveReal,
    Real,
)
from gpjax.typing import (
    Array,
    ScalarFloat,
)

# Enable Float64 for more stable matrix inversions.
from jax import config
import jax.numpy as jnp
import jax.tree_util as jtu
from jaxtyping import (
    Float,
    Num,
    install_import_hook,
)
import matplotlib as mpl
import matplotlib.pyplot as plt
import paramax
from paramax import AbstractUnwrappable

config.update("jax_enable_x64", True)


with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx


# set the default style for plotting
use_mpl_style()

cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
```

--------------------------------

### fit example

Source: https://docs.jaxgaussianprocesses.com/api/state_space/fit

Example usage of fit. Fits a state-space posterior with Optax gradient descent. Uses Adam optimizer and checks that the history is finite.

```python
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> import optax as ox
>>> import gpjax as gpx
>>> from gpjax.state_space import StateSpacePrior, fit
>>> X = jnp.linspace(0.0, 5.0, 20).reshape(-1, 1)
>>> y = jnp.sin(X)
>>> prior = StateSpacePrior(
...     mean_function=gpx.mean_functions.Zero(),
...     kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0),
... )
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1)
>>> posterior = prior * likelihood
>>> fitted, history = fit(
...     model=posterior,
...     train_data=gpx.Dataset(X=X, y=y),
...     optim=ox.adam(1e-2),
...     num_iters=2,
...     key=jr.key(0),
...     verbose=False,
... )
>>> bool(jnp.all(jnp.isfinite(history)))
True

```

--------------------------------

### create_oilmm_from_data usage example

Source: https://docs.jaxgaussianprocesses.com/api/models/oilmm

Example usage of create_oilmm_from_data with synthetic sine and cosine data for 2 latent GPs.

```python
import gpjax as gpx
import jax.numpy as jnp
import jax.random as jr
X = jnp.linspace(0, 1, 50).reshape(-1, 1)
y = jnp.column_stack([jnp.sin(X), jnp.cos(X)])
data = gpx.Dataset(X=X, y=y)
model = gpx.models.create_oilmm_from_data(
    dataset=data,
    num_latent_gps=2,
    key=jr.key(42)
)
```

--------------------------------

### create_oilmm example: 5 outputs, 2 latent GPs, Matern52 kernel

Source: https://docs.jaxgaussianprocesses.com/api/models/oilmm

Example usage of create_oilmm with 5 outputs, 2 latent GPs, a JAX PRNG key, and a Matern52 kernel.

```python
import gpjax as gpx
import jax.random as jr
model = gpx.models.create_oilmm(
    num_outputs=5,
    num_latent_gps=2,
    key=jr.key(42),
    kernel=gpx.kernels.Matern52()
)
```

--------------------------------

### Example: training a GP with gpx.fit

Source: https://docs.jaxgaussianprocesses.com/api/fit

Example usage of gpx.fit. It creates a Gaussian process posterior with an RBF kernel and constant mean, defines the negative marginal log-likelihood as the objective, and trains the model with an Adam optimiser for 100 iterations. The example enables x64 precision and uses a dataset of 50 points.

```python
>>> import jax
>>> jax.config.update("jax_enable_x64", True)
>>> import jax.numpy as jnp
>>> import optax as ox
>>> import gpjax as gpx
>>>
>>> xtrain = jnp.linspace(0, 1, 50).reshape(-1, 1)
>>> ytrain = jnp.sin(xtrain)
>>> D = gpx.Dataset(X=xtrain, y=ytrain)
>>>
>>> meanf = gpx.mean_functions.Constant()
>>> kernel = gpx.kernels.RBF()
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n)
>>> prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
>>> posterior = prior * likelihood
>>>
>>> nmll = lambda p, d: -gpx.objectives.conjugate_mll(p, d)
>>> trained_model, history = gpx.fit(
...     model=posterior, objective=nmll, train_data=D,
...     optim=ox.adam(0.01), num_iters=100, verbose=False,
... )
```

--------------------------------

### Create synthetic dataset for likelihood example

Source: https://docs.jaxgaussianprocesses.com/_examples/likelihoods_guide

Sets up a sinusoidal dataset with Gaussian noise, enabling Float64 for stability, and plots the observations and latent function.

```python
import jax

# Enable Float64 for more stable matrix inversions.
from jax import config
import jax.numpy as jnp
import jax.random as jr
import matplotlib.pyplot as plt

from examples.utils import use_mpl_style
import gpjax as gpx

config.update("jax_enable_x64", True)


# set the default style for plotting
use_mpl_style()
cols = plt.rcParams["axes.prop_cycle"].by_key()["color"]

key = jr.key(42)

n = 50
x = jnp.sort(jr.uniform(key=key, shape=(n, 1), minval=-3.0, maxval=3.0), axis=0)
xtest = jnp.linspace(-3, 3, 100)[:, None]
f = lambda x: jnp.sin(x)
y = f(x) + 0.1 * jr.normal(key, shape=x.shape)
D = gpx.Dataset(x, y)

fig, ax = plt.subplots()
ax.plot(x, y, "o", label="Observations")
ax.plot(x, f(x), label="Latent function")
ax.legend()
```

--------------------------------

### Progress bar output

Source: https://docs.jaxgaussianprocesses.com/_examples/uncollapsed_vi

Example progress bar output from the fitting loop, showing the initial state before iterations complete.

```text
  0%|          | 0/3000 [00:00<?, ?it/s]
```

--------------------------------

### Custom module before (flax.nnx.Module)

Source: https://docs.jaxgaussianprocesses.com/migration

Old way to define a custom kernel module in GPJax 0.13.x using flax.nnx.Module. This is the before example for the backend migration to equinox.Module.

```python
# Before (0.13.x)
from flax import nnx

class MyKernel(nnx.Module):
    def __init__(self, lengthscale):
        self.lengthscale = gpx.parameters.PositiveReal(lengthscale)

```

--------------------------------

### Enable Float64 and import modules

Source: https://docs.jaxgaussianprocesses.com/_examples/regression

Sets up the environment for the regression example. Enables Float64 for more stable matrix inversions, imports necessary libraries including GPJax, and sets a random key.

```python
# Enable Float64 for more stable matrix inversions.
from examples.utils import (
    use_mpl_style,
)
from jax import config
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import install_import_hook
import matplotlib as mpl
import matplotlib.pyplot as plt

from examples.utils import use_mpl_style, clean_legend

config.update("jax_enable_x64", True)


with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx


key = jr.key(123)

# set the default style for plotting
use_mpl_style()

cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]

```

--------------------------------

### Run MCMC with NUTS sampler

Source: https://docs.jaxgaussianprocesses.com/_examples/numpyro_integration

Sets up a NUTS kernel and MCMC sampler with 750 warmup and 500 samples, then runs it on the data and prints a summary. The comment notes that in practice more samples and multiple chains should be used.

```python
nuts_kernel = NUTS(model)
# In practice, one should run more samples from multiple chains.
mcmc = MCMC(
    nuts_kernel,
    num_warmup=750,
    num_samples=500,
)
mcmc.run(keys[2], x, y)
mcmc.print_summary()

```

--------------------------------

### Check GPJax installation version

Source: https://docs.jaxgaussianprocesses.com/installation

Prints the installed GPJax version to verify the installation.

```python
python -c 'import gpjax; print(gpjax.__version__)'
```

--------------------------------

### Import GPJax and configure environment

Source: https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels

Sets up the environment for the kernel guide: enables Float64 for stable matrix inversions, imports GPJax and plotting utilities, and sets a random key.

```python
# Enable Float64 for more stable matrix inversions.
from examples.utils import use_mpl_style
from gpjax.kernels.base import _val
from gpjax.kernels.computations import DenseKernelComputation
from gpjax.parameters import PositiveReal
from jax import config
from jax.nn import softplus
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import (
    Array,
    Float,
    install_import_hook,
)
import matplotlib.pyplot as plt

config.update("jax_enable_x64", True)


with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx


# set the default style for plotting
use_mpl_style()

key = jr.key(42)

cols = plt.rcParams["axes.prop_cycle"].by_key()["color"]

```

--------------------------------

### Configure JAX and Import GPJax

Source: https://docs.jaxgaussianprocesses.com/_examples/intro_to_kernels

Sets up the environment by enabling Float64 precision for numerical stability, configuring type checking with beartype, and initializing the random key.

```python
# Enable Float64 for more stable matrix inversions.
from examples.utils import use_mpl_style
from gpjax.typing import Array
from jax import config
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import (
    Float,
    install_import_hook,
)
import matplotlib as mpl
import matplotlib.pyplot as plt
import optax as ox
import pandas as pd
from sklearn.preprocessing import StandardScaler

config.update("jax_enable_x64", True)


with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx


key = jr.key(42)

# set the default style for plotting
use_mpl_style()
cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
```

--------------------------------

### variational_expectation example

Source: https://docs.jaxgaussianprocesses.com/api/objectives

Example usage of the variational_expectation function. It constructs the same GP model as the elbo example and computes the expectation of the log-likelihood under the variational distribution.

```python
import gpjax as gpx
import jax.numpy as jnp

xtrain = jnp.linspace(0, 1).reshape(-1, 1)
ytrain = jnp.sin(xtrain)
D = gpx.Dataset(X=xtrain, y=ytrain)

meanf = gpx.mean_functions.Constant()
kernel = gpx.kernels.RBF()
likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n)
prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
posterior = prior * likelihood

z = jnp.linspace(0, 1, 10).reshape(-1, 1)
q = gpx.variational_families.VariationalGaussian(
    posterior=posterior, inducing_inputs=z
)

gpx.objectives.variational_expectation(q, D)
```

--------------------------------

### StateSpaceConjugatePosterior usage example

Source: https://docs.jaxgaussianprocesses.com/api/state_space/gps

Example of constructing a StateSpaceConjugatePosterior by multiplying a StateSpacePrior with a Gaussian likelihood. The example uses a Matern32 kernel and verifies the class name.

```python
>>> import gpjax as gpx
>>> from gpjax.state_space import StateSpacePrior
>>> prior = StateSpacePrior(
...     mean_function=gpx.mean_functions.Zero(),
...     kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0),
... )
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1)
>>> posterior = prior * likelihood
>>> posterior.__class__.__name__
'StateSpaceConjugatePosterior'

```

--------------------------------

### Import libraries and configure JAX for count data regression

Source: https://docs.jaxgaussianprocesses.com/_examples/poisson

Sets up the environment for the notebook: imports GPJax, BlackJax, Equinox, and other dependencies, enables Float64 for stable matrix inversions, sets plotting style, and initializes a random key.

```python
from pathlib import Path

import blackjax
import equinox as eqx
from examples.utils import use_mpl_style
import jax
from jax import config
import jax.numpy as jnp
import jax.random as jr
import jax.tree_util as jtu
from jaxtyping import install_import_hook
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd
import paramax

with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx


# Enable Float64 for more stable matrix inversions.
config.update("jax_enable_x64", True)

# set the default style for plotting
use_mpl_style()
cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]

key = jr.key(42)

```

--------------------------------

### Import libraries and configure JAX for barycentre computation

Source: https://docs.jaxgaussianprocesses.com/_examples/barycentres

Sets up the environment for the notebook: enables JAX float64 for stable matrix inversions, installs the beartype hook for gpjax, and sets the plotting style. Requires the 'examples.utils' module for 'use_mpl_style'.

```python
import typing as tp

from examples.utils import use_mpl_style
import jax

# Enable Float64 for more stable matrix inversions.
from jax import config
import jax.numpy as jnp
import jax.random as jr
import jax.scipy.linalg as jsl
from jaxtyping import install_import_hook
import matplotlib.pyplot as plt
import numpyro.distributions as npd

config.update("jax_enable_x64", True)


with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx


key = jr.key(123)

# set the default style for plotting
use_mpl_style()

cols = plt.rcParams["axes.prop_cycle"].by_key()["color"]

```

--------------------------------

### TruncatedPeriodic usage example

Source: https://docs.jaxgaussianprocesses.com/api/state_space/kernels

Example of creating a TruncatedPeriodic kernel with lengthscale, variance, period, and truncation_order. The example shows that the truncation_order attribute is set to 6.

```python
>>> from gpjax.state_space import TruncatedPeriodic
>>> kernel = TruncatedPeriodic(
...     lengthscale=1.0, variance=1.0, period=1.0, truncation_order=6,
... )
>>> kernel.truncation_order
6
```

--------------------------------

### fit_scipy example

Source: https://docs.jaxgaussianprocesses.com/api/state_space/fit

Example usage of fit_scipy. Fits a state-space posterior with SciPy's L-BFGS-B. The example uses a Matern32 kernel and Gaussian likelihood, and checks that the history is finite.

```python
>>> import jax.numpy as jnp
>>> import gpjax as gpx
>>> from gpjax.state_space import StateSpacePrior, fit_scipy
>>> X = jnp.linspace(0.0, 5.0, 20).reshape(-1, 1)
>>> y = jnp.sin(X)
>>> prior = StateSpacePrior(
...     mean_function=gpx.mean_functions.Zero(),
...     kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0),
... )
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1)
>>> posterior = prior * likelihood
>>> fitted, history = fit_scipy(
...     model=posterior,
...     train_data=gpx.Dataset(X=X, y=y),
...     max_iters=2,
...     verbose=False,
... )
>>> bool(jnp.all(jnp.isfinite(history)))
True

```

--------------------------------

### fit_lbfgs example

Source: https://docs.jaxgaussianprocesses.com/api/state_space/fit

Example usage of fit_lbfgs. Fits a state-space posterior with Optax's L-BFGS. The example uses a Matern32 kernel and Gaussian likelihood, and checks that the fitted model is not None.

```python
>>> import jax.numpy as jnp
>>> import gpjax as gpx
>>> from gpjax.state_space import StateSpacePrior, fit_lbfgs
>>> X = jnp.linspace(0.0, 5.0, 20).reshape(-1, 1)
>>> y = jnp.sin(X)
>>> prior = StateSpacePrior(
...     mean_function=gpx.mean_functions.Zero(),
...     kernel=gpx.kernels.Matern32(lengthscale=1.0, variance=1.0),
... )
>>> likelihood = gpx.likelihoods.Gaussian(num_datapoints=20, obs_stddev=0.1)
>>> posterior = prior * likelihood
>>> fitted, history = fit_lbfgs(
...     model=posterior,
...     train_data=gpx.Dataset(X=X, y=y),
...     max_iters=2,
... )
>>> fitted is not None
True

```

--------------------------------

### Install stable GPJax with conda

Source: https://docs.jaxgaussianprocesses.com/installation

Installs GPJax from the conda-forge channel using conda.

```bash
conda install --channel conda-forge gpjax
```

--------------------------------

### ProductKernel

Source: https://docs.jaxgaussianprocesses.com/api/kernels/base

Documentation for the ProductKernel class, including its constructor and methods for kernel evaluation, cross-covariance, gram matrix, diagonal, input slicing, and arithmetic operations.

```APIDOC
## ProductKernel

### Description
A kernel that is the product of a set of kernels.

### Constructor
```
ProductKernel(
    kernels: list[AbstractKernel],
    compute_engine: AbstractKernelComputation = DenseKernelComputation(),
)
```

### Parameters
- **kernels** (list[AbstractKernel]) - Required - The list of kernels to multiply.
- **compute_engine** (AbstractKernelComputation) - Optional - The computation engine used for kernel operations. Defaults to `DenseKernelComputation()`.

### Methods

#### __call__

##### Description
Evaluate the kernel on a pair of inputs.

##### Signature
```
__call__(
    x: Float[Array, " D"], y: Float[Array, " D"]
) -> ScalarFloat
```

##### Parameters
- **x** (Float[Array, " D"]) - Required - The left hand input of the kernel function.
- **y** (Float[Array, " D"]) - Required - The right hand input of the kernel function.

##### Returns
- **ScalarFloat** - The evaluated kernel function at the supplied inputs.

#### cross_covariance

##### Description
Compute the cross-covariance matrix of the kernel.

##### Signature
```
cross_covariance(
    x: Num[Array, "N D"], y: Num[Array, "M D"]
) -> Float[Array, "N M"]
```

##### Parameters
- **x** (Num[Array, "N D"]) - Required - The first input matrix of shape `(N, D)`.
- **y** (Num[Array, "M D"]) - Required - The second input matrix of shape `(M, D)`.

##### Returns
- **Float[Array, "N M"]** - The cross-covariance matrix of the kernel of shape `(N, M)`.

#### gram

##### Description
Compute the gram matrix of the kernel.

##### Signature
```
gram(x: Num[Array, 'N D']) -> lx.AbstractLinearOperator
```

##### Parameters
- **x** (Num[Array, "N D"]) - Required - The input matrix of shape `(N, D)`.

##### Returns
- **AbstractLinearOperator** - The gram matrix of the kernel of shape `(N, N)`.

#### diagonal

##### Description
Compute the diagonal of the gram matrix of the kernel.

##### Signature
```
diagonal(x: Num[Array, 'N D']) -> lx.AbstractLinearOperator
```

##### Parameters
- **x** (Num[Array, "N D"]) - Required - The input matrix of shape `(N, D)`.

##### Returns
- **AbstractLinearOperator** - The diagonal of the gram matrix of the kernel of shape `(N,)`.

#### slice_input

##### Description
Slice out the relevant columns of the input matrix. Select the relevant columns of the supplied matrix to be used within the kernel's evaluation.

##### Signature
```
slice_input(
    x: Float[Array, "... D"],
) -> Float[Array, "... Q"]
```

##### Parameters
- **x** (Float[Array, "... D"]) - Required - The matrix or vector that is to be sliced.

##### Returns
- **Float[Array, "... Q"]** - The sliced form of the input matrix.

#### __add__

##### Description
Add two kernels together.

##### Signature
```
__add__(
    other: Union[AbstractKernel, ScalarFloat],
) -> AbstractKernel
```

##### Parameters
- **other** (Union[AbstractKernel, ScalarFloat]) - Required - The kernel to be added to the current kernel.

##### Returns
- **AbstractKernel** - A new kernel that is the sum of the two kernels.

#### __mul__

##### Description
Multiply two kernels together.

##### Signature
```
__mul__(
    other: Union[AbstractKernel, ScalarFloat],
) -> AbstractKernel
```

##### Parameters
- **other** (Union[AbstractKernel, ScalarFloat]) - Required - The kernel to be multiplied with the current kernel.

##### Returns
- **AbstractKernel** - A new kernel that is the product of the two kernels.
```

--------------------------------

### elbo example

Source: https://docs.jaxgaussianprocesses.com/api/objectives

Example usage of the elbo function. It constructs a GP model with a Bernoulli likelihood and a variational Gaussian family, then computes the ELBO. The example uses a constant mean function, RBF kernel, and inducing inputs.

```python
import gpjax as gpx
import jax.numpy as jnp

xtrain = jnp.linspace(0, 1).reshape(-1, 1)
ytrain = jnp.sin(xtrain)
D = gpx.Dataset(X=xtrain, y=ytrain)

meanf = gpx.mean_functions.Constant()
kernel = gpx.kernels.RBF()
likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n)
prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
posterior = prior * likelihood

z = jnp.linspace(0, 1, 10).reshape(-1, 1)
q = gpx.variational_families.VariationalGaussian(
    posterior=posterior, inducing_inputs=z
)

gpx.objectives.elbo(q, D)
```

--------------------------------

### Run all tests with uv

Source: https://docs.jaxgaussianprocesses.com/contributing

Runs the full test suite using uv and poe.

```bash
$ uv run poe all-tests

```

--------------------------------

### Install stable GPJax from PyPI

Source: https://docs.jaxgaussianprocesses.com/installation

Installs the latest stable release of GPJax from PyPI using pip.

```bash
pip install gpjax
```

--------------------------------

### vscan usage example

Source: https://docs.jaxgaussianprocesses.com/api/scan

A basic example demonstrating how to use vscan with a simple carry function and an array of inputs.

```python
>>> import jax.numpy as jnp ... def f(carry, x): ... return carry + x, carry + x init = 0 xs = jnp.arange(10) vscan(f, init, xs) # doctest: +SKIP (Array(45, dtype=int32), Array([ 0, 1, 3, 6, 10, 15, 21, 28, 36, 45], dtype=int32))
```

--------------------------------

### GaussianDistribution usage example

Source: https://docs.jaxgaussianprocesses.com/api/distributions

Example of creating a GaussianDistribution with a mean vector and a covariance matrix, then accessing its mean and variance.

```python
import jax.numpy as jnp
import lineax as lx
from gpjax.distributions import GaussianDistribution

mu = jnp.array([0.0, 1.0])
cov = lx.MatrixLinearOperator(jnp.eye(2))
dist = GaussianDistribution(loc=mu, scale=cov)

dist.mean  # doctest: +SKIP
Array([0., 1.], dtype=float32)

dist.variance  # doctest: +SKIP
Array([1., 1.], dtype=float32)

```

--------------------------------

### Sample from GP priors with Matérn and RBF kernels

Source: https://docs.jaxgaussianprocesses.com/_examples/intro_to_kernels

Creates a list of Matérn12, Matérn32, Matérn52, and RBF kernels, then samples 10 functions from each GP prior on a grid of 200 points between -3 and 3. Uses gpx.kernels, gpx.mean_functions.Zero, and gpx.gps.Prior.

```python
kernels = [
    gpx.kernels.Matern12(),
    gpx.kernels.Matern32(),
    gpx.kernels.Matern52(),
    gpx.kernels.RBF(),
]
fig, axes = plt.subplots(ncols=2, nrows=2, figsize=(7, 6), tight_layout=True)

x = jnp.linspace(-3.0, 3.0, num=200).reshape(-1, 1)

meanf = gpx.mean_functions.Zero()

for k, ax in zip(kernels, axes.ravel(), strict=False):
    prior = gpx.gps.Prior(mean_function=meanf, kernel=k)
    rv = prior(x)
    y = rv.sample(key=key, sample_shape=(10,))
    ax.plot(x, y.T, alpha=0.7)
    ax.set_title(type(k).__name__)

```

--------------------------------

### Example usage of log_posterior_density

Source: https://docs.jaxgaussianprocesses.com/api/objectives

Example showing how to compute the log-posterior density for a Bernoulli likelihood with an RBF kernel and constant mean function.

```python
import gpjax as gpx
import jax.numpy as jnp

xtrain = jnp.linspace(0, 1).reshape(-1, 1)
yt = jnp.sin(xtrain)
D = gpx.Dataset(X=xtrain, y=ytrain)

meanf = gpx.mean_functions.Constant()
kernel = gpx.kernels.RBF()
likelihood = gpx.likelihoods.Bernoulli(num_datapoints=D.n)
prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
posterior = prior * likelihood

gpx.objectives.log_posterior_density(posterior, D)
```

--------------------------------

### sample_approx example

Source: https://docs.jaxgaussianprocesses.com/api/state_space/gps

Builds and evaluates 10 approximate samples from a GP prior over the interval [0,1]. Uses an RBF kernel and Zero mean function. The sample function is evaluated at 100 points.

```python
>>> import gpjax as gpx
>>> import jax.numpy as jnp
>>> import jax.random as jr
>>> key = jr.key(123)
>>>
>>> meanf = gpx.mean_functions.Zero()
>>> kernel = gpx.kernels.RBF(n_dims=1)
>>> prior = gpx.gps.Prior(mean_function=meanf, kernel = kernel)
>>>
>>> sample_fn = prior.sample_approx(10, key)
>>> sample_fn(jnp.linspace(0, 1, 100).reshape(-1, 1))
```

--------------------------------

### Import dependencies and enable Float64 for state-space GP

Source: https://docs.jaxgaussianprocesses.com/_examples/state_space_gps

Set up the environment for the state-space GP notebook: enable Float64 for stable matrix factorisations, import JAX, GPJax, and plotting utilities, and create a random key. Also imports the state-space components StateSpacePrior, TruncatedPeriodic, and fit_scipy from gpjax.state_space.

```python
# Enable Float64 for more stable matrix factorisations.
from examples.utils import clean_legend, use_mpl_style
import jax
from jax import config
import jax.numpy as jnp
import jax.random as jr
from jaxtyping import install_import_hook
import matplotlib as mpl
import matplotlib.pyplot as plt
import pandas as pd

config.update("jax_enable_x64", True)


with install_import_hook("gpjax", "beartype.beartype"):
    import gpjax as gpx
    from gpjax.state_space import (
        StateSpacePrior,
        TruncatedPeriodic,
        fit_scipy,
    )


key = jr.key(123)

# set the default style for plotting
use_mpl_style()
cols = mpl.rcParams["axes.prop_cycle"].by_key()["color"]
```

--------------------------------

### fit_lbfgs usage example

Source: https://docs.jaxgaussianprocesses.com/api/fit

Example usage of fit_lbfgs. It sets up a Gaussian process posterior with a constant mean, RBF kernel, and Gaussian likelihood, then defines a negative marginal log-likelihood objective and calls fit_lbfgs to train the model. The example enables x64 precision and uses 20 training points.

```python
import jax
jax.config.update("jax_enable_x64", True)
import gpjax as gpx
import jax.numpy as jnp

xtrain = jnp.linspace(0, 1, 20).reshape(-1, 1)
ytrian = jnp.sin(xtrain)
D = gpx.Dataset(X=xtrain, y=ytrian)

meanf = gpx.mean_functions.Constant()
kernel = gpx.kernels.RBF()
likelihood = gpx.likelihoods.Gaussian(num_datapoints=D.n)
prior = gpx.gps.Prior(mean_function=meanf, kernel=kernel)
posterior = prior * likelihood

nmll = lambda p, d: -gpx.objectives.conjugate_mll(p, d)
trained_model, final_loss = gpx.fit_lbfgs(
    model=posterior, objective=nmll, train_data=D
)
```

--------------------------------

### ProductKernel combination of RBF, Polynomial, and Matern32 kernels

Source: https://docs.jaxgaussianprocesses.com/_examples/constructing_new_kernels

Creates a product kernel from RBF, Polynomial, and Matern32 kernels using gpx.kernels.ProductKernel, then plots the gram matrices of the individual kernels and the product kernel.

```python
k3 = gpx.kernels.Matern32()

prod_k = gpx.kernels.ProductKernel(kernels=[k1, k2, k3])

fig, ax = plt.subplots(ncols=4, figsize=(12, 3))
im0 = ax[0].matshow(k1.gram(x).as_matrix())
im1 = ax[1].matshow(k2.gram(x).as_matrix())
im2 = ax[2].matshow(k3.gram(x).as_matrix())
im3 = ax[3].matshow(prod_k.gram(x).as_matrix())

fig.colorbar(im0, ax=ax[0], fraction=0.05)
fig.colorbar(im1, ax=ax[1], fraction=0.05)
fig.colorbar(im2, ax=ax[2], fraction=0.05)
fig.colorbar(im3, ax=ax[3], fraction=0.05)

```

```output
<matplotlib.colorbar.Colorbar at 0x7f51dcadacd0>

```
Tokens:
Raw