code: def _image_tiles_to_mosaic(
      image: tf.Tensor, mosaic_size: int
  ) -> tf.Tensor:
      """Reshapes a batch of image tiles into a mosaic."""
      assert (
          len(image.shape) == 4
      ), f"Expect shape [num_tiles, h, w, c], got {image.shape}."
  
      # Get dynamic image shape:
      shape = tf.shape(image)
      num_tiles, h, w, c = shape[0], shape[1], shape[2], shape[3]
  
      tf.debugging.assert_equal(
          num_tiles,
          mosaic_size**2,
          "The first dimension must contain exactly self.mosaic_size ** 2 "
          f"elements, but got images of shape {image.shape}",
      )
  
      image = tf.reshape(image, [mosaic_size, mosaic_size, h, w, c])
      image = tf.transpose(image, [0, 2, 1, 3, 4])
      image = tf.reshape(image, [h * mosaic_size, w * mosaic_size, c])
      return image
  
file_path: scenic/scenic/projects/owl_vit/preprocessing/mosaic.py
imports: import dataclasses
  from typing import Tuple
  from scenic.projects.owl_vit.preprocessing import image_ops
  from scenic.projects.owl_vit.preprocessing import modalities
  import tensorflow
label_desc: Use image node_id 1 for calling a global function with example usage: _image_tiles_to_mosaic(image, mosaic_size) and returns: image