code: def conv_general_dilated_patches(
      lhs: jax.typing.ArrayLike,
      filter_shape: Sequence[int],
      window_strides: Sequence[int],
      padding: str | Sequence[tuple[int, int]],
      lhs_dilation: Sequence[int] | None = None,
      rhs_dilation: Sequence[int] | None = None,
      dimension_numbers: convolution.ConvGeneralDilatedDimensionNumbers
      | None = None,
      precision: lax.PrecisionType | None = None,
      preferred_element_type: DType | None = None,
  ) -> jax.Array:
      """Extract patches subject to the receptive field of `conv_general_dilated`.
  
      Runs the input through a convolution with given parameters. The kernel of the
      convolution is constructed such that the output channel dimension `"C"`
      contains flattened image patches, so instead a single `"C"` dimension
      represents, for example, three dimensions `"chw"` collapsed. The order of
      these dimensions is `"c" + ''.join(c for c in rhs_spec if c not in 'OI')`,
      where `rhs_spec == dimension_numbers[1]`, and the size of this `"C"`
      dimension is therefore the size of each patch, i.e.
      `np.prod(filter_shape) * lhs.shape[lhs_spec.index('C')]`, where
      `lhs_spec == dimension_numbers[0]`.
  
      Docstring below adapted from `jax.lax.conv_general_dilated`.
  
      See Also:
        https://www.tensorflow.org/xla/operation_semantics#conv_convolution
  
      Args:
        lhs: a rank `n+2` dimensional input array.
        filter_shape: a sequence of `n` integers, representing the receptive window
          spatial shape in the order as specified in
          `rhs_spec = dimension_numbers[1]`.
        window_strides: a sequence of `n` integers, representing the inter-window
          strides.
        padding: either the string `'SAME'`, the string `'VALID'`, or a sequence of
          `n` `(low, high)` integer pairs that give the padding to apply before and
          after each spatial dimension.
        lhs_dilation: `None`, or a sequence of `n` integers, giving the
          dilation factor to apply in each spatial dimension of `lhs`. LHS dilation
          is also known as transposed convolution.
        rhs_dilation: `None`, or a sequence of `n` integers, giving the
          dilation factor to apply in each spatial dimension of `rhs`. RHS dilation
          is also known as atrous convolution.
        dimension_numbers: either `None`, or a 3-tuple
          `(lhs_spec, rhs_spec, out_spec)`, where each element is a string
          of length `n+2`. `None` defaults to `("NCHWD..., OIHWD..., NCHWD...")`.
        precision: Optional. Either ``None``, which means the default precision for
          the backend, or a :class:`~jax.lax.Precision` enum value (``Precision.DEFAULT``,
          ``Precision.HIGH`` or ``Precision.HIGHEST``).
        preferred_element_type: Optional. Either ``None``, which means the default
          accumulation type for the input types, or a datatype, indicating to
          accumulate results to and return a result with that datatype.
  
      Returns:
        A rank `n+2` array containing the flattened image patches in the output
        channel (`"C"`) dimension. For example if
        `dimension_numbers = ("NcHW", "OIwh", "CNHW")`, the output has dimension
        numbers `"CNHW" = "{cwh}NHW"`, with the size of dimension `"C"` equal to
        the size of each patch
        (`np.prod(filter_shape) * lhs.shape[lhs_spec.index('C')]`).
  
      """
      lhs_array = jnp.asarray(lhs)
      filter_shape = tuple(filter_shape)
      dimension_numbers = convolution.conv_dimension_numbers(
          lhs_array.shape, (1, 1) + filter_shape, dimension_numbers
      )
  
      lhs_spec, rhs_spec, out_spec = dimension_numbers
  
      spatial_size = math.prod(filter_shape)
      n_channels = lhs_array.shape[lhs_spec[1]]
  
      # Move separate `lhs` spatial locations into separate `rhs` channels.
      rhs = jnp.eye(spatial_size, dtype=lhs_array.dtype).reshape(
          filter_shape * 2
      )
  
      rhs = rhs.reshape((spatial_size, 1) + filter_shape)
      rhs = jnp.tile(rhs, (n_channels,) + (1,) * (rhs.ndim - 1))
      rhs = jnp.moveaxis(rhs, (0, 1), (rhs_spec[0], rhs_spec[1]))
  
      out = convolution.conv_general_dilated(
          lhs=lhs_array,
          rhs=rhs,
          window_strides=window_strides,
          padding=padding,
          lhs_dilation=lhs_dilation,
          rhs_dilation=rhs_dilation,
          dimension_numbers=dimension_numbers,
          precision=None
          if precision is None
          else (precision, lax.Precision.DEFAULT),
          feature_group_count=n_channels,
          preferred_element_type=preferred_element_type,
      )
      return out
  
file_path: jax/jax/_src/lax/other.py
imports: from __future__ import annotations
  from collections.abc import Sequence
  import math
  from typing import Any, Union, cast
  import jax
  from jax._src.numpy import lax_numpy
  from jax._src.lax import lax
  from jax._src.lax import convolution
label_desc: Use image node_id 1 for calling a global function with example usage: conv_general_dilated_patches(lhs, filter_shape, window_strides, padding, lhs_dilation, rhs_dilation, dimension_numbers, precision, preferred_element_type) and returns: out