class_docstr: Well tested because this is a simple override of the base class.
code: def testSampleWithBroadcastScale(self):
      # mu corresponds to a 2-batch of 3-variate normals
      mu = np.ones([2, 3]) / 10
  
      # diag corresponds to no batches of 3-variate normals
      diag = np.ones([3]) / 2
  
      dist = mvn_diag.MultivariateNormalDiag(
          mu, diag, validate_args=True
      )
  
      mean = dist.mean()
      self.assertAllEqual([2, 3], mean.shape)
      self.assertAllClose(mu, self.evaluate(mean))
  
      n = int(1e3)
      samps = self.evaluate(dist.sample(n, seed=test_util.test_seed()))
      cov_mat = self.evaluate(tf.linalg.diag(diag)) ** 2
      sample_cov = (
          np.matmul(
              samps.transpose([1, 2, 0]), samps.transpose([1, 0, 2])
          )
          / n
      )
  
      self.assertAllClose(mu, samps.mean(axis=0), atol=0.10, rtol=0.05)
      self.assertAllClose(
          [cov_mat, cov_mat], sample_cov, atol=0.10, rtol=0.05
      )
  
imports: import os
  from absl.testing import parameterized
  import numpy
  from scipy import stats
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.bijectors import softplus
  from tensorflow_probability.python.distributions import kullback_leibler
  from tensorflow_probability.python.distributions import lognormal
  from tensorflow_probability.python.distributions import mvn_diag
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import transformed_distribution
  from tensorflow_probability.python.internal import tensorshape_util
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 10 for calling the MultivariateNormalDiagTest obj's underlying member method code with example usage: obj.testSampleWithBroadcastScale() without return types