code: def compute_metrics(logits, labels, weights, label_smoothing=0.0):
      """Compute summary metrics."""
      loss, normalizer = cross_entropy(
          logits, labels, weights, label_smoothing
      )
      acc, _ = accuracy(logits, labels, weights)
      metrics = {
          "loss": loss,
          "accuracy": acc,
          "normalizer": normalizer,
      }
      metrics = jax.lax.psum(metrics, axis_name="batch")
      return metrics
  
imports: import logging
  import os
  import sys
  from dataclasses import dataclass, field
  from pathlib import Path
  from typing import Dict, List, Optional, Tuple
  import numpy
  from datasets import load_dataset
  from tqdm import tqdm
  import jax
  import jax.numpy
  from flax import jax_utils
  from flax.optim import Adam
  from flax.training import common_utils
  from flax.training.common_utils import get_metrics
  from jax.nn import log_softmax
  from modeling_flax_performer import FlaxPerformerForMaskedLM
  from transformers import MODEL_FOR_MASKED_LM_MAPPING, AutoTokenizer, BertConfig, FlaxBertForMaskedLM, HfArgumentParser, PreTrainedTokenizerBase, TensorType, TrainingArguments, is_tensorboard_available, set_seed
label_desc: Use image node_id 2 for calling a global function with example usage: compute_metrics(logits, labels, weights, label_smoothing) and returns: metrics