code: def cross_entropy_loss(
      logits: Array, labels: Array, lengths: Array
  ) -> jax.Array:
      """Returns cross-entropy loss."""
      xe = jnp.sum(nn.log_softmax(logits) * labels, axis=-1)
      masked_xe = jnp.mean(mask_sequences(xe, lengths))
      return -masked_xe
  
imports: import functools
  from typing import Any, Dict, Tuple
  from absl import app
  from absl import flags
  from absl import logging
  from clu import metric_writers
  from flax import linen
  from flax.training import train_state
  import jax
  import jax.numpy
  import optax
  import models
  from input_pipeline import CharacterTable
  from input_pipeline import get_sequence_lengths
  from input_pipeline import mask_sequences
label_desc: Use image node_id 4 for calling a global function with example usage: cross_entropy_loss(logits, labels, lengths) and returns: unknown