From 546004907469d6a580fca629b0f28e4aed45adae Mon Sep 17 00:00:00 2001 From: Andreas Date: Sun, 1 Nov 2026 23:50:26 +0100 Subject: [PATCH 1/2] Raise in kl_divergence of TransformedDistributions with a non-injective bijector KL(g(P) || g(Q)) equals KL(P || Q) only for injective g. For a shared non-injective bijector such as AbsoluteValue, the base KL overstates the divergence: KL(|N(-1, 1)| || |N(1, 1)|) was returned as 2.0 although the two distributions are identical. Raise NotImplementedError instead, as TransformedDistribution._entropy already does for such bijectors. Fixes #2036 --- .../distributions/transformed_distribution.py | 4 ++++ .../transformed_distribution_test.py | 15 +++++++++++++++ 2 files changed, 19 insertions(+) diff --git a/tensorflow_probability/python/distributions/transformed_distribution.py b/tensorflow_probability/python/distributions/transformed_distribution.py index 8d05d607bc..9f496cd28f 100644 --- a/tensorflow_probability/python/distributions/transformed_distribution.py +++ b/tensorflow_probability/python/distributions/transformed_distribution.py @@ -747,6 +747,10 @@ def _kl_transformed_transformed(a, b, name=None): """ 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 ' diff --git a/tensorflow_probability/python/distributions/transformed_distribution_test.py b/tensorflow_probability/python/distributions/transformed_distribution_test.py index cbc02404a2..7fe0947cc0 100644 --- a/tensorflow_probability/python/distributions/transformed_distribution_test.py +++ b/tensorflow_probability/python/distributions/transformed_distribution_test.py @@ -577,6 +577,21 @@ def testTransformedKLDifferentBijectorFails(self): NotImplementedError, r'their bijectors are not equal'): kullback_leibler.kl_divergence(d1, d2) + def testTransformedKLNonInjectiveBijectorFails(self): + # |X| has the same law for X ~ N(-1, 1) and X ~ N(1, 1), while the KL of + # the base distributions is 2, so the base KL must not be returned when + # the shared bijector is not injective. + d1 = transformed_distribution.TransformedDistribution( + normal_lib.Normal(loc=-1., scale=1.), + bijector=absolute_value.AbsoluteValue(), + validate_args=True) + d2 = transformed_distribution.TransformedDistribution( + normal_lib.Normal(loc=1., scale=1.), + bijector=absolute_value.AbsoluteValue(), + validate_args=True) + with self.assertRaisesRegex(NotImplementedError, r'is not injective'): + kullback_leibler.kl_divergence(d1, d2) + def testTransformedNormalNormalKL(self): batch_size = 6 mu_a = np.array([3.0] * batch_size).astype(np.float32) From 56de420e48e7c670b62f4befc40ebcb3e49f8693 Mon Sep 17 00:00:00 2001 From: Andreas Date: Sun, 1 Nov 2026 23:56:22 +0100 Subject: [PATCH 2/2] Re-run the CLA check