code: def categorical_cross_entropy(labels: Array, logits: Array) -> Array:
      """Computes the softmax cross entropy between sets of logits and labels.
  
      See "Deep Learning" by Goodfellow et al.
      (http://www.deeplearningbook.org/contents/prob.html). The computation is
      equivalent to:
  
                      sum_i (labels_i * log_softmax(logits_i))
  
      Args:
        labels: a valid probability distribution (non-negative, sum to 1).
        logits: unnormalized log probabilities.
  
      Returns:
        a scalar loss.
      """
      warnings.warn(
          "Rlax categorical_cross_entropy will be deprecated. "
          "Please use distrax.Categorical.cross_entropy instead.",
          PendingDeprecationWarning,
          stacklevel=2,
      )
      return distrax.Categorical(probs=labels).cross_entropy(
          distrax.Categorical(logits=logits)
      )
  
file_path: rlax/rlax/_src/distributions.py
imports: import collections
  import warnings
  import chex
  import distrax
  import jax.numpy
label_desc: Use image node_id 9 for calling a global function with example usage: categorical_cross_entropy(labels, logits) and returns: distrax