Implement sample_and_log_prob for joint distributions.

PiperOrigin-RevId: 375534806

Author: davmre

tensorflow_probability/python/distributions/joint_distribution.py

@@ -98,6 +98,15 @@ def trace_values_only(dist, sample_shape, seed, value=None):
   return ValueWithTrace(value=value, traced=value)
 
 
+def trace_values_and_log_probs(dist, sample_shape, seed, value=None):
+  """Draws a sample, and traces both the sampled value and its log density."""
+  if value is None:
+    value, lp = dist.experimental_sample_and_log_prob(sample_shape, seed=seed)
+  else:
+    lp = dist.log_prob(value)
+  return ValueWithTrace(value=value, traced=(value, lp))
+
+
 CALLING_CONVENTION_DESCRIPTION = """
 The measure methods of `JointDistribution` (`log_prob`, `prob`, etc.)
 can be called either by passing a single structure of tensors or by using
@@ -570,6 +579,19 @@ def _sample_n(self, sample_shape, seed, value=None):
 
     return self._model_unflatten(xs)
 
+  # TODO(b/189122177): Implement _sample_and_log_prob for distributed JDs.
+  def _sample_and_log_prob(self, sample_shape, seed, value=None, **kwargs):
+    xs, lps = zip(
+        *self._call_execute_model(
+            sample_shape,
+            seed=seed,
+            value=self._resolve_value(value=value,
+                                      allow_partially_specified=True,
+                                      **kwargs),
+            sample_and_trace_fn=trace_values_and_log_probs))
+    return (self._model_unflatten(xs),
+            sum(maybe_check_wont_broadcast(lps, self.validate_args)))
+
   def _map_measure_over_dists(self, attr, value):
     if any(x is None for x in tf.nest.flatten(value)):
       raise ValueError('No `value` part can be `None`; saw: {}.'.format(value))

tensorflow_probability/python/distributions/joint_distribution_auto_batched_test.py

@@ -334,6 +334,66 @@ def coroutine_model():
 
     self.assertAllClose(*self.evaluate([log_prob, expected_log_prob]))
 
+  @parameterized.named_parameters(
+      {'testcase_name': 'coroutine',
+       'jd_class': tfd.JointDistributionCoroutineAutoBatched},
+      {'testcase_name': 'sequential',
+       'jd_class': tfd.JointDistributionSequentialAutoBatched},
+      {'testcase_name': 'named',
+       'jd_class': tfd.JointDistributionNamedAutoBatched})
+  def test_sample_and_log_prob(self, jd_class):
+
+    # Define a bijector to detect if/when `inverse` is called.
+    inverted_values = []
+
+    class InverseTracingExp(tfb.Exp):
+
+      def _inverse(self, y):
+        inverted_values.append(y)
+        return tf.math.log(y)
+
+    models = {}
+
+    def coroutine_model():
+      g = yield InverseTracingExp()(tfd.Normal(0., 1.), name='g')
+      df = yield tfd.Exponential(1., name='df')
+      loc = yield tfd.Sample(tfd.Normal(0, g), 20, name='loc')
+      yield tfd.StudentT(df, loc, 1, name='x')
+    models[tfd.JointDistributionCoroutineAutoBatched] = coroutine_model
+
+    models[tfd.JointDistributionSequentialAutoBatched] = [
+        InverseTracingExp()(tfd.Normal(0., 1.), name='g'),
+        tfd.Exponential(1., name='df'),
+        lambda _, g: tfd.Sample(tfd.Normal(0, g), 20, name='loc'),
+        lambda loc, df: tfd.StudentT(df, loc, 1, name='x')
+    ]
+
+    models[tfd.JointDistributionNamedAutoBatched] = collections.OrderedDict((
+        ('g', InverseTracingExp()(tfd.Normal(0., 1.))),
+        ('df', tfd.Exponential(1.)),
+        ('loc', lambda g: tfd.Sample(tfd.Normal(0, g), 20)),
+        ('x', lambda loc, df: tfd.StudentT(df, loc, 1))))
+
+    joint = jd_class(models[jd_class], validate_args=True)
+
+    for sample_shape in ([], [5]):
+      inverted_values.clear()
+      x1, lp1 = self.evaluate(
+          joint.experimental_sample_and_log_prob(
+              sample_shape,
+              seed=test_util.test_seed(sampler_type='seedless'),
+              df=2.7))  # Check that kwargs are supported.
+      x2 = self.evaluate(
+          joint.sample(sample_shape,
+                       seed=test_util.test_seed(sampler_type='seedless'),
+                       df=2.7))
+      self.assertAllCloseNested(x1, x2)
+
+      self.assertLen(inverted_values, 0)
+      lp2 = joint.log_prob(x1)
+      self.assertLen(inverted_values, 1)
+      self.assertAllClose(lp1, lp2)
+
   def test_sample_with_batch_value(self):
     @tfd.JointDistributionCoroutineAutoBatched
     def dist():

tensorflow_probability/python/distributions/joint_distribution_coroutine_test.py

@@ -334,6 +334,45 @@ def dist():
 
     self.assertAllClose(*self.evaluate([log_prob, expected_log_prob]))
 
+  def test_sample_and_log_prob(self):
+
+    # Define a bijector to detect if/when `inverse` is called.
+    inverted_values = []
+
+    class InverseTracingExp(tfb.Exp):
+
+      def _inverse(self, y):
+        inverted_values.append(y)
+        return tf.math.log(y)
+
+    def coroutine_model():
+      g = yield Root(InverseTracingExp()(tfd.Normal(0., 1.), name='g'))
+      df = yield Root(tfd.Exponential(1., name='df'))
+      loc = yield tfd.Sample(tfd.Normal(0, g), 20, name='loc')
+      yield tfd.Independent(tfd.StudentT(df[..., tf.newaxis], loc, 1, name='x'),
+                            reinterpreted_batch_ndims=1)
+
+    joint = tfd.JointDistributionCoroutine(coroutine_model, validate_args=True)
+
+    for sample_shape in ([], [5]):
+      inverted_values.clear()
+      x1, lp1 = self.evaluate(
+          joint.experimental_sample_and_log_prob(
+              sample_shape,
+              seed=test_util.test_seed(sampler_type='seedless'),
+              df=2.7 * tf.ones(sample_shape)  # Check that kwargs are supported.
+              ))
+      x2 = self.evaluate(
+          joint.sample(sample_shape,
+                       seed=test_util.test_seed(sampler_type='seedless'),
+                       df=2.7 * tf.ones(sample_shape)))
+      self.assertAllCloseNested(x1, x2)
+
+      self.assertLen(inverted_values, 0)
+      lp2 = joint.log_prob(x1)
+      self.assertLen(inverted_values, 1)
+      self.assertAllClose(lp1, lp2)
+
   def test_detect_missing_root(self):
     if not tf.executing_eagerly(): return
     # The joint distribution specified below is intended to

tensorflow_probability/python/distributions/joint_distribution_sample_path_mixin.py

@@ -23,6 +23,7 @@
 import tensorflow.compat.v2 as tf
 
 from tensorflow_probability.python import math as tfp_math
+from tensorflow_probability.python.distributions import joint_distribution as jd_lib
 from tensorflow_probability.python.internal import assert_util
 from tensorflow_probability.python.internal import prefer_static
 
@@ -123,9 +124,7 @@ def event_shape_tensor(self, sample_shape=(), name='event_shape_tensor'):
             d.event_shape_tensor()))
       return self._model_unflatten(component_shapes)
 
-  def _map_and_reduce_measure_over_dists(self, attr, reduce_fn, value):
-    """Reduces all non-batch dimensions of the provided measure."""
-    xs = list(self._map_measure_over_dists(attr, value))
+  def _reduce_measure_over_dists(self, xs, reduce_fn):
     num_trailing_batch_dims_treated_as_event = [
         prefer_static.rank_from_shape(
             d.batch_shape_tensor()) - self._batch_ndims
@@ -145,14 +144,31 @@ def _maybe_check_batch_shape(self):
             parts[0], s, message='Component batch shapes are inconsistent.'))
     return assertions
 
-  def _log_prob(self, value):
+  def _reduce_log_probs_over_dists(self, lps):
     if self._experimental_use_kahan_sum:
-      xs = self._map_and_reduce_measure_over_dists(
-          'log_prob', tfp_math.reduce_kahan_sum, value)
-      return sum(xs).total
-    xs = self._map_and_reduce_measure_over_dists(
-        'log_prob', tf.reduce_sum, value)
-    return sum(xs)
+      return sum(jd_lib.maybe_check_wont_broadcast(
+          self._reduce_measure_over_dists(
+              lps, reduce_fn=tfp_math.reduce_kahan_sum),
+          self.validate_args)).total
+    else:
+      return sum(jd_lib.maybe_check_wont_broadcast(
+          self._reduce_measure_over_dists(lps, reduce_fn=tf.reduce_sum),
+          self.validate_args))
+
+  def _sample_and_log_prob(self, sample_shape, seed, value=None, **kwargs):
+    xs, lps = zip(
+        *self._call_execute_model(
+            sample_shape,
+            seed=seed,
+            value=self._resolve_value(value=value,
+                                      allow_partially_specified=True,
+                                      **kwargs),
+            sample_and_trace_fn=jd_lib.trace_values_and_log_probs))
+    return self._model_unflatten(xs), self._reduce_log_probs_over_dists(lps)
+
+  def _log_prob(self, value):
+    return self._reduce_log_probs_over_dists(
+        self._map_measure_over_dists('log_prob', value))
 
   def log_prob_parts(self, value, name='log_prob_parts'):
     """Log probability density/mass function, part-wise.
