docstr: Checks that basic properties work with JointDistribution.
docstr_tok: ["Checks","that","basic","properties","work","with","JointDistribution","."]
code: def testJointDistribution(self, dist_fn, nums_elements, batch_shape):
      """Checks that basic properties work with JointDistribution."""
      base = dist_fn(self)
      num_elements = sum(nums_elements)
  
      flat = blockwise.Blockwise(base, validate_args=True)
      if self.use_static_shape:
          self.assertAllEqual([num_elements], flat.event_shape)
      self.assertAllEqual(
          [num_elements], self.evaluate(flat.event_shape_tensor())
      )
      if self.use_static_shape:
          self.assertAllEqual(batch_shape, flat.batch_shape)
      self.assertAllEqual(
          batch_shape, self.evaluate(flat.batch_shape_tensor())
      )
  
      base_sample = self.evaluate(
          base.sample(3, seed=test_util.test_seed())
      )
      base_sample_list = tf.nest.flatten(base_sample)
      flat_sample = self.evaluate(
          flat.sample(3, seed=test_util.test_seed())
      )
      self.assertAllEqual(
          [3] + batch_shape + [num_elements], flat_sample.shape
      )
  
      split_points = np.cumsum([0] + nums_elements)
      base_sample_list = [
          flat_sample[..., start:end].reshape(base_sample_part.shape)
          for start, end, base_sample_part in zip(
              split_points[:-1], split_points[1:], base_sample_list
          )
      ]
  
      base_sample = tf.nest.pack_sequence_as(
          base_sample, base_sample_list
      )
  
      base_log_prob = self.evaluate(base.log_prob(base_sample))
      flat_log_prob = self.evaluate(flat.log_prob(flat_sample))
      self.assertAllEqual([3] + batch_shape, flat_log_prob.shape)
      self.assertAllClose(base_log_prob, flat_log_prob)
  
imports: from absl.testing import parameterized
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import blockwise
  from tensorflow_probability.python.distributions import exponential
  from tensorflow_probability.python.distributions import gamma
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import joint_distribution_coroutine
  from tensorflow_probability.python.distributions import joint_distribution_named
  from tensorflow_probability.python.distributions import joint_distribution_sequential
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import mvn_linear_operator
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 4 for calling the BlockwiseTestStaticParams obj's underlying member method code with example usage: obj.testJointDistribution(dist_fn, nums_elements, batch_shape) without return types