Skip to content

Base

gpjax.kernels.base

SumKernel = partial(CombinationKernel, operator=jnp.sum) module-attribute
ProductKernel = partial(CombinationKernel, operator=jnp.prod) module-attribute
AbstractKernel dataclass

Bases: Module

Base kernel class.

compute_engine: AbstractKernelComputation = static_field(DenseKernelComputation()) class-attribute instance-attribute
active_dims: Optional[List[int]] = static_field(None) class-attribute instance-attribute
name: str = static_field('AbstractKernel') class-attribute instance-attribute
ndims property
spectral_density: Optional[tfd.Distribution] property
__init_subclass__(mutable: bool = False)
replace(**kwargs: Any) -> Self

Replace the values of the fields of the object.

Parameters:

Name Type Description Default
**kwargs Any

keyword arguments to replace the fields of the object.

 {}
Returns
Module: with the fields replaced.
replace_meta(**kwargs: Any) -> Self

Replace the metadata of the fields.

Parameters:

Name Type Description Default
**kwargs Any

keyword arguments to replace the metadata of the fields of the object.

 {}
Returns
Module: with the metadata of the fields replaced.
update_meta(**kwargs: Any) -> Self

Update the metadata of the fields. The metadata must already exist.

Parameters:

Name Type Description Default
**kwargs Any

keyword arguments to replace the fields of the object.

 {}
Returns
Module: with the fields replaced.
replace_trainable(**kwargs: Dict[str, bool]) -> Self

Replace the trainability status of local nodes of the Module.

replace_bijector(**kwargs: Dict[str, tfb.Bijector]) -> Self

Replace the bijectors of local nodes of the Module.

constrain() -> Self

Transform model parameters to the constrained space according to their defined bijectors.

Returns
Module: transformed to the constrained space.
unconstrain() -> Self

Transform model parameters to the unconstrained space according to their defined bijectors.

Returns
Module: transformed to the unconstrained space.
stop_gradient() -> Self

Stop gradients flowing through the Module.

Returns
Module: with gradients stopped.
trainables() -> Self
cross_covariance(x: Num[Array, 'N D'], y: Num[Array, 'M D'])
gram(x: Num[Array, 'N D'])
slice_input(x: Float[Array, '... D']) -> Float[Array, '... Q']

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:

Name Type Description Default
x Float[Array, '... D']

The matrix or vector that is to be sliced.

 required
Returns
Float[Array, "... Q"]: A sliced form of the input matrix.
__call__(x: Num[Array, ' D'], y: Num[Array, ' D']) -> ScalarFloat abstractmethod

Evaluate the kernel on a pair of inputs.

Parameters:

Name Type Description Default
x Num[Array, ' D']

The left hand input of the kernel function.

 required
y Num[Array, ' D']

The right hand input of the kernel function.

 required
Returns
ScalarFloat: The evaluated kernel function at the supplied inputs.
__add__(other: Union[AbstractKernel, ScalarFloat]) -> AbstractKernel

Add two kernels together. Args: other (AbstractKernel): The kernel to be added to the current kernel.

Returns
AbstractKernel: A new kernel that is the sum of the two kernels.
__radd__(other: Union[AbstractKernel, ScalarFloat]) -> AbstractKernel

Add two kernels together. Args: other (AbstractKernel): The kernel to be added to the current kernel.

Returns
AbstractKernel: A new kernel that is the sum of the two kernels.
__mul__(other: Union[AbstractKernel, ScalarFloat]) -> AbstractKernel

Multiply two kernels together.

Parameters:

Name Type Description Default
other AbstractKernel

The kernel to be multiplied with the current kernel.

 required
Returns
AbstractKernel: A new kernel that is the product of the two kernels.
Constant dataclass

Bases: AbstractKernel

A constant kernel. This kernel evaluates to a constant for all inputs. The scalar value itself can be treated as a model hyperparameter and learned during training.

compute_engine: AbstractKernelComputation = static_field(DenseKernelComputation()) class-attribute instance-attribute
active_dims: Optional[List[int]] = static_field(None) class-attribute instance-attribute
name: str = static_field('AbstractKernel') class-attribute instance-attribute
ndims property
spectral_density: Optional[tfd.Distribution] property
constant: ScalarFloat = param_field(jnp.array(0.0)) class-attribute instance-attribute
__init_subclass__(mutable: bool = False)
replace(**kwargs: Any) -> Self

Replace the values of the fields of the object.

Parameters:

Name Type Description Default
**kwargs Any

keyword arguments to replace the fields of the object.

 {}
Returns
Module: with the fields replaced.
replace_meta(**kwargs: Any) -> Self

Replace the metadata of the fields.

Parameters:

Name Type Description Default
**kwargs Any

keyword arguments to replace the metadata of the fields of the object.

 {}
Returns
Module: with the metadata of the fields replaced.
update_meta(**kwargs: Any) -> Self

Update the metadata of the fields. The metadata must already exist.

Parameters:

Name Type Description Default
**kwargs Any

keyword arguments to replace the fields of the object.

 {}
Returns
Module: with the fields replaced.
replace_trainable(**kwargs: Dict[str, bool]) -> Self

Replace the trainability status of local nodes of the Module.

replace_bijector(**kwargs: Dict[str, tfb.Bijector]) -> Self

Replace the bijectors of local nodes of the Module.

constrain() -> Self

Transform model parameters to the constrained space according to their defined bijectors.

Returns
Module: transformed to the constrained space.
unconstrain() -> Self

