code: def tf_apply_to_image_or_images(fn, image_or_images, **map_kw):
      """Applies a function to a single image or each image in a batch of them.
  
      Args:
        fn: the function to apply, receives an image, returns an image.
        image_or_images: Either a single image, or a batch of images.
        **map_kw: Arguments passed through to tf.map_fn if called.
  
      Returns:
        The result of applying the function to the image or batch of images.
  
      Raises:
        ValueError: if the input is not of rank 3 or 4.
      """
      static_rank = len(image_or_images.get_shape().as_list())
      if static_rank == 3:  # A single image: HWC
          return fn(image_or_images)
      elif static_rank == 4:  # A batch of images: BHWC
          return tf.map_fn(fn, image_or_images, **map_kw)
      elif static_rank > 4:  # A batch of images: ...HWC
          input_shape = tf.shape(image_or_images)
          h, w, c = image_or_images.get_shape().as_list()[-3:]
          image_or_images = tf.reshape(image_or_images, [-1, h, w, c])
          image_or_images = tf.map_fn(fn, image_or_images, **map_kw)
          return tf.reshape(image_or_images, input_shape)
      else:
          raise ValueError("Unsupported image rank: %d" % static_rank)
  
file_path: scenic/scenic/dataset_lib/big_transfer/preprocessing/utils.py
imports: from collections import abc
  import functools
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: tf_apply_to_image_or_images(fn, image_or_images) and returns: fn, tf, tf