code: def testPreconditionedMALT(self):
      step_size = self._constant(0.2)
      num_steps = 2000
      num_leapfrog_steps = 10
      damping = 0.5
      state = jnp.ones([16, 2], dtype=self._dtype)
  
      base_mean = self._constant([1.0, 0])
      base_cov = self._constant([[1, 0.5], [0.5, 1]])
  
      bijector = tfp.bijectors.Softplus()
      base_dist = tfp.distributions.MultivariateNormalFullCovariance(
          loc=base_mean, covariance_matrix=base_cov
      )
      target_dist = bijector(base_dist)
  
      def orig_target_log_prob_fn(x):
          return target_dist.log_prob(x), ()
  
      target_log_prob_fn, state = fun_mc.transform_log_prob_fn(
          orig_target_log_prob_fn, bijector, state
      )
  
      # pylint: disable=g-long-lambda
      def kernel(malt_state, seed):
          malt_seed, seed = util.split_seed(seed, 2)
          (
              malt_state,
              _,
          ) = malt.metropolis_adjusted_langevin_trajectories_step(
              malt_state,
              step_size=step_size,
              damping=damping,
              num_integrator_steps=num_leapfrog_steps,
              target_log_prob_fn=target_log_prob_fn,
              seed=malt_seed,
          )
          return (malt_state, seed), malt_state.state_extra[0]
  
      seed = self._make_seed(_test_seed())
  
      _, chain = jax.jit(
          lambda state, seed: fun_mc.trace(  # pylint: disable=g-long-lambda
              state=(
                  malt.metropolis_adjusted_langevin_trajectories_init(
                      state, target_log_prob_fn
                  ),
                  seed,
              ),
              fn=kernel,
              num_steps=num_steps,
          )
      )(state, seed)
      # Discard the warmup samples.
      chain = chain[1000:]
  
      sample_mean = jnp.mean(chain, axis=[0, 1])
      sample_cov = tfp.stats.covariance(chain, sample_axis=[0, 1])
  
      true_samples = target_dist.sample(
          4096, seed=self._make_seed(_test_seed())
      )
  
      true_mean = jnp.mean(true_samples, axis=0)
      true_cov = tfp.stats.covariance(chain, sample_axis=[0, 1])
  
      self.assertAllClose(true_mean, sample_mean, rtol=0.1, atol=0.1)
      self.assertAllClose(true_cov, sample_cov, rtol=0.1, atol=0.1)
  
imports: import functools
  import os
  import jax
  from jax import config
  import numpy
  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 malt
  from fun_mc import test_util
label_desc: Use image node_id 4 for calling the MALTTest obj's underlying member method code with example usage: obj.testPreconditionedMALT() and returns: target_dist,  , malt_state,  seed,  malt_state