code: def extract_image_patches(
      lhs, rhs_shape, strides, padding, rhs_dilation, data_format="NHWC"
  ):
      """Extract patches of size `rhs_shape` from `lhs`.
  
      Args:
        lhs: A 4-D Tensor;  With shape `[batch, in_rows, in_cols, depth].
        rhs_shape: tuple; Size of the sliding window for each dimension of `lhs`.
        strides: tuple; How far the centers of two consecutive patches are in the
          lhs. Must be: `[1, stride_rows, stride_cols, 1]`.
        padding: str; The type of padding algorithm to use.
          We specify the size-related attributes as: ```python ksizes = [1,
            ksize_rows, ksize_cols, 1] strides = [1, strides_rows, strides_cols, 1]
            rates = [1, rates_rows, rates_cols, 1]```
        rhs_dilation: A 1-D Tensor of length 4; Must be: `[1, rate_rows, rate_cols,
          1]`. This is the input stride, specifying how far two consecutive patch
          samples are in the input. Equivalent to extracting patches with
          `patch_sizes_eff = patch_sizes + (patch_sizes - 1) * (rates - 1)`,
          followed by subsampling them spatially by a factor of `rates`. This is
          equivalent to `rate` in dilated (a.k.a. Atrous) convolutions.
        data_format: str; The format of the `lhs`. Must be either `'NHWC'` or
          `'NCHW'`.
  
      Returns:
        A 4-D Tensor. Has the same type and data format as `lhs`, and with shape
        `[batch, num_patches_col, num_patches_row, rhs_shape[1], rhs_shape[2], C]`.
      """
      num_dims = lhs.ndim
      num_spatial_dims = num_dims - 2
  
      batch_dim = data_format.index("N")
      feature_dim = data_format.index("C")
      depth = lhs.shape[feature_dim]
  
      if rhs_shape[batch_dim] != 1 or rhs_shape[feature_dim] != 1:
          raise NotImplementedError(
              "Current implementation does not yet support window sizes > 1 in "
              "the batch and depth dimensions."
          )
  
      if strides[batch_dim] != 1 or strides[feature_dim] != 1:
          raise NotImplementedError(
              "Current implementation does not support strides in the batch "
              "and depth dimensions."
          )
  
      if rhs_dilation[batch_dim] != 1 or rhs_dilation[feature_dim] != 1:
          raise NotImplementedError(
              "Current implementation does not support dilations in the batch "
              "and depth dimensions."
          )
  
      # Replicating tensorflow's implementation.
      lhs_perm = lax.conv_general_permutations(
          (data_format, "HWIO", data_format)
      )[0]
      kernel_shape = [rhs_shape[i] for i in lhs_perm[2:]]
  
      kernel_size = np.prod(kernel_shape)
      conv_filter_shape = kernel_shape[:]
      conv_filter_shape.append(1)
      conv_filter_shape.append(kernel_size * depth)
  
      iota_kernel_shape = (kernel_size, depth, kernel_size)
  
      conv_filter = lax.eq(
          lax.broadcasted_iota(jnp.int32, iota_kernel_shape, 0),
          lax.broadcasted_iota(jnp.int32, iota_kernel_shape, 2),
      )
      conv_filter = lax.convert_element_type(conv_filter, lhs.dtype)
      conv_filter = lax.reshape(conv_filter, conv_filter_shape)
  
      dim_num = lax.conv_dimension_numbers(
          lhs.shape,
          conv_filter.shape,
          (data_format, "HWIO", data_format),
      )
      conv_strides = [0] * num_spatial_dims
      conv_rhs_dilation = [0] * num_spatial_dims
      for i in range(num_spatial_dims):
          dim = dim_num.lhs_spec[i + 2]
          conv_strides[i] = strides[dim]
          conv_rhs_dilation[i] = rhs_dilation[dim]
  
      conv = lax.conv_general_dilated(
          lhs,
          conv_filter,
          conv_strides,
          padding,
          None,
          conv_rhs_dilation,
          dim_num,
          depth,
      )
  
      conv_dims = list(conv.shape[:-1])
      conv_dims.append(depth)
      conv_dims.extend(kernel_shape)
      conv = lax.reshape(conv, conv_dims)
  
      permutation = list(range(len(conv_dims)))
      depth_dim = permutation.pop(-3)
      permutation.append(depth_dim)
  
      return lax.transpose(conv, permutation)
  
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 1 for calling a global function with example usage: extract_image_patches(lhs, rhs_shape, strides, padding, rhs_dilation, data_format) and returns: lax