code: def mixup(
      batch: Dict["str", jnp.ndarray],
      alpha: float = 1.0,
      image_format: str = "NHWC",
      rng: Optional[Any] = None,
  ) -> Dict["str", jnp.ndarray]:
      """Mixes images and labels within a single batch.
  
      Unlike dataset_utils.mixup, this implementation samples a different weight
      for each instance in the mini-batch.
  
      Args:
        batch: dict; A batch of data with 'inputs' and 'label'.
        alpha: float; Used to control the beta distribution that weight is sampled
          from.
        image_format: string; The format of the input images.
        rng: JAX rng key. If given, JAX numpy will be used as the backend, and if
          None (default value), normal numpy will be used.
  
      Returns:
        Tuple (mixed_images, mixed_labels).
      """
      images, labels = batch["inputs"], batch["label"]
      if labels.shape[-1] == 1:
          raise ValueError("Mixup requires one-hot targets.")
      if "N" not in image_format:
          raise ValueError(
              'Mixup requires "N" to be in "image_format".'
          )
  
      batch_size = labels.shape[0]
  
      # Setup the the numpy backend and prepare mixup weights.
      if rng is None:
          np_backend = np  # Ordinary numpy
          weight = np_backend.random.beta(
              alpha, alpha, size=(batch_size, 1)
          )
      else:
          np_backend = jnp  # JAX numpy
          weight = jax.random.beta(
              rng, alpha, alpha, shape=(batch_size, 1)
          )
  
      # Make sure that the original image has the higher weight during mixup
      weight = np_backend.maximum(weight, 1.0 - weight)
  
      # Mixup labels.
      batch["label"] = weight * labels + (1.0 - weight) * labels[::-1]
  
      # Mixup inputs.
      # Shape calculations use np to avoid device memory fragmentation:
      image_weight_shape = np.ones((images.ndim))
      image_weight_shape[image_format.index("N")] = batch_size
      weight = np_backend.reshape(
          weight, image_weight_shape.astype(np_backend.int32)
      )
      reverse = tuple(
          slice(images.shape[i]) if d != "N" else slice(-1, None, -1)
          for i, d in enumerate(image_format)
      )
      batch["inputs"] = (
          weight * images + (1.0 - weight) * images[reverse]
      )
  
      return batch
  
file_path: scenic/scenic/projects/ncr/utils.py
imports: import os
  from typing import Any, Dict, Optional, Tuple
  from flax import jax_utils
  from flax.training import checkpoints
  import jax
  import jax.numpy
  import numpy
  from scenic.train_lib import train_utils
  from tensorflow.io import gfile
label_desc: Use image node_id 5 for calling a global function with example usage: mixup(batch, alpha, image_format, rng) and returns: batch