code: def _integrator_conserves_energy(
      self, x, independent_chain_ndims, seed
  ):
      event_dims = ps.range(independent_chain_ndims, tf.rank(x))
  
      target_fn = lambda x: self._log_gamma_log_prob(x, event_dims)
  
      m = tf.random.normal(tf.shape(x), seed=seed)
      log_prob_0 = target_fn(x)
      old_energy = -log_prob_0 + 0.5 * tf.reduce_sum(
          m**2.0, axis=event_dims
      )
  
      event_size = np.prod(
          self.evaluate(x).shape[independent_chain_ndims:]
      )
  
      integrator = leapfrog_impl.SimpleLeapfrogIntegrator(
          target_fn, step_sizes=[0.09 / event_size], num_steps=1000
      )
  
      [[new_m], [_], log_prob_1, [_]] = integrator([m], [x])
  
      new_energy = -log_prob_1 + 0.5 * tf.reduce_sum(
          new_m**2.0, axis=event_dims
      )
  
      old_energy_, new_energy_ = self.evaluate([old_energy, new_energy])
      tf1.logging.vlog(
          1,
          "average energy relative change: {}".format(
              (1.0 - new_energy_ / old_energy_).mean()
          ),
      )
      self.assertAllClose(old_energy_, new_energy_, atol=0.0, rtol=0.02)
  
file_path: probability/tensorflow_probability/python/mcmc/internal/leapfrog_integrator_test.py
imports: import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import prefer_static
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.mcmc.internal import leapfrog_integrator
label_desc: Use image node_id 4 for calling the LeapfrogIntegratorTest obj's underlying member method code with example usage: obj._integrator_conserves_energy(x, independent_chain_ndims, seed) without return types