code: def accuracy(logits, targets, weights=None):
      """Compute weighted accuracy for log probs and targets.
      Args:
       logits: [batch, length, num_classes] float array.
       targets: categorical targets [batch, length] int array.
       weights: None or array of shape [batch, length]
      Returns:
        Tuple of scalar loss and batch normalizing factor.
      """
      if logits.ndim != targets.ndim + 1:
          raise ValueError(
              "Incorrect shapes. Got shape %s logits and %s targets"
              % (str(logits.shape), str(targets.shape))
          )
  
      loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)
      loss *= weights
  
      return loss.sum(), weights.sum()
  
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 3 for calling a global function with example usage: accuracy(logits, targets, weights) and returns: loss,  weights