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)