code: def mixup(
      batch: Dict["str", jnp.ndarray],
      alpha: float = 1.0,
      image_format: str = "NHWC",
      input_key: str = "inputs",
      label_key: str = "label",
      rng: Optional[Any] = None,
  ) -> Dict["str", jnp.ndarray]:
      """Mixes images and labels within a single batch.
  
      For more details, please see https://arxiv.org/abs/1710.09412.
  
      This function supports both using `numpy` to do mixup in the input-pipeline
      and `jax.numpy` to do mixup within a jitted/pmapped function (e.g. within
      a pmapped train step to apply mixup on device patch).
  
      Results in a batch with:
        mixed_images[idx] = weight * images[idx] + (1-weight) * images[-(idx+1)],
        where weight is sampled from a beta distribution with parameter alpha.
  
      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.
        input_key: The key in the `batch` dictionary corresponding to the input
          images. Default is `inputs`.
        label_key: The key in the `batch` dictionary corresponding to the labels.
          Default is `labels`.
        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[input_key], batch[label_key]
      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]
  
      # Set up the numpy backend and prepare mixup weights.
      if rng is None:
          np_backend = np  # Ordinary numpy
          weight = np_backend.random.beta(alpha, alpha)
      else:
          np_backend = jnp  # JAX numpy
          weight = jax.random.beta(rng, alpha, alpha)
      label_weight_shape = np.ones(labels.ndim)
      label_weight_shape[image_format.index("N")] = batch_size
      weight *= np_backend.ones(
          label_weight_shape.astype(np_backend.int32)
      )
  
      # Mixup labels.
      batch[label_key] = 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[input_key] = (
          weight * images + (1.0 - weight) * images[reverse]
      )
  
      return batch
  
imports: import collections
  import dataclasses
  import functools
  import itertools
  from typing import Any, Callable, Dict, Iterator, Optional, Sequence, Union
  from absl import logging
  from flax.training import common_utils
  import jax
  import jax.numpy
  import numpy
  import tensorflow
  import tensorflow_datasets
label_desc: Use image node_id 12 for calling a global function with example usage: mixup(batch, alpha, image_format, input_key, label_key, rng) and returns: batch