code: def compute_metrics(logits, labels, weights, label_smoothing=0.0):
      """Compute summary metrics."""
      loss, weight_sum = compute_weighted_cross_entropy(
          logits, labels, weights, label_smoothing
      )
      acc, _ = compute_weighted_accuracy(logits, labels, weights)
      metrics = {
          "loss": loss,
          "accuracy": acc,
          "denominator": weight_sum,
      }
      metrics = jax.lax.psum(metrics, axis_name="batch")
      return metrics
  
imports: import collections
  import functools
  import os
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import jax_utils
  from flax import linen
  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
  import jax.numpy
  import ml_collections
  import numpy
  import optax
  import tensorflow
  import bleu
  import decode
  import input_pipeline
  import models
label_desc: Use image node_id 5 for calling a global function with example usage: compute_metrics(logits, labels, weights, label_smoothing) and returns: metrics