code: def testHMCWithStateGrads(self):
      trajectory_length = 1.0
      epsilon = 1e-3
  
      seed = self._make_seed(_test_seed())
  
      def hmc_step(trajectory_length, axis_name=()):
          @tfp.experimental.distribute.JointDistributionCoroutine
          def model():
              z = yield Root(tfd.Normal(0.0, 1))
              yield tfp.experimental.distribute.Sharded(
                  tfd.Sample(tfd.Normal(z, 1.0), 8), axis_name
              )
  
          @tfp.experimental.distribute.JointDistributionCoroutine
          def momentum_dist():
              yield Root(tfd.Normal(0.0, 2))
              yield Root(
                  tfp.experimental.distribute.Sharded(
                      tfd.Sample(tfd.Normal(0.0, 3.0), 8), axis_name
                  )
              )
  
          def target_log_prob_fn(x):
              return model.log_prob(x), ()
  
          def kinetic_energy_fn(m):
              return -momentum_dist.log_prob(m), ()
  
          def momentum_sample_fn(seed):
              return momentum_dist.sample(2, seed=seed)
  
          state = model.sample(2, seed=seed)
          hmc_state = fun_mc.hamiltonian_monte_carlo_init(
              state, target_log_prob_fn
          )
          (
              hmc_state,
              hmc_extra,
          ) = sga_hmc.hamiltonian_monte_carlo_with_state_grads_step(
              hmc_state,
              trajectory_length=trajectory_length,
              scalar_step_size=epsilon,
              step_size_scale=util.map_tree(
                  lambda x: 1.0 + jnp.abs(x), state
              ),
              target_log_prob_fn=target_log_prob_fn,
              seed=seed,
              kinetic_energy_fn=kinetic_energy_fn,
              momentum_sample_fn=momentum_sample_fn,
              named_axis=model.experimental_shard_axis_names,
          )
  
          def sum_state(x, axis_name):
              res = jnp.sum(x**2)
              if axis_name:
                  res = backend.distribute_lib.psum(res, axis_name)
              return res
  
          sum_sq = util.map_tree_up_to(
              hmc_extra.proposed_state,
              sum_state,
              hmc_extra.proposed_state,
              model.experimental_shard_axis_names,
          )
          sum_sq = sum(util.flatten_tree(sum_sq))
          return sum_sq, ()
  
      def finite_diff_grad(f, epsilon, x):
          return (
              fun_mc.call_potential_fn(
                  f, util.map_tree(lambda x: x + epsilon, x)
              )[0]
              - fun_mc.call_potential_fn(
                  f, util.map_tree(lambda x: x - epsilon, x)
              )[0]
          ) / (2 * epsilon)
  
      f = jax.jit(hmc_step)
      auto_diff = util.value_and_grad(f, trajectory_length)[2]
      finite_diff = finite_diff_grad(f, epsilon, trajectory_length)
  
      self.assertAllClose(auto_diff, finite_diff, rtol=0.01)
  
      if BACKEND == "backend_jax":
  
          @functools.partial(jax.pmap, axis_name="i")
          def run(_):
              f = jax.jit(
                  lambda trajectory_length: hmc_step(  # pylint: disable=g-long-lambda
                      trajectory_length, axis_name="i"
                  )
              )
              auto_diff = util.value_and_grad(f, trajectory_length)[2]
              finite_diff = finite_diff_grad(
                  f, epsilon, trajectory_length
              )
              return auto_diff, finite_diff
  
          auto_diff, finite_diff = run(jnp.ones(4))
          self.assertAllClose(auto_diff, finite_diff, rtol=0.01)
  
imports: import functools
  import os
  from absl.testing import parameterized
  from jax import config
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import test_util
  from fun_mc import backend
  from fun_mc import fun_mc_lib
  from fun_mc import sga_hmc
  from fun_mc import test_util
label_desc: Use image node_id 4 for calling the SGAHMCTest obj's underlying member method code with example usage: obj.testHMCWithStateGrads() and returns: sum_sq,  , unknown, model,  , unknown,  , momentum_dist, res, auto_diff,  finite_diff