code: def _do_cutmix(
      inputs: jnp.ndarray, rng: Any, alpha: float
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Performs Cutmix.
  
      Args:
        inputs: The input images of shape NHWC or NTHWC.
        rng: A PRNGKey. Will be consumed by this function.
        alpha: The alpha value for cutmix.
  
      Returns:
        The modified images and label weights.
      """
      rng, beta_key = jax.random.split(rng)
      cutmix_lambda = jax.random.beta(beta_key, alpha, alpha)
      ratio = jnp.sqrt(1 - cutmix_lambda)
  
      # TODO(unterthiner): we are using the same bounding box for the whole batch
      y_min, y_max, x_min, x_max = get_random_bounding_box(
          inputs.shape[-3:-1], ratio, rng
      )
  
      height, width = inputs.shape[-3], inputs.shape[-2]
      y_idx = jnp.arange(height)
      x_idx = jnp.arange(width)
      mask0 = (y_min <= y_idx) & (y_idx < y_max)
      mask1 = (x_min <= x_idx) & (x_idx < x_max)
      mask = (~jnp.outer(mask0, mask1)).astype(int)
      if inputs.ndim == 4:  # image format NWHC
          mask = jnp.expand_dims(mask, axis=(0, -1))
      elif inputs.ndim == 5:  # image format NTWHC
          mask = jnp.expand_dims(mask, axis=(0, 1, -1))
      else:
          raise ValueError("Invalid image format")
  
      result_img = inputs * mask + jnp.flip(inputs, axis=0) * (
          1.0 - mask
      )
      box_area = (y_max - y_min) * (x_max - x_min)
      weight = 1.0 - box_area / float(height * width)
      weight *= jnp.ones((inputs.shape[0], 1))
      return result_img, weight
  
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 3 for calling a global function with example usage: _do_cutmix(inputs, rng, alpha) and returns: result_img,  weight