Forward kwargs into tangent space methods.

And fix how we check whether the tangent space has been specified already.

PiperOrigin-RevId: 429622219

Author: Googler

tensorflow_probability/python/bijectors/bijector.py

@@ -1684,7 +1684,8 @@ def forward_log_det_jacobian(self,
   def experimental_compute_density_correction(self,
                                               x,
                                               tangent_space,
-                                              backward_compat=False):
+                                              backward_compat=False,
+                                              **kwargs):
     """Density correction for this transformation wrt the tangent space, at x.
 
     Subclasses of Bijector may call the most specific applicable
@@ -1699,6 +1700,7 @@ def experimental_compute_density_correction(self,
         the support manifold at `x`.
       backward_compat: `bool` specifying whether to assume that the Bijector
         is dimension-preserving.
+      **kwargs: Optional keyword arguments forwarded to tangent space methods.
 
     Returns:
       density_correction: `Tensor` representing the density correction---in log
@@ -1710,7 +1712,7 @@ def experimental_compute_density_correction(self,
 
     """
     if backward_compat:
-      return tangent_space.transform_dimension_preserving(x, self)
+      return tangent_space.transform_dimension_preserving(x, self, **kwargs)
     else:
       raise TypeError(
           'Please call the `TangentSpace` method applicable to this Bijector.')

tensorflow_probability/python/distributions/distribution.py

@@ -1926,7 +1926,7 @@ def experimental_local_measure(self, value, backward_compat=False, **kwargs):
     """
     log_prob = self.log_prob(value, **kwargs)
     tangent_space = None
-    if getattr(self, '_experimental_tangent_space'):
+    if hasattr(self, '_experimental_tangent_space'):
       tangent_space = self._experimental_tangent_space
     elif backward_compat:
       # Import here rather than top-level to avoid circular import.

tensorflow_probability/python/experimental/tangent_spaces/spaces.py

@@ -49,7 +49,7 @@ class TangentSpace(object):
 
   """
 
-  def transform_general(self, x, f):
+  def transform_general(self, x, f, **kwargs):
     """Returns the density correction, in log space, corresponding to f at x.
 
     Also returns a new `TangentSpace` representing the tangent to fM at f(x).
@@ -58,6 +58,7 @@ def transform_general(self, x, f):
       x: `Tensor` (structure). The point at which to calculate the density.
       f: `Bijector` or one of its subclasses. The transformation that requires a
         density correction based on this tangent space.
+      **kwargs: Optional keyword arguments as part of the Bijector.
 
     Returns:
       log_density: A `Tensor` representing the log density correction of f at x
@@ -69,7 +70,7 @@ def transform_general(self, x, f):
     """
     raise NotImplementedError
 
-  def transform_dimension_preserving(self, x, f):
+  def transform_dimension_preserving(self, x, f, **kwargs):
     """Same as `transform_general`, assuming f goes from R^n to R^n.
 
     Default falls back to `transform_general`, which may be overridden
@@ -78,6 +79,7 @@ def transform_dimension_preserving(self, x, f):
     Args:
       x: same as in `transform_general`.
       f: same as in `transform_general`.
+      **kwargs: same as in `transform_general`.
 
     Returns:
       log_density: A `Tensor` representing the log density correction of f at x
@@ -88,9 +90,9 @@ def transform_dimension_preserving(self, x, f):
         `transform_general`.
 
     """
-    return self.transform_general(x, f)
+    return self.transform_general(x, f, **kwargs)
 
-  def transform_projection(self, x, f):
+  def transform_projection(self, x, f, **kwargs):
     """Same as `transform_general`, with f a projection (or its inverse).
 
     Default falls back to `transform_general`, which may be overridden
@@ -99,6 +101,7 @@ def transform_projection(self, x, f):
     Args:
       x: same as in `transform_general`.
       f: same as in `transform_general`.
+      **kwargs: same as in `transform_general`.
 
     Returns:
       log_density: A `Tensor` representing the log density correction of f at x
@@ -108,9 +111,9 @@ def transform_projection(self, x, f):
       NotImplementedError: if the `TangentSpace` subclass does not implement
         `transform_general`.
     """
-    return self.transform_general(x, f)
+    return self.transform_general(x, f, **kwargs)
 
