code: def compute_metrics(
      logits: Array, labels: Array, eos_id: int
  ) -> Dict[str, jax.Array]:
      """Computes metrics and returns them."""
      lengths = get_sequence_lengths(labels, eos_id)
      loss = cross_entropy_loss(logits, labels, lengths)
      # Computes sequence accuracy, which is the same as the accuracy during
      # inference, since teacher forcing is irrelevant when all output are correct.
      token_accuracy = jnp.argmax(logits, -1) == jnp.argmax(labels, -1)
      sequence_accuracy = (
          jnp.sum(mask_sequences(token_accuracy, lengths), axis=-1)
          == lengths
      )
      accuracy = jnp.mean(sequence_accuracy)
      metrics = {
          "loss": loss,
          "accuracy": accuracy,
      }
      return metrics
  
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 5 for calling a global function with example usage: compute_metrics(logits, labels, eos_id) and returns: metrics