code: def compute_metrics(logits, labels):
      loss = cross_entropy_loss(logits, labels)
      accuracy = jnp.mean(jnp.argmax(logits, -1) == labels)
      metrics = {
          "loss": loss,
          "accuracy": accuracy,
      }
      metrics = lax.pmean(metrics, axis_name="batch")
      return metrics
  
imports: import functools
  import time
  from typing import Any
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax.training import checkpoints
  from flax.training import common_utils
  from flax.training import dynamic_scale
  from flax.training import train_state
  import jax
  from jax import lax
  import jax.numpy
  from jax import random
  import ml_collections
  import optax
  import tensorflow
  import tensorflow_datasets
  import input_pipeline
  import models
label_desc: Use image node_id 4 for calling a global function with example usage: compute_metrics(logits, labels) and returns: metrics