code: def patch_image(
      inputs,
      inputs_shape,
      patch_size,
      strides=None,
      padding="VALID",
      mode="i2p",
  ):
      """Applies patching operation on the input.
  
      Args:
        inputs: Input data.
        inputs_shape: tuple; Shape of the input data.
        patch_size: tuple; size of the patch: (height, width).
        strides: tuple; Specifies how far two consecutive patches are in the
          input.
        padding: str; The type of padding algorithm to use.
        mode: str; Either 'i2p' to convert the input image to patches or 'p2i' to
          convert the patched image to the original shape.
  
      Returns:
        Patched image if mode='i2p', original image if mode='p2i'.
      """
      strides = strides or patch_size
  
      def i2p(x):
          return extract_image_patches(
              lhs=x.astype(jnp.float64),
              rhs_shape=(1,) + patch_size + (1,),
              strides=(1,) + strides + (1,),
              padding=padding,
              rhs_dilation=(1,) * inputs.ndim,
              data_format="NHWC",
          )
  
      if mode == "i2p":
          _, inputs_w, inputs_h, _ = inputs.shape
          patch_w, patch_h = patch_size
          if inputs_w < patch_w or inputs_h < patch_h:
              raise ValueError(
                  f"Patch height and width ({patch_w} and  {patch_h}) "
                  "should be smaller thatn inputs height and width"
                  f" ({inputs_w} and  {inputs_h})."
              )
          outputs = i2p(inputs)
  
      elif mode == "p2i":
          _, fn_vjp = jax.vjp(i2p, jnp.ones(inputs_shape))
          overlap_count = fn_vjp(jnp.ones_like(inputs))[0]
          outputs = fn_vjp(inputs)[0] / overlap_count
  
      else:
          raise ValueError()
      return outputs
  
file_path: scenic/scenic/model_lib/layers/nn_ops.py
imports: from typing import Callable, Sequence
  import flax.linen
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 4 for calling a global function with example usage: patch_image(inputs, inputs_shape, patch_size, strides, padding, mode) and returns: outputs, extract_image_patches