code: def test_diagonal_mass_matrix_independent(self):
      @tf.function(autograph=False)
      def run(seed):
          dist_seed, *seeds = samplers.split_seed(seed, 11)
          dist = sharded.Sharded(
              independent.Independent(
                  normal.Normal(tf.zeros(3), tf.ones(3)), 1
              ),
              shard_axis_name=self.axis_name,
          )
          state = dist.sample(seed=dist_seed)
          kernel = dmma.DiagonalMassMatrixAdaptation(
              EchoKernel(),
              sample_stats.RunningVariance.from_stats(
                  num_samples=10.0,
                  mean=tf.zeros(3),
                  variance=tf.ones(3),
              ),
          )
          pkr = kernel.bootstrap_results(state)
  
          def body(draw_pkr, i):
              seed = tf.gather(seeds, i)
              _, pkr = draw_pkr
              draw_seed, step_seed = samplers.split_seed(seed)
              draw = dist.sample(seed=draw_seed)
              _, pkr = kernel.one_step(draw, pkr, seed=step_seed)
              return draw, pkr
  
          (_, pkr), draws = loop_util.trace_scan(
              body,
              (tf.zeros(dist.event_shape), pkr),
              tf.range(len(seeds)),
              lambda v: v[0],
          )
  
          return draws, pkr
  
      draws, pkr = self.strategy_run(run, (self.key,), in_axes=None)
      running_variance = self.per_replica_to_composite_tensor(
          pkr.running_variance[0]
      )
      draws = self.per_replica_to_tensor(draws, axis=1)
      mean, sum_squared_residuals, draws = self.evaluate(
          (
              running_variance.mean,
              running_variance.sum_squared_residuals,
              draws,
          )
      )
      emp_mean = tf.reduce_sum(draws, axis=0) / 20.0
      emp_squared_residuals = (
          tf.reduce_sum((draws - emp_mean) ** 2, axis=0)
          + 10 * emp_mean**2
          + 10
      )
      self.assertAllClose(emp_mean, mean)
      self.assertAllClose(emp_squared_residuals, 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 2 for calling the DiagonalAdaptationTest obj's underlying member method code with example usage: obj.test_diagonal_mass_matrix_independent() and returns: draws,  pkr, draw,  pkr