diff --git a/tensorflow_probability/python/bijectors/bijector.py b/tensorflow_probability/python/bijectors/bijector.py index 1548f53cc3..6be2737c39 100644 --- a/tensorflow_probability/python/bijectors/bijector.py +++ b/tensorflow_probability/python/bijectors/bijector.py @@ -1130,7 +1130,7 @@ def _call_inverse_log_det_jacobian(self, y, event_ndims, name, **kwargs): allow_packing=True) if event_ndims is None: - event_ndims = self.forward_min_event_ndims + event_ndims = self.inverse_min_event_ndims reduce_shape, assertions = ldj_reduction_shape( nest.map_structure(ps.shape, y), diff --git a/tensorflow_probability/python/bijectors/bijector_test.py b/tensorflow_probability/python/bijectors/bijector_test.py index 0d5ddb619f..1acc40e782 100644 --- a/tensorflow_probability/python/bijectors/bijector_test.py +++ b/tensorflow_probability/python/bijectors/bijector_test.py @@ -551,6 +551,14 @@ def testNoReductionWhenEventNdimsIsOmitted(self): np.sum(-np.log(x), axis=-1), self.evaluate(bij.inverse_log_det_jacobian(x))) + def testInverseWithEventDimsOmitted(self): + bij = tfb.Split(2) + + self.assertAllEqual( + 0.0, + self.evaluate(bij.inverse_log_det_jacobian( + [tf.ones((3, 4, 5)), tf.ones((3, 4, 5))]))) + def testReduceEventNdimsForwardRaiseError(self): x = [[[1., 2.], [3., 4.]]] bij = ExpOnlyJacobian(forward_min_event_ndims=1)