Make IncrementLogProb a proper tfd.Distribution.

This enables slicing, returning the distribution from tf.function, etc.

Also, fix a bug with log_prob which didn't handle a callable log prob increment.

PiperOrigin-RevId: 429419375

Author: SiegeLordEx

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_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",