code: def extract_patches(lhs, rhs_shape, strides=(1, 1)):
      """Extracts patches from an image using a convolution operator.
  
      Args:
        lhs: A tensor of images of shapes (B, H, W, C).
        rhs_shape: The size of the patches to extract (h, w).
        strides: The shift between extracted patches (s1, s2)
  
      Returns:
        All the patches in a tensor of dimension
          (B, (H - h + 1) // s1, (W - w + 1) // s2, h, w, C).
      """
      # [batch, channels, height, width]
      lhs = jnp.moveaxis(lhs, -1, 1)
      d = lhs.shape[1]
      h, w = rhs_shape
  
      # Construct the lookup conv weights.
      dim_out = jnp.arange(d * h * w).reshape((-1, 1, 1, 1))
      dim_in = jnp.arange(d).reshape((1, -1, 1, 1))
      i = jnp.arange(h).reshape((1, 1, -1, 1))
      j = jnp.arange(w).reshape((1, 1, 1, -1))
      weights = ((w * i + j) * d + dim_in == dim_out).astype(
          jnp.float32
      )
  
      # [batch, h * w * d, (H - h + 1) // s1, (W - w + 1) // s2]
      concatenated_patches = lax.conv(
          lhs, weights, window_strides=strides, padding="VALID"
      )
  
      # [batch, (H - h + 1) // s1, (W - w + 1) // s2, h * w * d]
      concatenated_patches = jnp.moveaxis(concatenated_patches, 1, -1)
  
      # [batch, (H - h + 1) // s1, (W - w + 1) // s2, h, w, d]
      shape = concatenated_patches.shape[:3] + (h, w, d)
      return concatenated_patches.reshape(shape)
  
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 2 for calling a global function with example usage: extract_patches(lhs, rhs_shape, strides) and returns: concatenated_patches