cross_gram

pmrf.stats.cross_gram(kernel: Callable[[Array, Array], Array], x1: Array, x2: Array) → Array

Build the cross-Gram matrix of a kernel between two sets of inputs.

Entry [..., i, j] is kernel(x1[i], x2[j]), evaluated with a double jax.vmap(). Batching is preserved exactly as for gram(): if the kernel returns an array of shape batch_shape for a single pair of points, the result has shape (*batch_shape, N1, N2), so each event block gets its own (N1, N2) matrix.

Parameters:

• kernel (Callable[[jnp.ndarray, jnp.ndarray], jnp.ndarray]) – The covariance kernel, as for gram().

• x1 (jnp.ndarray) – The row input points, of shape (N1,) or (N1, d).

• x2 (jnp.ndarray) – The column input points, of shape (N2,) or (N2, d).

Returns:

The cross-Gram matrix, of shape (*batch_shape, N1, N2).

Return type:

jnp.ndarray