class_docstr: Tests for the MPO losses.
class_docstr_tok: ["Tests","for","the","MPO","losses","."]
docstr: Tests the gradients from the E-step against the analytic ones.
docstr_tok: ["Tests","the","gradients","from","the","E-step","against","the","analytic","ones","."]
code: def test_e_step_gradient_computation(
      self,
      e_step_fn,
      additional_inputs,
      expected_deriv_of_target,
      sample_dimension,
  ):
      """Tests the gradients from the E-step against the analytic ones."""
      # Target has shape [NUM_SAMPLES, T, B] => [1, 1, 1]
      target = jnp.array([[3]], jnp.float32)
      if sample_dimension:
          target = jnp.expand_dims(target, axis=0)
      temperature = jnp.array(0.1, jnp.float32)
  
      def fn(target_, temperature_):
          temperature_constraint = mpo_ops.LagrangePenalty(
              temperature_, _EPSILON_BOUND
          )
          temperature_loss, _, _ = e_step_fn(
              target_,
              temperature_constraint=temperature_constraint,
              projection_operator=_PROJECTION_OPERATOR,
              **additional_inputs,
          )
          return jnp.mean(temperature_loss)
  
      grad = jax.grad(fn, argnums=(0, 1))(target, temperature)
  
      np.testing.assert_almost_equal(
          np.array(grad[0]),
          np.array(expected_deriv_of_target, np.float32),
          decimal=4,
      )
      self.assertAlmostEqual(grad[1], _EPSILON_BOUND, places=4)
  
imports: import functools
  import math
  from typing import Tuple
  from absl.testing import absltest
  from absl.testing import parameterized
  import chex
  import haiku
  import jax
  import jax.numpy
  import numpy
  import optax
  from rlax._src import distributions
  from rlax._src import mpo_ops
label_desc: Use image node_id 2 for calling the MPOTest obj's underlying member method code with example usage: obj.test_e_step_gradient_computation(e_step_fn, additional_inputs, expected_deriv_of_target, sample_dimension) and returns: jnp