Add entropy for Wishart distribution.

numpyro/distributions/continuous.py

@@ -71,6 +71,7 @@
     gammaincinv,
     lazy_property,
     matrix_to_tril_vec,
+    multidigamma,
     promote_shapes,
     signed_stick_breaking_tril,
     validate_sample,
@@ -2706,6 +2707,16 @@ def infer_shapes(
             concentration, scale_matrix, rate_matrix, scale_tril
         )
 
+    def entropy(self):
+        p = self.event_shape[-1]
+        return (
+            (p + 1) * jnp.linalg.slogdet(self.scale_tril).logabsdet
+            + p * (p + 1) / 2 * jnp.log(2)
+            + multigammaln(self.concentration / 2, p)
+            - (self.concentration - p - 1) / 2 * multidigamma(self.concentration / 2, p)
+            + self.concentration * p / 2
+        )
+
 
 class WishartCholesky(Distribution):
     """

numpyro/distributions/util.py

@@ -12,6 +12,7 @@
 from jax import jit, lax, random, vmap
 import jax.numpy as jnp
 from jax.scipy.linalg import solve_triangular
+from jax.scipy.special import digamma
 
 # Parameters for Transformed Rejection with Squeeze (TRS) algorithm - page 3.
 _tr_params = namedtuple(
@@ -634,6 +635,13 @@ def assert_one_of(**kwargs):
         )
 
 
+def multidigamma(a: jnp.ndarray, d: jnp.ndarray) -> jnp.ndarray:
+    """
+    Derivative of the log of multivariate gamma.
+    """
+    return digamma(a[..., None] - 0.5 * jnp.arange(d)).sum(axis=-1)
+
+
 # The is sourced from: torch.distributions.util.py
 #
 # Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)