Customized Using Interface

Here is a complete code example demonstrating how to use Customized target distribution and base kernel function.

import jax
from jax import numpy as jnp
from jaxtyping import Array, Float

from ksd_metric.kernel import KernelJax
from ksd_metric.stein import KernelSteinDiscrepancyJax
from ksd_metric.target import TargetDistributionInterface

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,))


class CustomTargetDistribution(TargetDistributionInterface):
    def log_target_pdf(self, x: Float[Array, "num dim"]) -> Float[Array, "num dim"]:
        return -0.5 * x @ x

    def grad_log_target_pdf(
        self, x: Float[Array, "num dim"]
    ) -> Float[Array, "num dim"]:
        return -x


def custom_kernel(
    x: Float[Array, "num"],
    y: Float[Array, "num"],
) -> Float[Array, "num"]:
    dim = len(x)
    diff = x - y
    return (1.0 + (diff @ jnp.eye(dim) @ diff)) ** (-0.5)


# Define the target distribution and kernel
target = CustomTargetDistribution()
kernel = KernelJax(custom_kernel)
ksd = KernelSteinDiscrepancyJax(target=target, kernel=kernel)

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