class_docstr: Tests for the MPO losses.
class_docstr_tok: ["Tests","for","the","MPO","losses","."]
docstr: Tests no gradients flow through `weights` in the E-Step.
docstr_tok: ["Tests","no","gradients","flow","through","`","weights","`","in","the","E-Step","."]
code: def test_e_step_stop_gradient(
      self, e_step_fn, additional_inputs, sample_dimension
  ):
      """Tests no gradients flow through `weights` in the E-Step."""
      # 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)
  
      # pylint: disable=g-long-lambda
      def mean_weights_fn(target_, temperature_):
          temperature_constraint = mpo_ops.LagrangePenalty(
              temperature_, _EPSILON_BOUND
          )
          _, weights, _ = e_step_fn(
              target_,
              temperature_constraint=temperature_constraint,
              projection_operator=_PROJECTION_OPERATOR,
              **additional_inputs,
          )
          return jnp.mean(weights)
  
      grad = jax.grad(mean_weights_fn, argnums=(0, 1))(
          target, temperature
      )
      np.testing.assert_almost_equal(
          np.array(grad[0]), np.zeros_like(grad[0]), decimal=4
      )
      self.assertAlmostEqual(grad[1], 0.0, 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 3 for calling the MPOTest obj's underlying member method code with example usage: obj.test_e_step_stop_gradient(e_step_fn, additional_inputs, sample_dimension) and returns: jnp