gram

pmrf.covariance_kernels.gram(kernel: Callable[[Array, Array], Array], x: Array, *, jitter: float = 0.0) → Array

Build the Gram (covariance) matrix of a kernel evaluated at a set of inputs.

The kernel is evaluated for every pair of input points using a double jax.vmap(), producing an (N, N) matrix. Batching is preserved exactly: if the kernel returns an array of shape batch_shape for a single pair of points (for example a kernel whose parameters have shape (D,)), the result has shape (*batch_shape, N, N).

Parameters:

• kernel (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]) – The covariance kernel. Accepts two input points of identical shape and returns an array broadcastable to the kernel’s batch shape. Can be a function or a callable PyTree. See pmrf.covariance_kernels for built-in covariance kernels.

• x (jnp.ndarray) – The input points. An array of shape (N,) is treated as N one-dimensional features; an array of shape (N, d) is treated as N d-dimensional features.

• jitter (float, default=0.0) – A small scalar added to the diagonal of the matrix for numerical stability. The default of 0.0 returns the raw Gram matrix.

Returns:

The Gram matrix, of shape (*batch_shape, N, N), where batch_shape is the shape returned by the kernel for a single pair.

Return type:

jnp.ndarray