def LossInput(dist_inputs, actions, q_values, act_log_probs, mask):  # pylint: disable=invalid-name
      """Calculates action log probabilities and normalizes advantages."""
      # (batch_size, n_samples, ...) -> (n_samples, batch_size, ...)
      q_values = jnp.swapaxes(q_values, 0, 1)
      mask = jnp.swapaxes(mask, 0, 1)
      actions = jnp.swapaxes(actions, 0, 1)
      act_log_probs = jnp.swapaxes(act_log_probs, 0, 1)

      # TODO(pkozakowski,lukaszkaiser): Try max here, or reweighting?
      if self._sample_all_discrete_actions:
        values = jnp.sum(q_values * jnp.exp(act_log_probs), axis=0)
      else:
        values = jnp.mean(q_values, axis=0)
      advantages = q_values - values  # Broadcasting values over n_samples
      advantages = self._preprocess_advantages(advantages)

      # Broadcast inputs and calculate log-probs
      dist_inputs = jnp.broadcast_to(
          dist_inputs, (self._q_value_n_samples,) + dist_inputs.shape)
      log_probs = self._policy_dist.log_prob(dist_inputs, actions)
      return (log_probs, advantages, act_log_probs, mask)
Calculates action log probabilities and normalizes advantages.
trax/rl/actor_critic.py