code: def test_log_prob(self, is_static, get_distributions):
      dists, concat_dist = get_distributions(is_static=is_static)
      x_sample = concat_dist.sample(
          [32, 128], seed=test_util.test_seed()
      )
      log_prob = concat_dist.log_prob(x_sample)
      self.assertAllEqual(
          self.evaluate(tf.shape(log_prob)), [32, 128, 2, 6, 4]
      )
      self.assertAllEqual(
          log_prob[:, :, :, :3],
          dists[0].log_prob(x_sample[:, :, :, :3, ...]),
      )
      self.assertAllEqual(
          log_prob[:, :, :, 3:5],
          dists[1].log_prob(x_sample[:, :, :, 3:5, ...]),
      )
      self.assertAllEqual(
          log_prob[:, :, :, 5:],
          dists[2].log_prob(x_sample[:, :, :, 5:, ...]),
      )
  
imports: from absl.testing import parameterized
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import batch_concat
  from tensorflow_probability.python.distributions import dirichlet
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import logistic
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import onehot_categorical
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 9 for calling the BatchConcatTest obj's underlying member method code with example usage: obj.test_log_prob(is_static, get_distributions) without return types