code: def mixup_cutmix(
      batch: Dict["str", jnp.ndarray],
      rng: Any,
      mixup_alpha: float = 1.0,
      cutmix_alpha: float = 0.0,
      switch_prob: float = 0.5,
      label_smoothing: float = 0.0,
  ) -> Dict["str", jnp.ndarray]:
      """Performs Mixup or Cutmix within a single batch.
  
      For more details on Mixup, please see https://arxiv.org/abs/1710.09412.
      And for details on Cutmix, refer to https://arxiv.org/abs/1905.04899.
  
      Based on the implementation from:
      https://github.com/rwightman/pytorch-image-models/blob/master/timm/data/mixup.py
  
      This function supports `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'. 'inputs' is expected
          to have shape [batch, height, width, channels] or NTWHC.
        rng: JAX rng key. This key will be consumed by the function call.
        mixup_alpha: The alpha parameter of the beta distribution that the weight is
          sampled from.
        cutmix_alpha: The alpha parameter of the beta distribution that the cutmix
          weight is sampled from.
        switch_prob: The probability of switching to cutmix when both mixup and
          cutmix are enabled.
        label_smoothing: The co-efficient for label-smoothing. If using mixup or
          cutmix, this is done before mixing the labels.
  
      Returns:
        Tuple (mixed_images, mixed_labels).
      """
  
      if cutmix_alpha <= 0 and mixup_alpha <= 0:
          return batch
  
      images, labels = batch["inputs"], batch["label"]
      if labels.shape[-1] == 1:
          raise ValueError("Mixup requires one-hot targets.")
      if images.ndim not in (4, 5):
          raise ValueError(
              f"Unexpected shape: {images.shape}, wanted 4 or 5 dims."
          )
  
      rng, rng_coinflip = jax.random.split(rng)
      coin_flip = jax.random.bernoulli(rng_coinflip, p=switch_prob)
      pick_cutmix = cutmix_alpha > 0 and (mixup_alpha <= 0 or coin_flip)
  
      alpha = jax.lax.cond(
          pick_cutmix, lambda: cutmix_alpha, lambda: mixup_alpha
      )
      batch["inputs"], label_weight = jax.lax.cond(
          pick_cutmix, _do_cutmix, _do_mixup, images, rng, alpha
      )
  
      if label_smoothing > 0:
          labels = model_utils.apply_label_smoothing(
              labels, label_smoothing
          )
  
      batch["label"] = (
          label_weight * labels + (1.0 - label_weight) * labels[::-1]
      )
      return batch
  
imports: import functools
  from typing import Any, Callable, Dict, Optional, Tuple, Union
  from absl import logging
  from clu import platform
  import flax.linen
  import jax
  import jax.numpy
  import matplotlib.pyplot
  import numpy
  from scenic.model_lib.base_models import model_utils
  from scenic.model_lib.layers import nn_ops
  from scenic.train_lib import train_utils
  from scenic.train_lib_deprecated import train_utils
label_desc: Use image node_id 6 for calling a global function with example usage: mixup_cutmix(batch, rng, mixup_alpha, cutmix_alpha, switch_prob, label_smoothing) and returns: batch, batch