Transform model parameters to the unconstrained space according to their defined bijectors.

Returns
Module: transformed to the unconstrained space.
stop_gradient() -> Self

Stop gradients flowing through the Module.

Returns
Module: with gradients stopped.
trainables() -> Self
cross_covariance(x: Num[Array, 'N D'], y: Num[Array, 'M D'])
gram(x: Num[Array, 'N D'])
slice_input(x: Float[Array, '... D']) -> Float[Array, '... Q']

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:

Name Type Description Default
x Float[Array, '... D']

The matrix or vector that is to be sliced.

 required
Returns
Float[Array, "... Q"]: A sliced form of the input matrix.
__add__(other: Union[AbstractKernel, ScalarFloat]) -> AbstractKernel

Add two kernels together. Args: other (AbstractKernel): The kernel to be added to the current kernel.

Returns
AbstractKernel: A new kernel that is the sum of the two kernels.
__radd__(other: Union[AbstractKernel, ScalarFloat]) -> AbstractKernel

Add two kernels together. Args: other (AbstractKernel): The kernel to be added to the current kernel.

Returns
AbstractKernel: A new kernel that is the sum of the two kernels.
__mul__(other: Union[AbstractKernel, ScalarFloat]) -> AbstractKernel

Multiply two kernels together.

Parameters:

Name Type Description Default
other AbstractKernel

The kernel to be multiplied with the current kernel.

 required
Returns
AbstractKernel: A new kernel that is the product of the two kernels.
__call__(x: Float[Array, ' D'], y: Float[Array, ' D']) -> ScalarFloat

Evaluate the kernel on a pair of inputs.

Parameters:

Name Type Description Default
x Float[Array, ' D']

The left hand input of the kernel function.

 required
y Float[Array, ' D']

The right hand input of the kernel function.

 required
Returns
ScalarFloat: The evaluated kernel function at the supplied inputs.
CombinationKernel dataclass

Bases: AbstractKernel

A base class for products or sums of MeanFunctions.

compute_engine: AbstractKernelComputation = static_field(DenseKernelComputation()) class-attribute instance-attribute
active_dims: Optional[List[int]] = static_field(None) class-attribute instance-attribute
name: str = static_field('AbstractKernel') class-attribute instance-attribute
ndims property
spectral_density: Optional[tfd.Distribution] property
kernels: List[AbstractKernel] = None class-attribute instance-attribute
operator: Callable = static_field(None) class-attribute instance-attribute
__init_subclass__(mutable: bool = False)
replace(**kwargs: Any) -> Self

Replace the values of the fields of the object.

Parameters:

Name Type Description Default
**kwargs Any

keyword arguments to replace the fields of the object.

 {}
Returns
Module: with the fields replaced.
replace_meta(**kwargs: Any) -> Self

Replace the metadata of the fields.

Parameters:

Name Type Description Default
**kwargs Any

keyword arguments to replace the metadata of the fields of the object.

 {}
Returns
Module: with the metadata of the fields replaced.
update_meta(**kwargs: Any) -> Self

Update the metadata of the fields. The metadata must already exist.

Parameters:

Name Type Description Default
**kwargs Any

keyword arguments to replace the fields of the object.

 {}
Returns
Module: with the fields replaced.
replace_trainable(**kwargs: Dict[str, bool]) -> Self

Replace the trainability status of local nodes of the Module.

replace_bijector(**kwargs: Dict[str, tfb.Bijector]) -> Self

Replace the bijectors of local nodes of the Module.

constrain() -> Self

Transform model parameters to the constrained space according to their defined bijectors.

Returns
Module: transformed to the constrained space.
unconstrain() -> Self

Transform model parameters to the unconstrained space according to their defined bijectors.

Returns
Module: transformed to the unconstrained space.
stop_gradient() -> Self

Stop gradients flowing through the Module.

Returns
Module: with gradients stopped.
trainables() -> Self
cross_covariance(x: Num[Array, 'N D'], y: Num[Array, 'M D'])
gram(x: Num[Array, 'N D'])
slice_input(x: Float[Array, '... D']) -> Float[Array, '... Q']

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:

Name Type Description Default
x Float[Array, '... D']

The matrix or vector that is to be sliced.

 required
Returns
Float[Array, "... Q"]: A sliced form of the input matrix.
__add__(other: Union[AbstractKernel, ScalarFloat]) -> AbstractKernel

Add two kernels together. Args: other (AbstractKernel): The kernel to be added to the current kernel.

Returns
AbstractKernel: A new kernel that is the sum of the two kernels.
__radd__(other: Union[AbstractKernel, ScalarFloat]) -> AbstractKernel

Add two kernels together. Args: other (AbstractKernel): The kernel to be added to the current kernel.

Returns
AbstractKernel: A new kernel that is the sum of the two kernels.
__mul__(other: Union[AbstractKernel, ScalarFloat]) -> AbstractKernel

Multiply two kernels together.

Parameters:

Name Type Description Default
other AbstractKernel

The kernel to be multiplied with the current kernel.

 required
Returns
AbstractKernel: A new kernel that is the product of the two kernels.
__post_init__()
__call__(x: Float[Array, ' D'], y: Float[Array, ' D']) -> ScalarFloat

Evaluate the kernel on a pair of inputs.

Parameters:

Name Type Description Default
x Float[Array, ' D']

The left hand input of the kernel function.

 required
y Float[Array, ' D']

The right hand input of the kernel function.

 required
Returns
ScalarFloat: The evaluated kernel function at the supplied inputs.