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
  
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 5 for calling a global function with example usage: _do_cutmix(inputs, rng, alpha) and returns: result_img,  weight