Fix *LDJ ratio computation.

The old code was broken in many ways:

- Multipart bijectors didn't work
- Scalar LDJ bijectors didn't reduce correctly in all situations
- Subclass *LDJ ratio functions didn't correctly override parent class's

There are four major issues with the solution and one small one. The resultant solution feels unsatisfactory.

1) There is no unified interpretation of an LDJ of amulti-part bijector. tfb.Bijector has one piece of code to do it, and tfb.Composition has a different one which are not compatible (tfb.Bijector's is more restrictive). A simple disagreement I ran into is that tfb.Bijector forbids len(set(map(lambda e, me: e - me, event_ndims, min_event_ndims))) > 1 from being True, while tfb.JointMap will happily broadcast component LDJs. This forces me to use two code paths when computing the *ldj ratios rather than just one (tfb.Bijector's), as well as plumb event_ndims args to the custom LDJ ratio functions.

I don't know of this is WAI or a bug, but both behaviors are enshrined in tests.

2) JD's default bijector relies on kwargs to function, which forced me to plumb them throughout the entire machinery, which looks horrid and is likely fragile.

3) I had to put ldj_ratio.py into the `bijector` lib because of circular dependencies. This makes the already complicated code even harder to understand.

4) The kahan summation really should be decided by the bijector and not the LDJ ratio function. Since kahan summation has a small performance penalty and incomplete JAX support, I turned it off by default.

The small issue is that the seemingly innocuous refactor of tfb.Composition broke the testChainDynamicToStatic test, which enshrines some incorrect code's behavior. While I managed to get the test passing, it's concerning that we're testing extremely internal aspects of the code in an end-to-end test.

PiperOrigin-RevId: 380913022

Author: SiegeLordEx

tensorflow_probability/python/bijectors/BUILD

