Implement sample_and_log_prob for several meta distributions.

PiperOrigin-RevId: 375493200

Author: davmre

tensorflow_probability/python/distributions/blockwise.py

@@ -333,6 +333,11 @@ def _sample_n(self, n, seed=None):
     return self._flatten_and_concat_event(
         self._distribution.sample(n, seed=seed))
 
+  def _sample_and_log_prob(self, sample_shape, seed):
+    x, lp = self._distribution.experimental_sample_and_log_prob(
+        sample_shape, seed=seed)
+    return self._flatten_and_concat_event(x), lp
+
   def _log_prob(self, x):
     return self._distribution.log_prob(self._split_and_reshape_event(x))
 

tensorflow_probability/python/distributions/gaussian_process.py

@@ -532,6 +532,15 @@ def _event_shape(self, index_points=None):
   def _sample_n(self, n, seed=None, index_points=None):
     return self.get_marginal_distribution(index_points).sample(n, seed=seed)
 
+  def _sample_and_log_prob(self,
+                           sample_shape,
+                           seed,
+                           index_points=None,
+                           **kwargs):
+    return self.get_marginal_distribution(
+        index_points).experimental_sample_and_log_prob(
+            sample_shape, seed=seed, **kwargs)
+
   def _log_survival_function(self, value, index_points=None):
     return self.get_marginal_distribution(
         index_points).log_survival_function(value)

tensorflow_probability/python/distributions/independent.py

@@ -280,6 +280,11 @@ def _sum_fn(self):
       return lambda x, axis: tfp_math.reduce_kahan_sum(x, axis).total
     return tf.math.reduce_sum
 
+  def _sample_and_log_prob(self, sample_shape, seed, **kwargs):
+    x, lp = self.distribution.experimental_sample_and_log_prob(
+        sample_shape, seed=seed, **kwargs)
+    return x, self._reduce(self._sum_fn(), lp)
+
   def _log_prob(self, x, **kwargs):
     return self._reduce(
         self._sum_fn(), self.distribution.log_prob(x, **kwargs))

tensorflow_probability/python/distributions/matrix_normal_linear_operator.py

@@ -213,6 +213,13 @@ def _sample_n(self, n, seed=None):
     samples = self._as_multivariate_normal(loc=loc).sample(n, seed=seed)
     return _unvec(samples, self._event_shape_tensor(loc=loc))
 
+  def _sample_and_log_prob(self, sample_shape, seed):
+    loc = tf.convert_to_tensor(self.loc)
+    x, lp = self._as_multivariate_normal(
+        loc=loc).experimental_sample_and_log_prob(
+            sample_shape, seed=seed)
+    return _unvec(x, self._event_shape_tensor(loc=loc)), lp
+
   def _entropy(self):
     return self._as_multivariate_normal().entropy()
 

tensorflow_probability/python/distributions/matrix_t_linear_operator.py

@@ -242,6 +242,13 @@ def _sample_n(self, n, seed=None):
     samples = self._as_multivariate_t(loc=loc).sample(n, seed=seed)
     return _unvec(samples, self._event_shape_tensor(loc=loc))
 
+  def _sample_and_log_prob(self, sample_shape, seed):
+    loc = tf.convert_to_tensor(self.loc)
+    x, lp = self._as_multivariate_t(
+        loc=loc).experimental_sample_and_log_prob(
+            sample_shape, seed=seed)
+    return _unvec(x, self._event_shape_tensor(loc=loc)), lp
+
   def _entropy(self):
     return self._as_multivariate_t().entropy()