Implement sample_and_log_prob for BatchConcat and BatchReshape.

PiperOrigin-RevId: 375612805

Author: davmre

tensorflow_probability/python/distributions/batch_concat.py

@@ -423,6 +423,20 @@ def _sample_n(self, n, seed=None, **kwargs):
 
     return tf.concat(samples, axis=self._axis+1)
 
+  def _sample_and_log_prob(self, sample_shape, seed, **kwargs):
+    all_seeds = samplers.split_seed(
+        seed, len(self._distributions), salt='BatchConcat')
+    samples = []
+    log_probs = []
+    for d, s in zip(self._distributions, all_seeds):
+      x, lp = d.experimental_sample_and_log_prob(sample_shape, s)
+      samples.append(self._broadcast(x, sample_shape))
+      log_probs.append(self._broadcast(lp, sample_shape))
+
+    sample_shape_size = ps.rank_from_shape(sample_shape)
+    return (tf.concat(samples, axis=self._axis + sample_shape_size),
+            tf.concat(log_probs, axis=self._axis + sample_shape_size))
+
   def _call_split_concat(self, fn, x, **kwargs):
     sample_shape_size, split_x = self._split_sample(x)
     result = [

tensorflow_probability/python/distributions/batch_concat_test.py

@@ -128,6 +128,13 @@ def test_sample(self):
     samples = concat_dist.sample([12, 20], seed=seed)
     self.assertAllEqual(self.evaluate(tf.shape(samples)), [12, 20, 2, 6, 4, 2])
 
+  def test_sample_and_log_prob(self):
+    concat_dist = self.get_distributions()
+    seed = test_util.test_seed()
+    samples, lp = concat_dist.experimental_sample_and_log_prob(seed=seed)
+    self.assertAllEqual(self.evaluate(tf.shape(samples)), [2, 6, 4, 2])
+    self.assertAllClose(lp, concat_dist.log_prob(samples))
+
   def test_split_sample(self):
     concat_dist = self.get_distributions()
     x_sample = tf.ones([2, 6, 4, 2])

tensorflow_probability/python/distributions/batch_reshape.py

@@ -211,6 +211,16 @@ def _sample_n(self, n, seed=None, **kwargs):
         axis=0)
     return tf.reshape(x, new_shape)
 
+  def _sample_and_log_prob(self, sample_shape, seed=None, **kwargs):
+    x, lp = self.distribution.experimental_sample_and_log_prob(
+        sample_shape=sample_shape, seed=seed, **kwargs)
+    return (tf.reshape(x, tf.concat([sample_shape,
+                                     self._batch_shape_unexpanded,
+                                     self.event_shape_tensor()], axis=0)),
+            tf.reshape(lp, tf.concat([sample_shape,
+                                      self._batch_shape_unexpanded],
+                                     axis=0)))
+
   def _log_prob(self, x, **kwargs):
     return self._call_reshape_input_output(
         self.distribution.log_prob, x, extra_kwargs=kwargs)

tensorflow_probability/python/distributions/batch_reshape_test.py

@@ -200,7 +200,7 @@ def test_scalar_variate_sample_and_log_prob(self):
       # document that the test is not intended to run in eager mode.
       return
 
-    seed = test_util.test_seed()
+    seed = test_util.test_seed(sampler_type='stateless')
 
     new_batch_shape = [2, 2]
     old_batch_shape = [4]
@@ -220,27 +220,15 @@ def test_scalar_variate_sample_and_log_prob(self):
     expected_log_prob = tf.reshape(normal.log_prob(x), expected_log_prob_shape)
     actual_log_prob = reshape_normal.log_prob(expected_sample)
 
-    [
-        batch_shape_,
-        event_shape_,
-        expected_sample_,
-        actual_sample_,
-        expected_log_prob_,
-        actual_log_prob_,
-    ] = self.evaluate([
-        batch_shape,
-        event_shape,
-        expected_sample,
-        actual_sample,
-        expected_log_prob,
-        actual_log_prob,
-    ])
-    self.assertAllEqual(new_batch_shape, batch_shape_)
-    self.assertAllEqual([], event_shape_)
-    self.assertAllClose(expected_sample_, actual_sample_,
-                        atol=0., rtol=1e-6)
-    self.assertAllClose(expected_log_prob_, actual_log_prob_,
-                        atol=0., rtol=1e-6)
+    self.assertAllEqual(new_batch_shape, batch_shape)
+    self.assertAllEqual([], event_shape)
+    self.assertAllClose(expected_sample, actual_sample, atol=0., rtol=1e-6)
+    self.assertAllClose(expected_log_prob, actual_log_prob, atol=0., rtol=1e-6)
+
+    slp_sample, slp_lp = reshape_normal.experimental_sample_and_log_prob(
+        seed=seed)
+    self.assertAllClose(expected_sample, slp_sample, atol=0., rtol=1e-6)
+    self.assertAllClose(expected_log_prob, slp_lp, atol=0., rtol=1e-6)
     if not self.is_static_shape:
       return
     self.assertAllEqual(new_batch_shape, reshape_normal.batch_shape)