code: def cross_entropy(logits, targets, weights=None, label_smoothing=0.0):
      """Compute 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 * log_softmax(logits), axis=-1)
      loss = loss - normalizing_constant
  
      if weights is not None:
          loss = loss * weights
          normalizing_factor = weights.sum()
      else:
          normalizing_factor = np.prod(targets.shape)
  
      return loss.sum(), normalizing_factor
  
imports: import logging
  import os
  import sys
  from dataclasses import dataclass, field
  from pathlib import Path
  from typing import Dict, List, Optional, Tuple
  import jax
  import jax.numpy
  import numpy
  from datasets import load_dataset
  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 tqdm import tqdm
  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 4 for calling a global function with example usage: cross_entropy(logits, targets, weights, label_smoothing) and returns: loss,  normalizing_factor