Fix bug in Biject.inverse_log_det_jacobian when no event_ndims specif…

…ied. PiperOrigin-RevId: 369272955
tensorflow · Apr 19, 2021 · 008f362 · 008f362
1 parent 1aee21d
commit 008f362
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 1 deletion.
diff --git 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
@@ -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)