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

Open Beginner friendly
  • Is something wrong?
#2,036 0 comments 0 reactions 0 assignees View on GitHub

@Andreas-Hum is already working on this.

Since Nov 1, 2026.

  • #2037 by @Andreas-Hum — open

Assessment

Difficulty
2/5
Estimated time
1-3 hours
Newbie friendliness
72/100
Issue type
Bug
Clarity
Clearly specified
Activity status
Active
Tech stack
python
Domain
machine-learning

Research direction

Start in tensorflow_probability/python/distributions/transformed_distribution.py at _kl_transformed_transformed, where the a.bijector == b.bijector branch returns the base-distribution KL. Check the bijector's _is_injective flag and raise NotImplementedError for non-injective bijectors such as AbsoluteValue, leaving injective cases unchanged. Done when the folded-normal pairs in the reproduction raise and a test covers that path.

Written by the indexing model from the issue text.

Description

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.

Dominant language
Jupyter Notebook
Stars
4.4k
Forks
1.1k
No merged PRs in 30d

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from tensorflow/probability

All issues in tensorflow/probability

Similar issues

More Machine Learning issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.