Merge branch 'tensorflow:main' into frighterafix#1505

Author: Frightera

tensorflow_probability/python/bijectors/bijector.py

@@ -1370,7 +1370,7 @@ def forward(self, x, name='forward', **kwargs):
   @classmethod
   def _is_increasing(cls, **kwargs):
     """Subclass implementation for `is_increasing` public function."""
-    raise NotImplementedError('`_is_increasing` not implemented.')
+    raise NotImplementedError(f'`_is_increasing` not implemented in {cls}.')
 
   def _call_is_increasing(self, name, **kwargs):
     """Wraps call to _is_increasing, allowing extra shared logic."""
@@ -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/bijectors/ffjord.py

@@ -356,7 +356,7 @@ def _solve_ode(self, ode_fn, state, **kwargs):
   def _augmented_forward(self, x, **condition_kwargs):
     """Computes forward and forward_log_det_jacobian transformations."""
     augmented_ode_fn = self._trace_augmentation_fn(
-        self._state_time_derivative_fn, x.shape, x.dtype)
+        self._state_time_derivative_fn, prefer_static.shape(x), x.dtype)
     augmented_x = (x, tf.zeros_like(x))
     if condition_kwargs:
       y, fldj = self._solve_ode(augmented_ode_fn, augmented_x,
@@ -368,7 +368,7 @@ def _augmented_forward(self, x, **condition_kwargs):
   def _augmented_inverse(self, y, **condition_kwargs):
     """Computes inverse and inverse_log_det_jacobian transformations."""
     augmented_inv_ode_fn = self._trace_augmentation_fn(
-        self._inv_state_time_derivative_fn, y.shape, y.dtype)
+        self._inv_state_time_derivative_fn, prefer_static.shape(y), y.dtype)
     augmented_y = (y, tf.zeros_like(y))
     if condition_kwargs:
       x, ildj = self._solve_ode(augmented_inv_ode_fn, augmented_y,

tensorflow_probability/python/distributions/BUILD

@@ -934,6 +934,7 @@ multi_substrate_py_library(
         # tensorflow dep,
         "//tensorflow_probability",
         "//tensorflow_probability/python/bijectors:hypothesis_testlib",
+        "//tensorflow_probability/python/experimental/distributions",
         "//tensorflow_probability/python/internal:hypothesis_testlib",
         "//tensorflow_probability/python/internal:tensorshape_util",
     ],

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/distributions/distribution_properties_test.py

@@ -58,7 +58,8 @@
 INSTANTIABLE_BUT_NOT_SLICABLE = (
     'BatchBroadcast',
     'BatchReshape',
-    'Sample'  # TODO(b/204210361)
+    'IncrementLogProb',
+    'Sample',  # TODO(b/204210361)
 )
 
 
@@ -350,6 +351,7 @@ def testCanConstructAndSampleDistribution(self, data):
                                     'TruncatedNormal', 'Uniform')
     not_annotated_dists = ('Empirical|event_ndims=0', 'Empirical|event_ndims=1',
                            'Empirical|event_ndims=2', 'FiniteDiscrete',
+                           'IncrementLogProb',
                            # cov_perturb_factor is not annotated since its shape
                            # could be a vector or a matrix.
                            'MultivariateNormalDiagPlusLowRankCovariance',

tensorflow_probability/python/distributions/hypothesis_testlib.py

@@ -29,6 +29,7 @@
 from tensorflow_probability.python import distributions as tfd
 from tensorflow_probability.python import util as tfp_util
 from tensorflow_probability.python.bijectors import hypothesis_testlib as bijector_hps
+from tensorflow_probability.python.experimental import distributions as tfed
 from tensorflow_probability.python.internal import hypothesis_testlib as tfp_hps
 from tensorflow_probability.python.internal import tensorshape_util
 
@@ -448,6 +449,10 @@ def _instantiable_base_dists():
       functools.partial(tfd.Empirical, event_ndims=1), dict(samples=2))
   result['Empirical|event_ndims=2'] = DistInfo(  #
       functools.partial(tfd.Empirical, event_ndims=2), dict(samples=3))
