code: def cross_entropy_logits_sparse(
      logits: JaxArray, labels: Union[JaxArray, int]
  ) -> JaxArray:
      """Computes the softmax cross-entropy loss.
  
      Args:
          logits: (batch, ..., #class) tensor of logits.
          labels: (batch, ...) integer tensor of label indexes in {0, ...,#nclass-1} or just a single integer.
  
      Returns:
          (batch, ...) tensor of the cross-entropies for each entry.
      """
  
      if isinstance(labels, int):
          labeled_logits = logits[..., labels]
      else:
          labeled_logits = jn.take_along_axis(
              logits, labels[..., None], -1
          ).squeeze(-1)
  
      return logsumexp(logits, axis=-1) - labeled_logits
  
file_path: objax/objax/functional/loss.py
imports: from typing import Union, Iterable, Optional
  import jax.numpy
  from objax.functional.core import logsumexp
  from objax.typing import JaxArray
  import warnings
label_desc: Use image node_id 2 for calling a global function with example usage: cross_entropy_logits_sparse(logits, labels) and returns: unknown