-  def transform_coordinatewise(self, x, f):
+  def transform_coordinatewise(self, x, f, **kwargs):
     """Same as `transform_dimension_preserving`, for a coordinatewise f.
 
     Default falls back to `transform_dimension_preserving`, which may
@@ -119,6 +122,7 @@ def transform_coordinatewise(self, x, f):
     Args:
       x: same as in `transform_dimension_preserving`.
       f: same as in `transform_dimension_preserving`.
+      **kwargs: same as in `transform_dimension_preserving`.
 
     Returns:
       log_density: A `Tensor` representing the log density correction of f at x
@@ -129,7 +133,7 @@ def transform_coordinatewise(self, x, f):
         `transform_dimension_preserving`.
 
     """
-    return self.transform_dimension_preserving(x, f)
+    return self.transform_dimension_preserving(x, f, **kwargs)
 
 
 def unit_basis():
@@ -161,31 +165,32 @@ def __init__(self, axis_mask):
     """
     self.axis_mask = axis_mask
 
-  def transform_general(self, x, f):
+  def transform_general(self, x, f, **kwargs):
     as_general_space = GeneralSpace(unit_basis_on(self.axis_mask), 1)
-    return as_general_space.transform_general(x, f)
+    return as_general_space.transform_general(x, f, **kwargs)
 
-  def transform_projection(self, x, f):
+  def transform_projection(self, x, f, **kwargs):
     if not hasattr(f, 'experimental_update_live_dimensions'):
       msg = ('When calling `transform_projection` the Bijector must implement '
              'the `experimental_update_live_dimensions` method.')
       raise NotImplementedError(msg)
-    new_live_dimensions = f.experimental_update_live_dimensions(self.axis_mask)
+    new_live_dimensions = f.experimental_update_live_dimensions(
+        self.axis_mask, **kwargs)
     if all(tf.get_static_value(new_live_dimensions)):
       # Special-case a bijector (direction) that knows that the result
       # of the projection will be a full space
       return 0, FullSpace()
     else:
       return 0, AxisAlignedSpace(new_live_dimensions)
 
-  def transform_coordinatewise(self, x, f):
+  def transform_coordinatewise(self, x, f, **kwargs):
     # TODO(pravnar): compute the derivative of f along x along the
     # live dimensions.
     raise NotImplementedError
 
 
-def jacobian_determinant(x, f):
-  return f.forward_log_det_jacobian(x)
+def jacobian_determinant(x, f, **kwargs):
+  return f.forward_log_det_jacobian(x, **kwargs)
 
 
 class FullSpace(TangentSpace):
@@ -197,16 +202,17 @@ class FullSpace(TangentSpace):
   at all.
   """
 
-  def transform_general(self, x, f):
+  def transform_general(self, x, f, **kwargs):
     """If the bijector is weird, fall back to the general case."""
     as_general_space = GeneralSpace(unit_basis(), 1)
-    return as_general_space.transform_general(x, f)
+    return as_general_space.transform_general(x, f, **kwargs)
 
-  def transform_dimension_preserving(self, x, f):
-    return jacobian_determinant(x, f), FullSpace()
+  def transform_dimension_preserving(self, x, f, **kwargs):
+    return jacobian_determinant(x, f, **kwargs), FullSpace()
 
-  def transform_projection(self, x, f):
-    return AxisAlignedSpace(tf.ones_like(x)).transform_projection(x, f)
+  def transform_projection(self, x, f, **kwargs):
+    return AxisAlignedSpace(tf.ones_like(x)).transform_projection(
+        x, f, **kwargs)
 
 
 def volume_coefficient(basis):
@@ -223,7 +229,7 @@ def __init__(self, basis, computed_volume=None):
       computed_volume = volume_coefficient(basis)
     self.volume = computed_volume
 
-  def transform_general(self, x, f):
+  def transform_general(self, x, f, **kwargs):
     raise NotImplementedError
 
 
@@ -236,7 +242,7 @@ class ZeroSpace(TangentSpace):
 
   """
 
-  def transform_general(self, x, f):
+  def transform_general(self, x, f, **kwargs):
     del x, f
     return 0, ZeroSpace()