+
+  # We use a special strategy for instantiating this, so event_dims is set to a
+  # dummy value.
+  result['IncrementLogProb'] = DistInfo(tfed.IncrementLogProb, None)
   return result
 
 
@@ -680,6 +685,12 @@ class names so they will not be drawn at the top level.
     return draw(spherical_uniforms(
         batch_shape=batch_shape, event_dim=event_dim,
         validate_args=validate_args))
+  elif dist_name == 'IncrementLogProb':
+    return draw(
+        increment_log_probs(
+            batch_shape=batch_shape,
+            enable_vars=enable_vars,
+            validate_args=validate_args))
 
   if params is None:
     params_unconstrained, batch_shape = draw(
@@ -738,6 +749,45 @@ def spherical_uniforms(
   return result_dist
 
 
+@hps.composite
+def increment_log_probs(draw,
+                        batch_shape=None,
+                        enable_vars=False,
+                        validate_args=True):
+  """Strategy for drawing `IncrementLogProb` distributions.
+
+  Args:
+    draw: Hypothesis strategy sampler supplied by `@hps.composite`.
+    batch_shape: An optional `TensorShape`.  The batch shape of the resulting
+      `IncrementLogProb` distribution.
+    enable_vars: TODO(b/181859346): Make this `True` all the time and put
+      variable initialization in slicing_test.  If `False`, the returned
+      parameters are all `tf.Tensor`s and not {`tf.Variable`,
+      `tfp.util.DeferredTensor` `tfp.util.TransformedVariable`}
+    validate_args: Python `bool`; whether to enable runtime assertions.
+
+  Returns:
+    dists: A strategy for drawing `IncrementLogProb` distributions with the
+      specified `batch_shape` (or an arbitrary one if omitted).
+  """
+  if batch_shape is None:
+    batch_shape = draw(tfp_hps.shapes(min_ndims=0, max_side=4))
+
+  log_prob_value = draw(
+      tfp_hps.maybe_variable(
+          tfp_hps.constrained_tensors(tfp_hps.identity_fn,
+                                      tensorshape_util.as_list(batch_shape)),
+          enable_vars=enable_vars))
+
+  if draw(hps.booleans()):
+    return tfed.IncrementLogProb(log_prob_value, validate_args=validate_args)
+  else:
+    return tfed.IncrementLogProb(
+        lambda v: v,
+        validate_args=validate_args,
+        log_prob_increment_kwargs={'v': log_prob_value})
+
+
 @hps.composite
 def batch_broadcasts(
     draw, batch_shape=None, event_dim=None,

tensorflow_probability/python/distributions/jax_transformation_test.py

@@ -82,6 +82,7 @@
 JVP_LOGPROB_SAMPLE_BLOCKLIST = (
     'BetaQuotient',  # https://b/178552958
     'GeneralizedExtremeValue',  # http://b/175654800
+    'IncrementLogProb',  # Sample and log prob are independent.
     'NegativeBinomial',  # Too slow: http://b/170871051
 )
 JVP_LOGPROB_PARAM_BLOCKLIST = (
@@ -95,6 +96,7 @@
 VJP_LOGPROB_SAMPLE_BLOCKLIST = (
     'BetaQuotient',  # https://b/178552958
     'GeneralizedExtremeValue',  # http://b/175654800
+    'IncrementLogProb',  # Sample and log prob are independent.
     'NegativeBinomial',  # Too slow: http://b/170871051
 )
 VJP_LOGPROB_PARAM_BLOCKLIST = (

tensorflow_probability/python/experimental/distributions/BUILD

@@ -78,6 +78,7 @@ multi_substrate_py_library(
         # numpy dep,
         # tensorflow dep,
         "//tensorflow_probability/python/bijectors:identity",
+        "//tensorflow_probability/python/distributions:distribution",
         "//tensorflow_probability/python/internal:callable_util",
         "//tensorflow_probability/python/internal:prefer_static",
         "//tensorflow_probability/python/internal:reparameterization",