code: def augment_example(
      example, dtype=tf.float32, data_augmentations=None
  ):
      """Apply data augmentation on the given training example.
  
      Args:
        example: dict; Example that has an 'image' and a 'label'.
        dtype: Tensorflow data type; Data type of the image.
        data_augmentations: list(str); Types of data augmentation applied on
          training data.
  
      Returns:
        An augmented training example.
      """
      image = tf.cast(example["inputs"], dtype=dtype)
      if data_augmentations is not None:
          if "cifar_default" in data_augmentations:
              image = dataset_utils.augment_random_crop_flip(
                  image,
                  HEIGHT,
                  WIDTH,
                  NUM_CHANNELS,
                  crop_padding=4,
                  flip=True,
              )
          image = tf.cast(image, dtype=dtype)
      return {"inputs": image, "label": example["label"]}
  
file_path: scenic/scenic/dataset_lib/cifar10_dataset.py
imports: import functools
  from typing import Optional
  from absl import logging
  import jax.numpy
  from scenic.dataset_lib import dataset_utils
  from scenic.dataset_lib import datasets
  import tensorflow
label_desc: Use image node_id 2 for calling a global function with example usage: augment_example(example, dtype, data_augmentations) and returns: dict