SharedIndependentKernel

class pmrf.stats.covariance_kernels.SharedIndependentKernel(base_kernel: AbstractCovarianceKernel, num_shared_axes: int = 1, *, name: str | None = None, metadata: Any = None)

Bases: AbstractCovarianceKernel

Shares a base kernel across trailing independent batch axes (e.g. the real and imaginary parts), so that every element on those axes uses the same hyperparameters.

The shared axes are appended to the base kernel’s output as size-1 axes rather than materialized. A base output of shape batch_shape becomes (*batch_shape, 1, ..., 1), which broadcasts to any size along the shared axes, so a batched base kernel (e.g. AutoCrossKernel) keeps routing along its own leading axes.

Parameters:

  • base_kernel (AbstractCovarianceKernel) – The underlying kernel whose parameters are shared.

  • num_shared_axes (int, default=1) – The number of trailing batch axes the base kernel is shared across.

__call__(x1, x2, key=None)

Evaluate the kernel between two points.

Parameters:

  • x1 (jnp.ndarray) – First input point.

  • x2 (jnp.ndarray) – Second input point.

  • key (jax.random.PRNGKey, optional) – Random key for stochastic kernels.

Returns:

Kernel covariance scalar.

Return type:

jnp.ndarray

base_kernel: AbstractCovarianceKernel

The base kernel.

num_shared_axes: int = 1

The number of trailing shared axes.