code: def softmax_cross_entropy(
      logits: chex.Array, target_probs: chex.Array
  ) -> chex.Array:
      """Compute cross entropy loss between logits and target probabilities."""
      chex.assert_equal_shape([target_probs, logits])
      return -jnp.sum(
          target_probs * jax.nn.log_softmax(logits), axis=-1
      )
  
file_path: acme/acme/agents/jax/mpo/rollout_loss.py
imports: from typing import Tuple
  from acme import types
  from acme.adders import reverb
  from acme.agents.jax.mpo import categorical_mpo
  from acme.agents.jax.mpo import networks
  from acme.agents.jax.mpo import types
  from acme.agents.jax.mpo import utils
  from acme.jax import networks
  import chex
  import jax
  import jax.numpy
  import rlax
label_desc: Use image node_id 1 for calling a global function with example usage: softmax_cross_entropy(logits, target_probs) and returns: unknown