Add KL Divergence helper

.github/workflows/tests.yml

@@ -76,6 +76,7 @@ jobs:
             tests/stats/test_log_likelihood.py
             tests/distributions/test_distribution.py
             tests/distributions/test_discrete.py
+            tests/distributions/stats/test_kl_divergence.py
 
           - |
             tests/tuning/test_scaling.py

pymc/distributions/__init__.py

@@ -114,6 +114,12 @@
 )
 from pymc.distributions.truncated import Truncated
 
+# now dispatched stats are being initialized
+# isort: off
+import pymc.distributions.stats
+
+# isort: on
+
 __all__ = [
     "Uniform",
     "Flat",

pymc/distributions/stats/__init__.py

@@ -0,0 +1,14 @@
+#   Copyright 2024 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.
+from pymc.distributions.stats import kl_divergence

pymc/distributions/stats/kl_divergence.py

@@ -0,0 +1,35 @@
+#   Copyright 2024 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 pytensor.tensor as pt
+
+from pymc.distributions.continuous import Normal
+from pymc.logprob.abstract import _kl_div
+
+
+@_kl_div.register(Normal, Normal)
+def _normal_normal_kl(
+    q_dist: Normal,
+    p_dist: Normal,
+    q_inputs: list[pt.TensorVariable],
+    p_inputs: list[pt.TensorVariable],
+):
+    _, _, _, q_mu, q_sigma = q_inputs
+    _, _, _, p_mu, p_sigma = p_inputs
+    diff_log_scale = pt.log(q_sigma) - pt.log(p_sigma)
+    return (
+        0.5 * (q_mu / p_sigma - p_mu / p_sigma) ** 2
+        + 0.5 * pt.expm1(2.0 * diff_log_scale)
+        - diff_log_scale
+    )

pymc/logprob/__init__.py

@@ -37,6 +37,7 @@
 from pymc.logprob.basic import (
     conditional_logp,
     icdf,
+    kl_div,
     logcdf,
     logp,
     transformed_conditional_logp,
@@ -58,4 +59,5 @@
     "logp",
     "logcdf",
     "icdf",
+    "kl_div",
 )

pymc/logprob/abstract.py

@@ -39,6 +39,7 @@
 from collections.abc import Sequence
 from functools import singledispatch
 
+from multipledispatch import dispatch as multipledispatch
 from pytensor.graph.op import Op
 from pytensor.graph.utils import MetaType
 from pytensor.tensor import TensorVariable
@@ -153,3 +154,27 @@
 
 
 MeasurableVariable.register(MeasurableElemwise)
+
+
+@multipledispatch(Op, Op)
+def _kl_div(
+    q_dist: Op,
+    p_dist: Op,
+    q_inputs: list[TensorVariable],
+    p_inputs: list[TensorVariable],
+) -> TensorVariable:
+    raise NotImplementedError(f"KL Divergence is not implemented for {q_dist}, {p_dist}")
+
+
+def _kl_div_helper(
+    q_rv: TensorVariable,
+    p_rv: TensorVariable,
+) -> TensorVariable:
+    kl = _kl_div(
+        q_rv.owner.op,
+        p_rv.owner.op,
+        q_inputs=q_rv.owner.inputs,
+        p_inputs=p_rv.owner.inputs,
+    )
+    kl.name = f"{q_rv.name}_{p_rv.name}_kl"
+    return kl

pymc/logprob/basic.py

@@ -59,6 +59,7 @@
 from pymc.logprob.abstract import (
     MeasurableVariable,
     _icdf_helper,
+    _kl_div_helper,
     _logcdf_helper,
     _logprob,
     _logprob_helper,
@@ -633,6 +634,27 @@
     return logp_terms_list
 
 
+def kl_div(
+    q_rv: TensorVariable,
+    p_rv: TensorVariable,
+) -> TensorVariable:
+    r"""KL Divergence
+
+    Parameters
+    ----------
+    q_rv : TensorVariable
+        :math:`Q` Distribution
+    p_rv : TensorVariable
+        :math:`P` Distribution
+
+    Returns
+    -------
+    TensorVariable
+        :math::KL\{Q||P\}
+    """
+    return _kl_div_helper(q_rv, p_rv)
+
+
 def factorized_joint_logprob(*args, **kwargs):
     warnings.warn(
         "`factorized_joint_logprob` was renamed to `conditional_logp`. "

requirements-dev.txt

@@ -30,4 +30,4 @@ sphinx>=1.5
 sphinxext-rediraffe
 types-cachetools
 typing-extensions>=3.7.4
-watermark
+watermark

requirements.txt

@@ -2,6 +2,7 @@ arviz>=0.13.0
 cachetools>=4.2.1
 cloudpickle
 fastprogress>=0.2.0
+multipledispatch
 numpy>=1.15.0
 pandas>=0.24.0
 pytensor>=2.18.1,<2.19

tests/distributions/stats/test_kl_divergence.py

@@ -0,0 +1,25 @@
+#   Copyright 2024 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 pymc as pm
+
+from pymc.logprob import kl_div
+
+
+def test_normal_normal():
+    q = pm.Normal.dist(0, 1)
+    p = pm.Normal.dist(1, 2)
+    kl = kl_div(q, p)
+    np.testing.assert_allclose(kl.eval(), 0.44314718)