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 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 4 for calling a global function with example usage: cross_entropy(logits, targets, weights, label_smoothing) and returns: loss,  normalizing_factor