code: def compute_weighted_cross_entropy(
      logits, targets, weights=None, label_smoothing=0.0
  ):
      """Compute weighted cross entropy and entropy 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].
       label_smoothing: label smoothing constant, used to determine the on and off
         values.
  
      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))
          )
      vocab_size = logits.shape[-1]
      confidence = 1.0 - label_smoothing
      low_confidence = (1.0 - confidence) / (vocab_size - 1)
      normalizing_constant = -(
          confidence * jnp.log(confidence)
          + (vocab_size - 1)
          * low_confidence
          * jnp.log(low_confidence + 1e-20)
      )
      soft_targets = common_utils.onehot(
          targets,
          vocab_size,
          on_value=confidence,
          off_value=low_confidence,
      )
  
      loss = -jnp.sum(soft_targets * nn.log_softmax(logits), axis=-1)
      loss = loss - normalizing_constant
  
      normalizing_factor = np.prod(targets.shape)
      if weights is not None:
          loss = loss * weights
          normalizing_factor = weights.sum()
  
      return loss.sum(), normalizing_factor
  
imports: import collections
  import os
  from absl import logging
  from clu import metric_writers
  from clu import periodic_actions
  from flax import linen
  from flax.training import checkpoints
  from flax.training import common_utils
  import jax
  from jax import random
  import jax.numpy
  from jax.sharding import PartitionSpec, Mesh, NamedSharding
  import ml_collections
  import numpy
  import optax
  import tensorflow
  import input_pipeline
  import models
  import temperature_sampler
  import utils
label_desc: Use image node_id 3 for calling a global function with example usage: compute_weighted_cross_entropy(logits, targets, weights, label_smoothing) and returns: loss,  normalizing_factor