Entropy estimator. Source: https://x.com/adad8m/status/1745109776458187138?s=20

Kozachenko_Leonenko_entropy_est.py

import jax 
import jax.numpy as jnp


def ent(X):
    """Calculates the entropy of a dataset X."""

    N, D = X.shape  # Get the number of samples and dimensions

    # Compute pairwise squared distances
    dist_sq = jnp.sum((X[:, jnp.newaxis, :] - X[jnp.newaxis, :, :]) ** 2, axis=-1)

    # Set the diagonal to a large number so it doesn't affect the min calculation
    dist_sq = fill_diagonal(dist_sq, jnp.inf)

    # Find the minimum distance for each point
    min_dist = jnp.sqrt(jnp.min(dist_sq, axis=1))

    # Kozachenko-Leonenko estimator of the entropy (up to irrelevant constant)
    return jnp.mean(jnp.log((N - 1) * min_dist ** D))


@jax.jit
def functional(X):
    """Calculates the functional value for a dataset X."""

    return -ent(X) + jnp.mean(-log_density(X))