docstr: Evaluates the criterion in sharded and non-sharded contexts.
docstr_tok: ["Evaluates","the","criterion","in","sharded","and","non-sharded","contexts","."]
code: def testCriterion(
      self, criterion_fn, with_trajectory, with_state_mean
  ):
      """Evaluates the criterion in sharded and non-sharded contexts."""
      seed = self._make_seed(_test_seed())
      seeds = util.split_seed(seed, 4)
      previous_state = {
          "global": self._constant(
              util.random_normal([2, 3], self._dtype, seeds[0])
          ),
          "local": self._constant(
              util.random_normal([2, 2, 3], self._dtype, seeds[1])
          ),
      }
      named_axis = util.map_tree(lambda _: [], previous_state)
      chain_named_axis = []
      trajectory_length = self._constant(0.5)
      accept_prob = self._constant([0.1, 0.5])
      state_mean = {
          "global": self._constant(
              util.random_normal([3], self._dtype, seeds[2])
          ),
          "local": self._constant(
              util.random_normal([2, 3], self._dtype, seeds[3])
          ),
      }
  
      def eval_criterion(
          trajectory_length,
          previous_state,
          accept_prob,
          chain_named_axis,
          state_mean,
          named_axis,
      ):
          extra_kwargs = {}
          if with_trajectory:
              extra_kwargs.update(trajectory_length=trajectory_length)
          if with_state_mean:
              extra_kwargs.update(
                  state_mean=state_mean, state_mean_weight=0.5
              )
  
          def proposed_state_part(previous_state, named_axis):
              if BACKEND == "backend_jax":
                  part_trajectory_length = distribute_lib.pbroadcast(
                      trajectory_length, [chain_named_axis, named_axis]
                  )
              else:
                  part_trajectory_length = trajectory_length
              return (part_trajectory_length + 1) * previous_state
  
          proposed_state = util.map_tree_up_to(
              previous_state,
              proposed_state_part,
              previous_state,
              named_axis,
          )
  
          return criterion_fn(
              previous_state=previous_state,
              proposed_state=proposed_state,
              accept_prob=accept_prob,
              named_axis=named_axis,
              chain_named_axis=chain_named_axis,
              **extra_kwargs,
          )
  
      value, _, grad = fun_mc.call_potential_fn_with_grads(
          functools.partial(
              eval_criterion,
              previous_state=previous_state,
              chain_named_axis=chain_named_axis,
              accept_prob=accept_prob,
              state_mean=state_mean,
              named_axis=named_axis,
          ),
          trajectory_length,
      )
  
      self.assertEqual(self._dtype, value.dtype)
      self.assertEqual(self._dtype, grad.dtype)
      self.assertAllGreater(jnp.abs(grad), 0.0)
  
      if BACKEND == "backend_jax":
          named_axis = {
              "global": [],
              "local": "local",
          }
          chain_named_axis = "chain"
          in_axes = {
              "global": None,
              "local": 0,
          }
  
          @functools.partial(jax.pmap, axis_name="chain")
          def run_chain(previous_state, accept_prob):
              @functools.partial(
                  jax.pmap,
                  axis_name="local",
                  in_axes=(in_axes, in_axes),
              )
              def run_state(previous_state, state_mean):
                  value, _, grad = fun_mc.call_potential_fn_with_grads(
                      functools.partial(
                          eval_criterion,
                          previous_state=previous_state,
                          chain_named_axis=chain_named_axis,
                          accept_prob=accept_prob,
                          state_mean=state_mean,
                          named_axis=named_axis,
                      ),
                      trajectory_length,
                  )
                  return value, grad
  
              return run_state(previous_state, state_mean)
  
          sharded_value, sharded_grad = run_chain(
              previous_state, accept_prob
          )
          self.assertAllClose(value, sharded_value[0, 0])
          self.assertAllClose(grad, sharded_grad[0, 0])
  
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 5 for calling the SGAHMCTest obj's underlying member method code with example usage: obj.testCriterion(criterion_fn, with_trajectory, with_state_mean) and returns: criterion_fn, unknown, run_state, value,  grad