MultiOutputKernelComputation

class gpjax.kernels.MultiOutputKernelComputation

Bases: AbstractKernelComputation

Compute engine for multi-output kernels.

Iterates over kernel.components — a sequence of (CoregionalizationMatrix, kernel) pairs — to build structured covariance matrices. Single-component kernels (ICM) retain Kronecker structure; multi-component kernels (LCM) materialise the sum to Dense.

cross_covariance(kernel, x, y)

Override to bypass [N, M] return type annotation for multi-output.

Parameters:

• x (Num[jaxlib._jax.Array, 'N D'] | Num[ndarray, 'N D'])

• y (Num[jaxlib._jax.Array, 'M D'] | Num[ndarray, 'M D'])

Return type:

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

diagonal(kernel, inputs)

For a given kernel, compute the elementwise diagonal of the NxN gram matrix on an input matrix of shape (N, D).

Parameters:

• kernel – the kernel function.

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

Returns:

The computed diagonal variance as a Diagonal linear operator.

Return type:

AbstractLinearOperator

gram(kernel, x)

For a given kernel, compute Gram covariance operator of the kernel function on an input matrix of shape (N, D).

Parameters:

• kernel – the kernel function.

• x (Num[jaxlib._jax.Array, 'N D'] | Num[ndarray, 'N D']) – the inputs to the kernel function of shape (N, D).

Returns:

The Gram covariance of the kernel function as a linear operator.

Return type:

AbstractLinearOperator