KL divergence between two `TransformedDistribution`s with a shared non-injective bijector (`AbsoluteValue`) returns the base KL

tensorflow_probability/python/distributions/transformed_distribution.py の _kl_transformed_transformed から始めてください。ここでは a.bijector == b.bijector の分岐が基底分布のKLを返しています。bijectorの _is_injective フラグを確認し、AbsoluteValue のような非単射のbijectorに対しては NotImplementedError を発生させ、単射のケースは変更しないでください。再現例のfolded-normalのペアがエラーを発生させ、そのパスをテストがカバーしていれば完了です。

Description

tfd.kl_divergence(a, b) for two TransformedDistributions returns kl_divergence(a.distribution, b.distribution) whenever a.bijector == b.bijector. That identity, KL(g(P) || g(Q)) = KL(P || Q), holds only for injective g. For a non-injective bijector the KL of the pushforwards can be strictly smaller, because information is lost: by the chain rule of KL, the difference is E_{Y ~ g(P)} KL(P(. | g = Y) || Q(. | g = Y)) >= 0.

tfb.AbsoluteValue is documented as "This non-injective bijector allows for transformations of scalar distributions with the absolute value function". TransformedDistribution.log_prob handles it correctly by summing over both preimages, but the KL rule does not check injectivity. For example KL(|N(-1, 1)| || |N(1, 1)|) returns 2.0, although the two distributions are identical. TransformedDistribution._entropy already refuses non-injective bijectors (L570-L572).

Reproduction

NumPy substrate (tensorflow_probability.substrates.numpy); the KL rule is the same code for TF and JAX:

import numpy as np
from scipy.integrate import quad
import tensorflow_probability.substrates.numpy as tfp

tfd, tfb = tfp.distributions, tfp.bijectors
print("tensorflow_probability", tfp.__version__, "(NumPy substrate)")


def folded(m, s=1.0):  # |X|, X ~ N(m, s^2)
    return tfd.TransformedDistribution(tfd.Normal(np.float64(m), np.float64(s)), tfb.AbsoluteValue())


def true_kl(p, q):  # KL(p || q) by quadrature, using TFP's own (correct, two-branch) log_prob
    return quad(lambda y: np.exp(p.log_prob(y)) * (p.log_prob(y) - q.log_prob(y)), 0, 40, limit=200)[0]


p0 = folded(-1.0)
print(f"log_prob of |N(-1,1)| integrates to {quad(lambda y: np.exp(p0.log_prob(y)), 0, 40)[0]:.6f}")
for name, p, q in [("|N(-1,1)| vs |N(1,1)| (same law)", folded(-1.0), folded(1.0)),
                   ("|N(3,1)| vs |N(0,1)|", folded(3.0), folded(0.0)),
                   ("|N(0.5,1)| vs |N(0,1)| (= HalfNormal(1))", folded(0.5), folded(0.0))]:
    kl, kl_true = float(tfd.kl_divergence(p, q)), true_kl(p, q)
    print(f"{name:42s} kl_divergence = {kl:.4f}   true KL = {kl_true:.4f}")
    assert kl > kl_true + 0.1
print("Bug reproduced: the base-distribution KL is returned for a non-injective bijector")

Output:

tensorflow_probability 0.25.0 (NumPy substrate)
log_prob of |N(-1,1)| integrates to 1.000000
|N(-1,1)| vs |N(1,1)| (same law)           kl_divergence = 2.0000   true KL = 0.0000
|N(3,1)| vs |N(0,1)|                       kl_divergence = 4.5000   true KL = 3.8107
|N(0.5,1)| vs |N(0,1)| (= HalfNormal(1))   kl_divergence = 0.1250   true KL = 0.0136
Bug reproduced: the base-distribution KL is returned for a non-injective bijector

Expected behavior

Either the true KL or NotImplementedError (there is no closed form for the folded-normal KL). The base KL is right only when the conditional distribution of the sign given |X| is the same under both distributions, e.g. when both bases are centred at 0.

Root cause

distributions/transformed_distribution.py#L731-L754 (v0.25.0, 9709569; identical on main, 6fa068e):

  with tf.name_scope(name or 'kl_transformed_transformed'):
    if a.bijector == b.bijector:
      return kullback_leibler.kl_divergence(a.distribution, b.distribution)

Proposed fix

--- a/tensorflow_probability/python/distributions/transformed_distribution.py
+++ b/tensorflow_probability/python/distributions/transformed_distribution.py
@@ -747,6 +747,10 @@
   """
   with tf.name_scope(name or 'kl_transformed_transformed'):
     if a.bijector == b.bijector:
+      if not a.bijector._is_injective:  # pylint: disable=protected-access
+        raise NotImplementedError(
+            'Unable to calculate KL divergence between {} and {} because '
+            'their bijector {} is not injective.'.format(a, b, a.bijector))
       return kullback_leibler.kl_divergence(a.distribution, b.distribution)
   raise NotImplementedError(
       'Unable to calculate KL divergence between {} and {} because '

Verifying the fix

fix_check.py applies this change to the installed 0.25.0 function (NumPy substrate) in memory and re-registers it. It then checks that the folded pairs raise NotImplementedError, and that pairs with injective bijectors (Exp, Shift(Scale)) give exactly the same values as before:

tensorflow_probability 0.25.0 (NumPy substrate), patched _kl_transformed_transformed
|N(-1,1)| vs |N(1,1)|    -> NotImplementedError: ...value", batch_shape=[], min_event_ndims=0) is not injective.
|N(0.5,1)| vs |N(0,1)|   -> NotImplementedError: ...value", batch_shape=[], min_event_ndims=0) is not injective.
LogNormal via Exp            -> fixed 0.443147, released 0.443147
Shift(1)(Scale(2)) of Normal -> fixed 0.183243, released 0.183243
FIX OK: non-injective bijectors raise NotImplementedError; injective bijectors unchanged

Versions

tensorflow-probability 0.25.0 (NumPy substrate; the JAX substrate of 0.25.0 does not import with the installed jax 0.12), numpy 2.5.4, scipy 1.18.1 for the reference quadrature, Python 3.13, Linux, CPU only.

I would be happy to open a PR with this fix and a test.