@@ -75,7 +75,6 @@ multi_substrate_py_library(
         ":joint_map",
         ":kumaraswamy_cdf",
         ":lambertw_transform",
-        ":ldj_ratio",
         ":masked_autoregressive",
         ":matrix_inverse_tril",
         ":moyal_cdf",
@@ -121,6 +120,7 @@ multi_substrate_py_library(
         "bijector.py",
         "chain.py",
         "composition.py",
+        "ldj_ratio.py",
     ],
     deps = [
         # numpy dep,
@@ -265,11 +265,11 @@ multi_substrate_py_library(
     name = "scale_matvec_diag",
     srcs = ["scale_matvec_diag.py"],
     deps = [
-        ":ldj_ratio",
         ":scale_matvec_linear_operator",
         # tensorflow dep,
         "//tensorflow_probability/python/internal:dtype_util",
         "//tensorflow_probability/python/internal:parameter_properties",
+        "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:tensor_util",
     ],
 )
@@ -604,14 +604,6 @@ multi_substrate_py_library(
     ],
 )
 
-multi_substrate_py_library(
-    name = "ldj_ratio",
-    srcs = ["ldj_ratio.py"],
-    deps = [
-        # tensorflow dep,
-    ],
-)
-
 multi_substrate_py_library(
     name = "masked_autoregressive",
     srcs = ["masked_autoregressive.py"],

tensorflow_probability/python/bijectors/bijector.py

@@ -1422,7 +1422,7 @@ def _call_inverse_log_det_jacobian(self, y, event_ndims, name, **kwargs):
           # Non-injective bijectors don't use caching, and the resulting
           # LDJ is a tuple of LDJ over possible partitions on `x`.
           return tuple(
-              self._reduce_jacobian_det_over_shape(ildj, reduce_shape)
+              reduce_jacobian_det_over_shape(ildj, reduce_shape)
               for ildj in self._inverse_log_det_jacobian(y, **kwargs))
 
         # Make sure the unreduced ILDJ is in the cache.
@@ -1457,7 +1457,7 @@ def _call_inverse_log_det_jacobian(self, y, event_ndims, name, **kwargs):
               'Neither _forward_log_det_jacobian nor _inverse_log_det_jacobian '
               'is implemented. One or the other is required.')
 
-        return self._reduce_jacobian_det_over_shape(ildj, reduce_shape)
+        return reduce_jacobian_det_over_shape(ildj, reduce_shape)
 
   def inverse_log_det_jacobian(self,
                                y,
@@ -1582,7 +1582,7 @@ def _call_forward_log_det_jacobian(self, x, event_ndims, name, **kwargs):
               'Neither _forward_log_det_jacobian nor _inverse_log_det_jacobian '
               'is implemented. One or the other is required.')
 
-        return self._reduce_jacobian_det_over_shape(-ildj, reduce_shape)
+        return reduce_jacobian_det_over_shape(-ildj, reduce_shape)
 
   def forward_log_det_jacobian(self,
                                x,
@@ -1724,14 +1724,6 @@ def _name_and_control_scope(self, name=None):
         with tf.control_dependencies(deps) as deps_scope:
           yield deps_scope
 
-  def _reduce_jacobian_det_over_shape(self, unreduced, reduce_shape):
-    """Reduce LDJ over the rightmost `reduce_shape.ndims` dimensions."""
-    # Broadcast LDJ to the reduce shape (in case of is_constant_jacobian)
-    # and reduce over the trailing dimensions.
-    ones = tf.ones(reduce_shape, unreduced.dtype)
-    reduce_dims = ps.range(-ps.size(reduce_shape), 0)
-    return tf.reduce_sum(ones * unreduced, axis=reduce_dims)
-
   def _parameter_control_dependencies(self, is_init):
     """Returns a list of ops to be executed in members with graph deps.
 
@@ -2167,6 +2159,17 @@ def ldj_reduction_shape(shape_structure,
     return ldj_reduce_shape, assertions
 
 
+def reduce_jacobian_det_over_shape(unreduced,
+                                   reduce_shape,
+                                   sum_fn=tf.reduce_sum):
+  """Reduce LDJ over the rightmost `reduce_shape.ndims` dimensions."""
+  # Broadcast LDJ to the reduce shape (in case of is_constant_jacobian)
+  # and reduce over the trailing dimensions.
+  ones = tf.ones(reduce_shape, unreduced.dtype)
+  reduce_dims = ps.range(-ps.size(reduce_shape), 0)
+  return sum_fn(ones * unreduced, axis=reduce_dims)
+
+
 def _autodiff_log_det_jacobian(fn, x):
   """Automatically compute the log det jacobian of a scalar function."""
   # Note: x must be fully broadcast (`shape(x) == shape(fn(x))`); otherwise

tensorflow_probability/python/bijectors/chain.py

@@ -22,7 +22,6 @@
 import tensorflow.compat.v2 as tf
 from tensorflow_probability.python.bijectors import bijector as bijector_lib
 from tensorflow_probability.python.bijectors import composition
-from tensorflow_probability.python.bijectors import ldj_ratio
 from tensorflow_probability.python.internal import parameter_properties
 from tensorflow_probability.python.internal import prefer_static as ps
 
@@ -184,37 +183,3 @@ def __new__(cls, *args, **kwargs):
     'of `bijectors` is not a `CompositeTensor`, then a non-`CompositeTensor` '
     '`_Chain` instance is created instead. Bijector subclasses that inherit '
     'from `Chain` will also inherit from `CompositeTensor`.')
-
-
-@ldj_ratio.RegisterFLDJRatio(_Chain)
-def _fldj_ratio_chain(p, x, q, y):
-  """Sum-of-diffs FLDJRatio for Chains."""
-  if len(p.bijectors) != len(q.bijectors):
-    raise ValueError('Mismatched lengths of bijectors: `p` has '
-                     f'{len(p.bijectors)} but `q` has {len(q.bijectors)}.')
-  ratios = []
-  max_shp = []
-  for p, q in zip(reversed(p.bijectors), reversed(q.bijectors)):
-    ratios.append(ldj_ratio.forward_log_det_jacobian_ratio(
-        p, x, q, y, p.forward_min_event_ndims))
-    max_shp = ps.broadcast_shape(max_shp, ps.shape(ratios[-1]))
-    x, y = p.forward(x), q.forward(y)
-  ratios = [tf.broadcast_to(r, max_shp) for r in ratios]
-  return tf.add_n(ratios)
-
-
-@ldj_ratio.RegisterILDJRatio(_Chain)
-def _ildj_ratio_chain(p, x, q, y):
-  """Sum-of-diffs ILDJRatio for Chains."""
-  if len(p.bijectors) != len(q.bijectors):
-    raise ValueError('Mismatched lengths of bijectors: `p` has '
-                     f'{len(p.bijectors)} but `q` has {len(q.bijectors)}.')
-  ratios = []
-  max_shp = []
-  for p, q in zip(p.bijectors, q.bijectors):
-    ratios.append(ldj_ratio.inverse_log_det_jacobian_ratio(
-        p, x, q, y, p.inverse_min_event_ndims))
-    max_shp = ps.broadcast_shape(max_shp, ps.shape(ratios[-1]))
-    x, y = p.inverse(x), q.inverse(y)
-  ratios = [tf.broadcast_to(r, max_shp) for r in ratios]
-  return tf.add_n(ratios)

tensorflow_probability/python/bijectors/chain_test.py

@@ -330,6 +330,9 @@ def xform_dynamic(x):
       return tf1.placeholder_with_default(x, shape=None)
 
     def xform_static(x):
+      # Copy the Tensor, because otherwise the set_shape can pass information
+      # into the past.
+      x = tf.identity(x)
       tensorshape_util.set_shape(x, [1])
       return x