def sequence_cross_entropy_loss(*, labels: jnp.ndarray,
                                logits: jnp.ndarray) -> jnp.ndarray:
  """Computes the mean cross-entropy for the sequence predictions.

  Args:
    labels: 2D int array of shape (B, T) where each value is in [0, C-1].
    logits: 3D array of shape (B, T, C) where C is number of classes.

  Returns:
    float loss.
  """
  _, _, num_classes = logits.shape
  labels_one_hot = jax.nn.one_hot(labels, num_classes)
  logprobs = jax.nn.log_softmax(logits, axis=-1)
  return -jnp.mean(jnp.sum(logprobs * labels_one_hot, axis=-1))
Computes the mean cross-entropy for the sequence predictions.

  Args:
    labels: 2D int array of shape (B, T) where each value is in [0, C-1].
    logits: 3D array of shape (B, T, C) where C is number of classes.

  Returns:
    float loss.
  
videogvt/train_lib/losses.py