code: def get_top_k_weights(
      top_k_fraction: float,
      restarting_weights: Array,
      scaled_advantages: Array,
      axis_name: Optional[str] = None,
      use_stop_gradient: bool = True,
  ) -> Array:
      """Get the weights for the top top_k_fraction of advantages.
  
      Args:
        top_k_fraction: The fraction of weights to use.
        restarting_weights: Restarting weights, shape E*, 0 means that this step is
          the start of a new episode and we ignore losses at this step because the
          agent cannot influence these.
        scaled_advantages: The advantages for each example (shape E*), scaled by
          temperature.
        axis_name: Optional axis name for `pmap`. If `None`, computations are
          performed locally on each device.
        use_stop_gradient: bool indicating whether or not to apply stop gradient.
  
      Returns:
        Weights for the top top_k_fraction of advantages
      """
      chex.assert_equal_shape([scaled_advantages, restarting_weights])
      chex.assert_type([scaled_advantages, restarting_weights], float)
  
      if not 0.0 < top_k_fraction <= 1.0:
          raise ValueError(
              f"`top_k_fraction` must be in (0, 1], got {top_k_fraction}"
          )
      logging.info("[vmpo_e_step] top_k_fraction: %f", top_k_fraction)
  
      if top_k_fraction < 1.0:
          # Don't include the restarting samples in the determination of top-k.
          valid_scaled_advantages = (
              scaled_advantages - (1.0 - restarting_weights) * _INFINITY
          )
          # Determine the minimum top-k value across all devices,
          if axis_name:
              all_valid_scaled_advantages = jax.lax.all_gather(
                  valid_scaled_advantages, axis_name=axis_name
              )
          else:
              all_valid_scaled_advantages = valid_scaled_advantages
          top_k = int(
              top_k_fraction * jnp.size(all_valid_scaled_advantages)
          )
          if top_k == 0:
              raise ValueError(
                  "top_k_fraction too low to get any valid scaled advantages."
              )
          # TODO(b/160450251): Use jnp.partition(all_valid_scaled_advantages, top_k)
          #   when this is implemented in jax.
          top_k_min = jnp.sort(
              jnp.reshape(all_valid_scaled_advantages, [-1])
          )[-top_k]
          # Fold the top-k into the restarting weights.
          top_k_weights = jnp.greater_equal(
              valid_scaled_advantages, top_k_min
          ).astype(jnp.float32)
          top_k_weights = jax.lax.select(
              use_stop_gradient,
              jax.lax.stop_gradient(top_k_weights),
              top_k_weights,
          )
          top_k_restarting_weights = restarting_weights * top_k_weights
      else:
          top_k_restarting_weights = restarting_weights
  
      return top_k_restarting_weights
  
file_path: rlax/rlax/_src/mpo_ops.py
imports: import functools
  from typing import Callable, NamedTuple, Optional, Sequence, Tuple
  from absl import logging
  import chex
  import jax
  import jax.numpy
  from rlax._src import base
label_desc: Use image node_id 5 for calling a global function with example usage: get_top_k_weights(top_k_fraction, restarting_weights, scaled_advantages, axis_name, use_stop_gradient) and returns: top_k_restarting_weights