def sequence_soft_cross_entropy_loss(*, target_labels: jnp.ndarray,
                                     target_probs: jnp.ndarray,
                                     logits: jnp.ndarray) -> jnp.ndarray:
  """Computes the mean cross-entropy for the sequence soft-label predictions.

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

  Returns:
    float loss.
  """
  logprobs = jax.nn.log_softmax(logits, axis=-1)
  cross_entropy = jnp.sum(
      target_probs * jnp.take_along_axis(logprobs, target_labels, axis=-1),
      axis=-1)
  return -jnp.mean(cross_entropy)
Computes the mean cross-entropy for the sequence soft-label predictions.

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

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