obsweatherscale.kernels.neural_kernel#

Classes#

NeuralKernel

Neural Kernel class.

Module Contents#

class NeuralKernel(net: torch.nn.Module, kernel: gpytorch.kernels.Kernel)[source]#

Bases: gpytorch.kernels.Kernel

Neural Kernel class.

This class implements a neural kernel that transforms the input data to a learned feature space using a neural network before applying the kernel function.

TODO: add reference

net#
kernel#
forward(x1: torch.Tensor, x2: torch.Tensor, *params: Any, **kwargs: Any) → torch.Tensor | linear_operator.operators.LinearOperator[source]#

Computes the covariance between \(\mathbf x_1\) and \(\mathbf x_2\). This method should be implemented by all Kernel subclasses.

Parameters:

  • x1 – First set of data (… x N x D).

  • x2 – Second set of data (… x M x D).

  • diag – Should the Kernel compute the whole kernel, or just the diag? If True, it must be the case that x1 == x2. (Default: False.)

  • last_dim_is_batch – If True, treat the last dimension of x1 and x2 as another batch dimension. (Useful for additive structure over the dimensions). (Default: False.)

Returns:

The kernel matrix or vector. The shape depends on the kernel’s evaluation mode:

  • full_covar: … x N x M

  • full_covar with last_dim_is_batch=True: … x K x N x M

  • diag: … x N

  • diag with last_dim_is_batch=True: … x K x N