docstr: Nothing distributed. Make sure EchoKernel works.
docstr_tok: ["Nothing","distributed",".","Make","sure","EchoKernel","works","."]
code: def test_diagonal_mass_matrix_no_distribute(self):
      """Nothing distributed. Make sure EchoKernel works."""
      kernel = dmma.DiagonalMassMatrixAdaptation(
          EchoKernel(),
          sample_stats.RunningVariance.from_stats(
              num_samples=10.0, mean=tf.zeros(3), variance=tf.ones(3)
          ),
      )
      state = tf.zeros(3)
      pkr = kernel.bootstrap_results(state)
      draws = np.random.randn(10, 3).astype(np.float32)
  
      def body(pkr_seed, draw):
          pkr, seed = pkr_seed
          seed, kernel_seed = samplers.split_seed(seed)
          _, pkr = kernel.one_step(draw, pkr, seed=kernel_seed)
          return (pkr, seed)
  
      (pkr, _), _ = loop_util.trace_scan(
          body,
          (pkr, samplers.sanitize_seed(self.key)),
          draws,
          lambda _: (),
      )
  
      running_variance = pkr.running_variance[0]
      emp_mean = draws.sum(axis=0) / 20.0
      emp_squared_residuals = (
          np.sum((draws - emp_mean) ** 2, axis=0)
          + 10 * emp_mean**2
          + 10
      )
      self.assertAllClose(emp_mean, running_variance.mean)
      self.assertAllClose(
          emp_squared_residuals, running_variance.sum_squared_residuals
      )
  
imports: import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.distributions import independent
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.distributions import sample
  from tensorflow_probability.python.experimental.distribute import sharded
  from tensorflow_probability.python.experimental.mcmc import diagonal_mass_matrix_adaptation
  from tensorflow_probability.python.experimental.mcmc import preconditioned_hmc
  from tensorflow_probability.python.experimental.stats import sample_stats
  from tensorflow_probability.python.internal import distribute_test_lib
  from tensorflow_probability.python.internal import loop_util
  from tensorflow_probability.python.internal import samplers
  from tensorflow_probability.python.internal import test_util
label_desc: Use image node_id 1 for calling the DiagonalAdaptationTest obj's underlying member method code with example usage: obj.test_diagonal_mass_matrix_no_distribute() and returns: pkr,  seed