From 5946ebda957e258d942763fd30ec954b4fdb8ee7 Mon Sep 17 00:00:00 2001
From: uditjainstjis <uditj668@gmail.com>
Date: Fri, 24 Jul 2026 04:54:37 +0530
Subject: [PATCH 1/2] Add entropy dispatch for distributions

Implements an `entropy` dispatcher under `pymc.distributions.moments`,
mirroring the existing `mean` dispatch (#7530). It returns the differential
entropy for continuous distributions and the Shannon entropy (in nats) for
discrete ones, following the `scipy.stats` convention.

Closed-form entropies are registered for the univariate distributions
(Normal, Uniform, Exponential, Laplace, AsymmetricLaplace, Cauchy,
HalfCauchy, HalfNormal, Gamma, InverseGamma, Beta, Logistic, LogNormal,
Gumbel, Pareto, Weibull, StudentT, Triangular, Moyal, VonMises, Bernoulli,
Geometric, DiscreteUniform, Categorical) as well as MvNormal and Dirichlet.
ChiSquared is covered automatically through its Gamma parametrization.

Every value is checked against `scipy.stats`'s `entropy`, and the result is
verified to be differentiable w.r.t. the distribution parameters (the
regularization use case motivating the issue). Distributions without a
closed-form entropy raise NotImplementedError.

Closes #8085
---
 pymc/distributions/moments/__init__.py      |   3 +-
 pymc/distributions/moments/entropy.py       | 261 ++++++++++++++++++++
 tests/distributions/moments/test_entropy.py | 228 +++++++++++++++++
 3 files changed, 491 insertions(+), 1 deletion(-)
 create mode 100644 pymc/distributions/moments/entropy.py
 create mode 100644 tests/distributions/moments/test_entropy.py

diff --git a/pymc/distributions/moments/__init__.py b/pymc/distributions/moments/__init__.py
index b61e8b6400..19c3532dbb 100644
--- a/pymc/distributions/moments/__init__.py
+++ b/pymc/distributions/moments/__init__.py
@@ -14,6 +14,7 @@
 
 """Moments dispatchers for pymc random variables."""
 
+from pymc.distributions.moments.entropy import entropy
 from pymc.distributions.moments.means import mean
 
-__all__ = ["mean"]
+__all__ = ["entropy", "mean"]
diff --git a/pymc/distributions/moments/entropy.py b/pymc/distributions/moments/entropy.py
new file mode 100644
index 0000000000..5df8a1c07a
--- /dev/null
+++ b/pymc/distributions/moments/entropy.py
@@ -0,0 +1,261 @@
+#   Copyright 2024 - present The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+"""Entropy dispatcher for pymc random variables.
+
+The entropy of a random variable is the differential entropy for continuous
+variables and the Shannon entropy (in nats) for discrete variables, matching
+the convention of :meth:`scipy.stats.rv_continuous.entropy`.
+"""
+
+from functools import singledispatch
+
+import numpy as np
+
+from pytensor import tensor as pt
+from pytensor.tensor.random.basic import (
+    BernoulliRV,
+    BetaRV,
+    CategoricalRV,
+    CauchyRV,
+    DirichletRV,
+    ExponentialRV,
+    GammaRV,
+    GeometricRV,
+    GumbelRV,
+    HalfNormalRV,
+    InvGammaRV,
+    LaplaceRV,
+    LogisticRV,
+    LogNormalRV,
+    MvNormalRV,
+    NormalRV,
+    ParetoRV,
+    StudentTRV,
+    TriangularRV,
+    UniformRV,
+    VonMisesRV,
+)
+from pytensor.tensor.variable import TensorVariable
+
+from pymc.distributions.continuous import (
+    AsymmetricLaplaceRV,
+    HalfCauchyRV,
+    MoyalRV,
+    WeibullBetaRV,
+)
+from pymc.distributions.discrete import DiscreteUniformRV
+from pymc.distributions.shape_utils import maybe_resize, rv_size_is_none
+
+__all__ = ["entropy"]
+
+
+def _betaln(a, b):
+    return pt.gammaln(a) + pt.gammaln(b) - pt.gammaln(a + b)
+
+
+@singledispatch
+def _entropy(op, rv, *rv_inputs) -> TensorVariable:
+    raise NotImplementedError(f"Variable {rv} of type {op} has no entropy implementation.")
+
+
+def entropy(rv: TensorVariable) -> TensorVariable:
+    """Compute the entropy of a random variable.
+
+    The entropy is the differential entropy for continuous distributions and
+    the Shannon entropy (in nats) for discrete distributions. This matches the
+    convention used by ``scipy.stats``'s ``entropy`` method.
+
+    The only parameter to this function is the RandomVariable
+    for which the entropy is to be derived.
+    """
+    return _entropy(rv.owner.op, rv, *rv.owner.inputs)
+
+
+# --- Continuous univariate ---
+
+
+@_entropy.register(NormalRV)
+def normal_entropy(op, rv, rng, size, mu, sigma):
+    return maybe_resize(0.5 * pt.log(2 * np.pi * np.e) + pt.log(sigma), size)
+
+
+@_entropy.register(UniformRV)
+def uniform_entropy(op, rv, rng, size, lower, upper):
+    return maybe_resize(pt.log(upper - lower), size)
+
+
+@_entropy.register(ExponentialRV)
+def exponential_entropy(op, rv, rng, size, mu):
+    # ``mu`` is the mean (scale) of the distribution
+    return maybe_resize(1 + pt.log(mu), size)
+
+
+@_entropy.register(LaplaceRV)
+def laplace_entropy(op, rv, rng, size, mu, b):
+    return maybe_resize(1 + pt.log(2 * b), size)
+
+
+@_entropy.register(AsymmetricLaplaceRV)
+def asymmetric_laplace_entropy(op, rv, rng, size, b, kappa, mu):
+    # scale = 1 / b
+    return maybe_resize(1 + pt.log((kappa + 1 / kappa) / b), size)
+
+
+@_entropy.register(CauchyRV)
+def cauchy_entropy(op, rv, rng, size, alpha, beta):
+    return maybe_resize(pt.log(4 * np.pi * beta), size)
+
+
+@_entropy.register(HalfCauchyRV)
+def halfcauchy_entropy(op, rv, rng, size, beta):
+    return maybe_resize(pt.log(2 * np.pi * beta), size)
+
+
+@_entropy.register(HalfNormalRV)
+def halfnormal_entropy(op, rv, rng, size, loc, sigma):
+    return maybe_resize(0.5 * pt.log(np.pi * sigma**2 / 2) + 0.5, size)
+
+
+@_entropy.register(GammaRV)
+def gamma_entropy(op, rv, rng, size, alpha, inv_beta):
+    # ``inv_beta`` is the scale (1 / rate)
+    return maybe_resize(
+        alpha + pt.log(inv_beta) + pt.gammaln(alpha) + (1 - alpha) * pt.digamma(alpha),
+        size,
+    )
+
+
+@_entropy.register(InvGammaRV)
+def invgamma_entropy(op, rv, rng, size, alpha, beta):
+    # ``beta`` is the scale
+    return maybe_resize(
+        alpha + pt.log(beta) + pt.gammaln(alpha) - (1 + alpha) * pt.digamma(alpha),
+        size,
+    )
+
+
+@_entropy.register(BetaRV)
+def beta_entropy(op, rv, rng, size, alpha, beta):
+    return maybe_resize(
+        _betaln(alpha, beta)
+        - (alpha - 1) * pt.digamma(alpha)
+        - (beta - 1) * pt.digamma(beta)
+        + (alpha + beta - 2) * pt.digamma(alpha + beta),
+        size,
+    )
+
+
+@_entropy.register(LogisticRV)
+def logistic_entropy(op, rv, rng, size, mu, s):
+    return maybe_resize(pt.log(s) + 2, size)
+
+
+@_entropy.register(LogNormalRV)
+def lognormal_entropy(op, rv, rng, size, mu, sigma):
+    return maybe_resize(mu + 0.5 * pt.log(2 * np.pi * np.e * sigma**2), size)
+
+
+@_entropy.register(GumbelRV)
+def gumbel_entropy(op, rv, rng, size, mu, beta):
+    return maybe_resize(pt.log(beta) + np.euler_gamma + 1, size)
+
+
+@_entropy.register(ParetoRV)
+def pareto_entropy(op, rv, rng, size, alpha, m):
+    return maybe_resize(pt.log(m / alpha) + 1 / alpha + 1, size)
+
+
+@_entropy.register(WeibullBetaRV)
+def weibull_entropy(op, rv, rng, size, alpha, beta):
+    # ``alpha`` is the shape, ``beta`` is the scale
+    return maybe_resize(np.euler_gamma * (1 - 1 / alpha) + pt.log(beta / alpha) + 1, size)
+
+
+@_entropy.register(StudentTRV)
+def studentt_entropy(op, rv, rng, size, nu, mu, sigma):
+    return maybe_resize(
+        pt.log(sigma)
+        + 0.5 * (nu + 1) * (pt.digamma((nu + 1) / 2) - pt.digamma(nu / 2))
+        + pt.log(pt.sqrt(nu))
+        + _betaln(nu / 2, 0.5),
+        size,
+    )
+
+
+@_entropy.register(TriangularRV)
+def triangular_entropy(op, rv, rng, size, lower, c, upper):
+    return maybe_resize(0.5 + pt.log((upper - lower) / 2), size)
+
+
+@_entropy.register(MoyalRV)
+def moyal_entropy(op, rv, rng, size, mu, sigma):
+    return maybe_resize(
+        pt.log(sigma) + 0.5 * pt.log(2 * np.pi) + 0.5 * (np.euler_gamma + np.log(2) + 1),
+        size,
+    )
+
+
+@_entropy.register(VonMisesRV)
+def vonmises_entropy(op, rv, rng, size, mu, kappa):
+    return maybe_resize(
+        pt.log(2 * np.pi * pt.i0(kappa)) - kappa * pt.i1(kappa) / pt.i0(kappa), size
+    )
+
+
+# --- Discrete univariate ---
+
+
+@_entropy.register(BernoulliRV)
+def bernoulli_entropy(op, rv, rng, size, p):
+    return maybe_resize(-p * pt.log(p) - (1 - p) * pt.log1p(-p), size)
+
+
+@_entropy.register(GeometricRV)
+def geometric_entropy(op, rv, rng, size, p):
+    return maybe_resize((-(1 - p) * pt.log1p(-p) - p * pt.log(p)) / p, size)
+
+
+@_entropy.register(DiscreteUniformRV)
+def discrete_uniform_entropy(op, rv, rng, size, lower, upper):
+    return maybe_resize(pt.log(upper - lower + 1), size)
+
+
+@_entropy.register(CategoricalRV)
+def categorical_entropy(op, rv, rng, size, p):
+    return maybe_resize(-pt.sum(p * pt.log(p), axis=-1), size)
+
+
+# --- Multivariate ---
+
+
+@_entropy.register(MvNormalRV)
+def mvnormal_entropy(op, rv, rng, size, mu, cov):
+    k = cov.shape[-1]
+    _, logdet = pt.linalg.slogdet(cov)
+    res = 0.5 * k * pt.log(2 * np.pi * np.e) + 0.5 * logdet
+    if rv_size_is_none(size):
+        return res
+    return maybe_resize(res, size)
+
+
+@_entropy.register(DirichletRV)
+def dirichlet_entropy(op, rv, rng, size, a):
+    a0 = pt.sum(a, axis=-1)
+    k = a.shape[-1]
+    log_beta = pt.sum(pt.gammaln(a), axis=-1) - pt.gammaln(a0)
+    res = log_beta + (a0 - k) * pt.digamma(a0) - pt.sum((a - 1) * pt.digamma(a), axis=-1)
+    if rv_size_is_none(size):
+        return res
+    return maybe_resize(res, size)
diff --git a/tests/distributions/moments/test_entropy.py b/tests/distributions/moments/test_entropy.py
new file mode 100644
index 0000000000..7d42564f23
--- /dev/null
+++ b/tests/distributions/moments/test_entropy.py
@@ -0,0 +1,228 @@
+#   Copyright 2024 - present The PyMC Developers
+#
+#   Licensed under the Apache License, Version 2.0 (the "License");
+#   you may not use this file except in compliance with the License.
+#   You may obtain a copy of the License at
+#
+#       http://www.apache.org/licenses/LICENSE-2.0
+#
+#   Unless required by applicable law or agreed to in writing, software
+#   distributed under the License is distributed on an "AS IS" BASIS,
+#   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#   See the License for the specific language governing permissions and
+#   limitations under the License.
+
+import numpy as np
+import pytest
+
+from pytensor import function
+from pytensor import tensor as pt
+from pytensor.compile.mode import Mode
+from scipy import stats
+from scipy.stats import (
+    bernoulli,
+    beta,
+    cauchy,
+    chi2,
+    dirichlet,
+    expon,
+    gamma,
+    geom,
+    gumbel_r,
+    halfcauchy,
+    halfnorm,
+    invgamma,
+    laplace,
+    laplace_asymmetric,
+    logistic,
+    lognorm,
+    moyal,
+    multivariate_normal,
+    norm,
+    pareto,
+    randint,
+    t,
+    triang,
+    uniform,
+    vonmises,
+    weibull_min,
+)
+
+from pymc import (
+    AsymmetricLaplace,
+    Bernoulli,
+    Beta,
+    Binomial,
+    Categorical,
+    Cauchy,
+    ChiSquared,
+    Dirichlet,
+    DiscreteUniform,
+    ExGaussian,
+    Exponential,
+    Flat,
+    Gamma,
+    Geometric,
+    Gumbel,
+    HalfCauchy,
+    HalfFlat,
+    HalfNormal,
+    HalfStudentT,
+    InverseGamma,
+    Kumaraswamy,
+    Laplace,
+    Logistic,
+    LogitNormal,
+    LogNormal,
+    Moyal,
+    MvNormal,
+    MvStudentT,
+    Normal,
+    Pareto,
+    Poisson,
+    Rice,
+    SkewNormal,
+    StudentT,
+    Triangular,
+    Uniform,
+    VonMises,
+    Wald,
+    Weibull,
+)
+from pymc.distributions.moments.entropy import entropy
+
+mode = Mode(linker="py", optimizer=None)
+
+
+@pytest.mark.parametrize(
+    ["dist", "scipy_equiv", "dist_params", "scipy_params"],
+    [
+        [
+            AsymmetricLaplace,
+            laplace_asymmetric,
+            {"kappa": 2, "mu": 0.2, "b": 1 / 1.2},
+            {"kappa": 2, "loc": 0.2, "scale": 1.2},
+        ],
+        [Bernoulli, bernoulli, {"p": 0.6}, {"p": 0.6}],
+        [Beta, beta, {"alpha": 3, "beta": 2}, {"a": 3, "b": 2}],
+        [Cauchy, cauchy, {"alpha": 2, "beta": 1.5}, {"loc": 2, "scale": 1.5}],
+        [ChiSquared, chi2, {"nu": 6}, {"df": 6}],
+        [DiscreteUniform, randint, {"lower": 2, "upper": 8}, {"low": 2, "high": 9}],
+        [Exponential, expon, {"lam": 0.5}, {"scale": 2}],
+        [Gamma, gamma, {"alpha": 4, "beta": 3}, {"a": 4, "scale": 1 / 3}],
+        [Geometric, geom, {"p": 0.1}, {"p": 0.1}],
+        [Gumbel, gumbel_r, {"mu": 2, "beta": 1}, {"loc": 2, "scale": 1}],
+        [HalfCauchy, halfcauchy, {"beta": 1.5}, {"scale": 1.5}],
+        [HalfNormal, halfnorm, {"sigma": 3}, {"scale": 3}],
+        [InverseGamma, invgamma, {"alpha": 2, "beta": 2}, {"a": 2, "scale": 2}],
+        [Laplace, laplace, {"mu": 2, "b": 2}, {"loc": 2, "scale": 2}],
+        [Logistic, logistic, {"mu": 2, "s": 1.5}, {"loc": 2, "scale": 1.5}],
+        [LogNormal, lognorm, {"mu": 0.3, "sigma": 0.6}, {"scale": np.exp(0.3), "s": 0.6}],
+        [Moyal, moyal, {"mu": 2, "sigma": 2}, {"loc": 2, "scale": 2}],
+        [Normal, norm, {"mu": 2, "sigma": 3}, {"loc": 2, "scale": 3}],
+        [Pareto, pareto, {"alpha": 5, "m": 2}, {"b": 5, "scale": 2}],
+        [StudentT, t, {"nu": 6, "mu": 0, "sigma": 2}, {"df": 6, "loc": 0, "scale": 2}],
+        [Triangular, triang, {"lower": 0, "c": 0.5, "upper": 2}, {"c": 0.25, "loc": 0, "scale": 2}],
+        [Uniform, uniform, {"lower": -1, "upper": 4}, {"loc": -1, "scale": 5}],
+        [VonMises, vonmises, {"mu": 1, "kappa": 2.5}, {"loc": 1, "kappa": 2.5}],
+        [Weibull, weibull_min, {"alpha": 1.5, "beta": 2}, {"c": 1.5, "scale": 2}],
+    ],
+)
+def test_entropy_equal_to_scipy(dist, scipy_equiv, dist_params, scipy_params):
+    rv = dist.dist(**dist_params)
+    pymc_entropy = entropy(rv).eval(mode=mode)
+    scipy_entropy = scipy_equiv(**scipy_params).entropy()
+    assert np.asarray(pymc_entropy).shape == np.asarray(scipy_entropy).shape
+    np.testing.assert_allclose(pymc_entropy, scipy_entropy, rtol=1e-6)
+
+    # entropy of a batched distribution broadcasts over the batch dimension
+    pymc_entropy_tiled = entropy(dist.dist(shape=(3,), **dist_params)).eval()
+    np.testing.assert_allclose(pymc_entropy_tiled, np.tile(pymc_entropy, 3), rtol=1e-6)
+
+
+@pytest.mark.parametrize(
+    ["dist", "dist_params", "expected"],
+    [
+        [Categorical, {"p": [0.1, 0.2, 0.3, 0.4]}, stats.entropy([0.1, 0.2, 0.3, 0.4])],
+    ],
+)
+def test_entropy_equal_expected(dist, dist_params, expected):
+    rv = dist.dist(**dist_params)
+    pymc_entropy = entropy(rv).eval(mode=mode)
+    np.testing.assert_allclose(pymc_entropy, expected, rtol=1e-6)
+
+
+@pytest.mark.parametrize(
+    ["dist", "dist_params", "scipy_entropy"],
+    [
+        [
+            MvNormal,
+            {
+                "mu": np.zeros(3),
+                "cov": np.array([[2.0, 0.3, 0.1], [0.3, 1.5, 0.2], [0.1, 0.2, 1.0]]),
+            },
+            multivariate_normal(
+                np.zeros(3), np.array([[2.0, 0.3, 0.1], [0.3, 1.5, 0.2], [0.1, 0.2, 1.0]])
+            ).entropy(),
+        ],
+        [
+            Dirichlet,
+            {"a": np.array([2.0, 3.0, 4.0, 1.0])},
+            dirichlet(np.array([2.0, 3.0, 4.0, 1.0])).entropy(),
+        ],
+    ],
+)
+def test_entropy_multivariate(dist, dist_params, scipy_entropy):
+    rv = dist.dist(**dist_params)
+    pymc_entropy = entropy(rv).eval(mode=mode)
+    assert np.asarray(pymc_entropy).shape == ()
+    np.testing.assert_allclose(pymc_entropy, scipy_entropy, rtol=1e-6)
+
+    # a batch of independent multivariate distributions yields one entropy per batch element
+    pymc_entropy_batched = entropy(dist.dist(size=(4,), **dist_params)).eval()
+    np.testing.assert_allclose(pymc_entropy_batched, np.full(4, pymc_entropy), rtol=1e-6)
+
+
+@pytest.mark.parametrize(
+    ["dist", "dist_params"],
+    [
+        [Binomial, {"n": 5, "p": 0.6}],
+        [ExGaussian, {"mu": 0, "sigma": 1, "nu": 1}],
+        [Flat, {}],
+        [HalfFlat, {}],
+        [HalfStudentT, {"nu": 3, "sigma": 1}],
+        [Kumaraswamy, {"a": 2, "b": 2}],
+        [LogitNormal, {"mu": 0, "sigma": 1}],
+        [MvStudentT, {"mu": np.zeros(3), "scale": np.eye(3), "nu": 4}],
+        [Poisson, {"mu": 3}],
+        [Rice, {"nu": 1, "sigma": 1}],
+        [SkewNormal, {"mu": 0, "sigma": 1, "alpha": 2}],
+        [Wald, {"mu": 1, "lam": 1}],
+    ],
+)
+def test_no_entropy(dist, dist_params):
+    with pytest.raises(NotImplementedError):
+        entropy(dist.dist(**dist_params))
+
+
+@pytest.mark.parametrize(
+    ["dist", "param", "value"],
+    [
+        # A regularization-by-entropy term (e.g. in RL policies) requires the
+        # entropy to be differentiable w.r.t. the distribution parameters.
+        [lambda p: Normal.dist(mu=0.0, sigma=p), "sigma", 2.0],
+        [lambda p: Gamma.dist(alpha=p, beta=1.0), "alpha", 3.0],
+        [lambda p: Beta.dist(alpha=p, beta=2.0), "alpha", 3.0],
+        [lambda p: StudentT.dist(nu=p, mu=0.0, sigma=1.0), "nu", 5.0],
+        [lambda p: VonMises.dist(mu=0.0, kappa=p), "kappa", 2.5],
+    ],
+)
+def test_entropy_is_differentiable(dist, param, value):
+    p = pt.scalar(param)
+    grad = pt.grad(entropy(dist(p)).sum(), p)
+    grad_fn = function([p], grad, mode=mode, on_unused_input="ignore")
+    analytic = grad_fn(value)
+    h = 1e-5
+    entropy_fn = function([p], entropy(dist(p)), mode=mode, on_unused_input="ignore")
+    finite_diff = (entropy_fn(value + h) - entropy_fn(value - h)) / (2 * h)
+    np.testing.assert_allclose(analytic, finite_diff, rtol=1e-4)

From 329b038203a105eca757485d2b772060573e507f Mon Sep 17 00:00:00 2001
From: Udit Jain <uditjainstjis@gmail.com>
Date: Fri, 24 Jul 2026 08:14:24 +0530
Subject: [PATCH 2/2] Register test_entropy.py in CI test matrix

---
 .github/workflows/tests.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 40473a014c..75f7bac323 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -76,6 +76,7 @@ jobs:
             tests/distributions/test_continuous.py
             tests/distributions/test_multivariate.py
             tests/distributions/moments/test_means.py
+            tests/distributions/moments/test_entropy.py
 
           - |
             tests/distributions/test_censored.py