code: def _compute_advantages(
      logits_t: Array, q_t: Array, use_stop_gradient=True
  ) -> Tuple[Array, Array]:
      """Computes summed advantage using logits and action values."""
      policy_t = jax.nn.softmax(logits_t, axis=1)
  
      # Avoid computing gradients for action_values.
      q_t = jax.lax.select(
          use_stop_gradient, jax.lax.stop_gradient(q_t), q_t
      )
      baseline_t = jnp.sum(policy_t * q_t, axis=1)
  
      adv_t = q_t - jnp.expand_dims(baseline_t, 1)
      return policy_t, adv_t
  
file_path: rlax/rlax/_src/policy_gradients.py
imports: from typing import Optional, Tuple
  import chex
  import jax
  import jax.numpy
  from rlax._src import distributions
  from rlax._src import losses
label_desc: Use image node_id 5 for calling a global function with example usage: _compute_advantages(logits_t, q_t, use_stop_gradient) and returns: policy_t,  adv_t