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
  
file_path: scenic/scenic/projects/objectvivit/train_utils.py
imports: from typing import Any, Dict, Tuple
  from absl import logging
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.base_models import model_utils
label_desc: Use image node_id 4 for calling a global function with example usage: mixup_cutmix(batch, rng, mixup_alpha, cutmix_alpha, switch_prob, label_smoothing) and returns: batch, batch