code: def kl_alpha_loss(
      restarting_weights: Array,
      kl_constraints: Sequence[Tuple[Array, LagrangePenalty]] = (),
      axis_name: Optional[str] = None,
  ):
      """Calculates the losses for multiple KL constraints.
  
      Args:
        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.
        kl_constraints: KL and variables for applying Lagrangian penalties to bound
          them in the M-step, KLs are [E*, A?]. Here A is the action dimension
          in the case of per-dimension KL constraints.
        axis_name: Optional axis name for `pmap`. If `None`, computations are
          performed locally on each device.
  
      Returns:
        The kl loss and dual variable loss both shape E*.
      """
      chex.assert_type(restarting_weights, float)
      if kl_constraints:
          for kl, penalty in kl_constraints:
              chex.assert_rank(penalty.epsilon, 0)
              chex.assert_type(
                  [kl, penalty.alpha, penalty.epsilon], float
              )
              chex.assert_equal_shape_prefix(
                  [kl, restarting_weights], restarting_weights.ndim
              )
  
          # Implement decoupled KL constraints.
          kl_alpha_losses = [
              kl_constraint_loss(kl, penalty, lambda x: x)[:2]
              for kl, penalty in kl_constraints
          ]
          kl_loss, alpha_loss = [
              sum(losses) for losses in zip(*kl_alpha_losses)
          ]
          all_sum = base.AllSum(axis_name)
          num_samples = all_sum(restarting_weights) + _EPSILON
          # Reshape in case KL is per dimension.
          kl_loss = all_sum(kl_loss * restarting_weights) / num_samples
          alpha_loss = (
              all_sum(alpha_loss * restarting_weights) / num_samples
          )
      else:
          # No M-step constraint.
          kl_loss = jnp.asarray(0.0)
          alpha_loss = jnp.asarray(0.0)
      return kl_loss, alpha_loss
  
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 8 for calling a global function with example usage: kl_alpha_loss(restarting_weights, kl_constraints, axis_name) and returns: kl_loss,  alpha_loss