code: def test_factored_joint_mvn_broadcast_batch_shape(self):
      # Test that combining MVNs with different but broadcast-compatible
      # batch shapes yields an MVN with the correct broadcast batch shape.
      random_with_shape = lambda shape: np.random.standard_normal(
          shape
      ).astype(np.float32)
  
      event_shape = [3]
      # mvn with batch shape [2]
      mvn1 = mvn_diag.MultivariateNormalDiag(
          loc=random_with_shape([2] + event_shape),
          scale_diag=tf.exp(random_with_shape([2] + event_shape)),
      )
  
      # mvn with batch shape [3, 2]
      mvn2 = mvn_diag.MultivariateNormalDiag(
          loc=random_with_shape([3, 2] + event_shape),
          scale_diag=tf.exp(random_with_shape([1, 2] + event_shape)),
      )
  
      # mvn with batch shape [1, 2]
      mvn3 = mvn_diag.MultivariateNormalDiag(
          loc=random_with_shape([1, 2] + event_shape),
          scale_diag=tf.exp(random_with_shape([2] + event_shape)),
      )
  
      joint = sts_util.factored_joint_mvn([mvn1, mvn2, mvn3])
      self.assertAllEqual(
          self.evaluate(joint.batch_shape_tensor()), [3, 2]
      )
  
      joint_mean_ = self.evaluate(joint.mean())
      broadcast_means = tf.ones_like(joint.mean()[..., 0:1])
      self.assertAllEqual(
          joint_mean_[..., :3],
          self.evaluate(broadcast_means * mvn1.mean()),
      )
      self.assertAllEqual(
          joint_mean_[..., 3:6],
          self.evaluate(broadcast_means * mvn2.mean()),
      )
      self.assertAllEqual(
          joint_mean_[..., 6:9],
          self.evaluate(broadcast_means * mvn3.mean()),
      )
  
      joint_cov_ = self.evaluate(joint.covariance())
      broadcast_covs = tf.ones_like(joint.covariance()[..., :1, :1])
      self.assertAllEqual(
          joint_cov_[..., :3, :3],
          self.evaluate(broadcast_covs * mvn1.covariance()),
      )
      self.assertAllEqual(
          joint_cov_[..., 3:6, 3:6],
          self.evaluate(broadcast_covs * mvn2.covariance()),
      )
      self.assertAllEqual(
          joint_cov_[..., 6:9, 6:9],
          self.evaluate(broadcast_covs * mvn3.covariance()),
      )
  
imports: from absl.testing import parameterized
  import numpy
  import pandas
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import mvn_tril
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.sts.components import local_level
  from tensorflow_probability.python.sts.internal import missing_values_util
  from tensorflow_probability.python.sts.internal import util
label_desc: Use image node_id 2 for calling the MultivariateNormalUtilsTest obj's underlying member method code with example usage: obj.test_factored_joint_mvn_broadcast_batch_shape() without return types