AbstractKernel

class gpjax.kernels.AbstractKernel(active_dims=None, n_dims=None, compute_engine=<gpjax.kernels.computations.dense.DenseKernelComputation object>)

Bases: _SummaryMixin, Module

Base kernel class.

This class is the base class for all kernels in GPJax. It provides the basic functionality for evaluating a kernel function on a pair of inputs, as well as the ability to combine kernels using addition and multiplication.

The class also provides a method for slicing the input matrix to select the relevant columns for the kernel’s evaluation.

Parameters:

• active_dims (list[int] | slice)

• n_dims (int | None)

• compute_engine (AbstractKernelComputation)

cross_covariance(x, y)

Compute the cross-covariance matrix of the kernel.

Parameters:

• x (Num[jaxlib._jax.Array, 'N D'] | Num[ndarray, 'N D']) – the first input matrix of shape (N, D).

• y (Num[jaxlib._jax.Array, 'M D'] | Num[ndarray, 'M D']) – the second input matrix of shape (M, D).

Returns:

The cross-covariance matrix of the kernel of shape (N, M).

Return type:

Float[jaxlib._jax.Array, ‘N M’] | Float[ndarray, ‘N M’]

diagonal(x)

Compute the diagonal of the gram matrix of the kernel.

Parameters:

x (Num[jaxlib._jax.Array, 'N D'] | Num[ndarray, 'N D']) – the input matrix of shape (N, D).

Returns:

The diagonal of the gram matrix of the kernel of shape (N,).

Return type:

AbstractLinearOperator

gram(x)

Compute the gram matrix of the kernel.

Parameters:

x (Num[jaxlib._jax.Array, 'N D'] | Num[ndarray, 'N D']) – the input matrix of shape (N, D).

Returns:

The gram matrix of the kernel of shape (N, N).

Return type:

AbstractLinearOperator

slice_input(x)

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.

Parameters:

x (Float[jaxlib._jax.Array, '... D'] | Float[ndarray, '... D']) – the matrix or vector that is to be sliced.

Returns:

The sliced form of the input matrix.

Return type:

Float[jaxlib._jax.Array, ‘… Q’] | Float[ndarray, ‘… Q’]