Quick Start

Here is a complete code example demonstrating how to use KernelSteinDiscrepancyJax:

import jax
from jax import numpy as jnp
from jax.scipy.stats import multivariate_normal

from ksd_metric.kernel import KernelJax
from ksd_metric.stein import KernelSteinDiscrepancyJax
from ksd_metric.target import TargetDistributionJax
from ksd_metric.utils import JaxKernelFunction

# Example usage of KernelSteinDiscrepancyJax
key = jax.random.PRNGKey(42)
dim = 2
mean = jnp.zeros(dim)
cov = jnp.eye(dim)
N = 10_000
x = jax.random.multivariate_normal(key, mean, cov, shape=(N,))

# Define the target distribution and kernel
log_target_pdf = lambda x: multivariate_normal.logpdf(x, mean=mean, cov=cov)
target = TargetDistributionJax(log_target_pdf=log_target_pdf)
kernel = KernelJax(lambda x, y: JaxKernelFunction.imq(x, y, jnp.eye(dim)))
ksd = KernelSteinDiscrepancyJax(target=target, kernel=kernel)

# Calculate kernel Stein discrepancy
res = ksd.kernel_stein_discrepancy(x)
print(res)