code: def softmax_cross_entropy_loss_with_logits(
      labels: jnp.ndarray, logits: jnp.ndarray
  ) -> jnp.ndarray:
      """Implementation of the softmax cross entropy loss."""
      return -jnp.sum(labels * nn.log_softmax(logits))
  
imports: from flax import linen
  import jax.numpy
label_desc: Use image node_id 3 for calling a global function with example usage: softmax_cross_entropy_loss_with_logits(labels, logits) and returns: unknown