code: def _do_mixup(
      inputs: jnp.ndarray, rng: Any, alpha: float
  ) -> Tuple[jnp.ndarray, jnp.ndarray]:
      """Performs Mixup.
  
      Args:
        inputs: The input images of shape NHWC or NTHWC. Mixup is always performed
          along the leading axis of the array, i.e., along the "N" dimension.
        rng: A PRNGKey. Will be consumed by this function.
        alpha: The alpha value for mixup.
  
      Returns:
        The modified images and label weights.
      """
      batch_size = inputs.shape[0]
      weight = jax.random.beta(rng, alpha, alpha)
      weight *= jnp.ones((batch_size, 1))
  
      # Mixup inputs.
      # Shape calculations use np to avoid device memory fragmentation:
      image_weight_shape = np.ones(inputs.ndim, np.int32)
      image_weight_shape[0] = batch_size
      image_weight = weight.reshape(image_weight_shape)
      reverse = tuple(
          slice(inputs.shape[i]) if i > 0 else slice(-1, None, -1)
          for i in range(inputs.ndim)
      )
      result_img = (
          image_weight * inputs + (1.0 - image_weight) * inputs[reverse]
      )
      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 2 for calling a global function with example usage: _do_mixup(inputs, rng, alpha) and returns: result_img,  weight