class_docstr: Tests for the MPO losses.
class_docstr_tok: ["Tests","for","the","MPO","losses","."]
docstr: Test that only the top k fraction are used.
docstr_tok: ["Test","that","only","the","top","k","fraction","are","used","."]
code: def test_top_k_fraction(
      self, top_k_fraction, scaled_advantages, expected_top_k_weights
  ):
      """Test that only the top k fraction are used."""
      top_k_weights = mpo_ops.get_top_k_weights(
          top_k_fraction,
          jnp.ones_like(scaled_advantages),
          scaled_advantages,
      )
      np.testing.assert_allclose(top_k_weights, expected_top_k_weights)
  
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 7 for calling the MPOTest obj's underlying member method code with example usage: obj.test_top_k_fraction(top_k_fraction, scaled_advantages, expected_top_k_weights) without return types