@@ -172,18 +188,14 @@ def log_prob_parts(self, value, name='log_prob_parts'):
       sum_fn = tf.reduce_sum
       if self._experimental_use_kahan_sum:
         sum_fn = lambda x, axis: tfp_math.reduce_kahan_sum(x, axis=axis).total
-      xs = self._map_and_reduce_measure_over_dists(
-          'log_prob', sum_fn, value)
-      return self._model_unflatten(xs)
+      return self._model_unflatten(
+          self._reduce_measure_over_dists(
+              self._map_measure_over_dists('log_prob', value),
+              sum_fn))
 
   def _unnormalized_log_prob(self, value):
-    if self._experimental_use_kahan_sum:
-      xs = self._map_and_reduce_measure_over_dists(
-          'unnormalized_log_prob', tfp_math.reduce_kahan_sum, value)
-      return sum(xs).total
-    xs = self._map_and_reduce_measure_over_dists(
-        'unnormalized_log_prob', tf.reduce_sum, value)
-    return sum(xs)
+    return self._reduce_log_probs_over_dists(
+        self._map_measure_over_dists('unnormalized_log_prob', value))
 
   def unnormalized_log_prob_parts(self, value, name=None):
     """Unnormalized log probability density/mass function, part-wise.
@@ -203,9 +215,10 @@ def unnormalized_log_prob_parts(self, value, name=None):
       sum_fn = tf.reduce_sum
       if self._experimental_use_kahan_sum:
         sum_fn = lambda x, axis: tfp_math.reduce_kahan_sum(x, axis=axis).total
-      xs = self._map_and_reduce_measure_over_dists(
-          'unnormalized_log_prob', sum_fn, value)
-      return self._model_unflatten(xs)
+      return self._model_unflatten(
+          self._reduce_measure_over_dists(
+              self._map_measure_over_dists('unnormalized_log_prob', value),
+              sum_fn))
 
   def prob_parts(self, value, name='prob_parts'):
     """Log probability density/mass function.
@@ -221,9 +234,10 @@ def prob_parts(self, value, name='prob_parts'):
         each `distribution_fn` evaluated at each corresponding `value`.
     """
     with self._name_and_control_scope(name):
-      xs = self._map_and_reduce_measure_over_dists(
-          'prob', tf.reduce_prod, value)
-      return self._model_unflatten(xs)
+      return self._model_unflatten(
+          self._reduce_measure_over_dists(
+              self._map_measure_over_dists('prob', value),
+              tf.reduce_prod))
 
   def is_scalar_batch(self, name='is_scalar_batch'):
     """Indicates that `batch_shape == []`.