code: def compute_weighted_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)
      normalizing_factor = np.prod(logits.shape[:-1])
      if weights is not None:
          loss = loss * weights
          normalizing_factor = weights.sum()
  
      return loss.sum(), normalizing_factor
  
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 4 for calling a global function with example usage: compute_weighted_accuracy(logits, targets, weights) and returns: loss,  normalizing_factor