code: def _test_sample_and_log_prob(self, dist_name, dist):
      seed = test_util.test_seed(sampler_type="stateless")
      num_samples = 3
  
      # Sample from the distribution before composite tensoring
      sample1 = self.evaluate(dist.sample(num_samples, seed=seed))
      hp.note("Drew samples {}".format(sample1))
  
      # Sample from the distribution after composite tensoring
      flat = tf.nest.flatten(dist, expand_composites=True)
      unflat = tf.nest.pack_sequence_as(
          dist, flat, expand_composites=True
      )
      sample2 = self.evaluate(unflat.sample(num_samples, seed=seed))
      hp.note("Drew samples {}".format(sample2))
  
      # Check that the samples are the same
      self.assertAllClose(sample1, sample2)
  
      # Check that all the log_probs agree for the samples from before ...
      ct_lp1 = unflat.log_prob(sample1)
      orig_lp1 = dist.log_prob(sample1)
      ct_lp1_, orig_lp1_ = self.evaluate((ct_lp1, orig_lp1))
      self.assertAllClose(
          ct_lp1_,
          orig_lp1_,
          rtol=COMPOSITE_TENSOR_LOGPROB_RTOL[dist_name],
      )
  
      # ... and after.  (Even though they're supposed to be the same anyway.)
      ct_lp2 = unflat.log_prob(sample2)
      orig_lp2 = dist.log_prob(sample2)
      ct_lp2_, orig_lp2_ = self.evaluate((ct_lp2, orig_lp2))
      self.assertAllClose(
          ct_lp2_,
          orig_lp2_,
          rtol=COMPOSITE_TENSOR_LOGPROB_RTOL[dist_name],
      )
  
imports: import collections
  from absl.testing import parameterized
  import hypothesis
  from hypothesis import strategies
  import numpy
  import six
  import tensorflow.compat.v2
  from tensorflow_probability.python import distributions
  from tensorflow_probability.python.distributions import hypothesis_testlib
  from tensorflow_probability.python.internal import hypothesis_testlib
  from tensorflow_probability.python.internal import tensor_util
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the DistributionCompositeTensorTest obj's underlying member method code with example usage: obj._test_sample_and_log_prob(dist_name, dist) without return types