class_docstr: General N-dimensional convolution.
docstr: **Arguments:**
  
  - `num_spatial_dims`: The number of spatial dimensions. For example traditional
      convolutions for image processing have this set to `2`.
  - `in_channels`: The number of input channels.
  - `out_channels`: The number of output channels.
  - `kernel_size`: The size of the convolutional kernel.
  - `stride`: The stride of the convolution.
  - `padding`: The amount of padding to apply before and after each spatial
      dimension.
  - `dilation`: The dilation of the convolution.
  - `groups`: The number of input channel groups. At `groups=1`,
      all input channels contribute to all output channels. Values
      higher than `1` are equivalent to running `groups` independent
      `Conv` operations side-by-side, each having access only to
      `in_channels` // `groups` input channels, and
      concatenating the results along the output channel dimension.
      `in_channels` must be divisible by `groups`.
  - `use_bias`: Whether to add on a bias after the convolution.
  - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
      initialisation. (Keyword only argument.)
  
  !!! info
  
      All of `kernel_size`, `stride`, `padding`, `dilation` can be either an
      integer or a sequence of integers. If they are a sequence then the sequence
      should be of length equal to `num_spatial_dims`, and specify the value of
      each property down each spatial dimension in turn.
  
      If they are an integer then the same kernel size / stride / padding /
      dilation will be used along every spatial dimension.
  
      `padding` can alternatively also be a sequence of 2-element tuples,
      each representing the padding to apply before and after each spatial
      dimension.
code: def __init__(
      self,
      num_spatial_dims: int,
      in_channels: int,
      out_channels: int,
      kernel_size: Union[int, Sequence[int]],
      stride: Union[int, Sequence[int]] = 1,
      padding: Union[int, Sequence[int], Sequence[tuple[int, int]]] = 0,
      dilation: Union[int, Sequence[int]] = 1,
      groups: int = 1,
      use_bias: bool = True,
      *,
      key: PRNGKeyArray,
  ):
      """**Arguments:**
  
      - `num_spatial_dims`: The number of spatial dimensions. For example traditional
          convolutions for image processing have this set to `2`.
      - `in_channels`: The number of input channels.
      - `out_channels`: The number of output channels.
      - `kernel_size`: The size of the convolutional kernel.
      - `stride`: The stride of the convolution.
      - `padding`: The amount of padding to apply before and after each spatial
          dimension.
      - `dilation`: The dilation of the convolution.
      - `groups`: The number of input channel groups. At `groups=1`,
          all input channels contribute to all output channels. Values
          higher than `1` are equivalent to running `groups` independent
          `Conv` operations side-by-side, each having access only to
          `in_channels` // `groups` input channels, and
          concatenating the results along the output channel dimension.
          `in_channels` must be divisible by `groups`.
      - `use_bias`: Whether to add on a bias after the convolution.
      - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
          initialisation. (Keyword only argument.)
  
      !!! info
  
          All of `kernel_size`, `stride`, `padding`, `dilation` can be either an
          integer or a sequence of integers. If they are a sequence then the sequence
          should be of length equal to `num_spatial_dims`, and specify the value of
          each property down each spatial dimension in turn.
  
          If they are an integer then the same kernel size / stride / padding /
          dilation will be used along every spatial dimension.
  
          `padding` can alternatively also be a sequence of 2-element tuples,
          each representing the padding to apply before and after each spatial
          dimension.
  
      """
      wkey, bkey = jrandom.split(key, 2)
  
      parse = _ntuple(num_spatial_dims)
      kernel_size = parse(kernel_size)
      stride = parse(stride)
      dilation = parse(dilation)
  
      if in_channels % groups != 0:
          raise ValueError(
              f"`in_channels` (={in_channels}) must be divisible "
              f"by `groups` (={groups})."
          )
  
      grouped_in_channels = in_channels // groups
      lim = 1 / np.sqrt(grouped_in_channels * math.prod(kernel_size))
      self.weight = jrandom.uniform(
          wkey,
          (out_channels, grouped_in_channels) + kernel_size,
          minval=-lim,
          maxval=lim,
      )
      if use_bias:
          self.bias = jrandom.uniform(
              bkey,
              (out_channels,) + (1,) * num_spatial_dims,
              minval=-lim,
              maxval=lim,
          )
      else:
          self.bias = None
  
      self.num_spatial_dims = num_spatial_dims
      self.in_channels = in_channels
      self.out_channels = out_channels
      self.kernel_size = kernel_size
      self.stride = stride
      if isinstance(padding, int):
          self.padding = tuple(
              (padding, padding) for _ in range(num_spatial_dims)
          )
      elif (
          isinstance(padding, Sequence)
          and len(padding) == num_spatial_dims
      ):
          if all_sequences(padding):
              self.padding = tuple(padding)
          else:
              self.padding = tuple((p, p) for p in padding)
      else:
          raise ValueError(
              "`padding` must either be an int or tuple of length "
              f"{num_spatial_dims} containing ints or tuples of length 2."
          )
      self.dilation = dilation
      self.groups = groups
      self.use_bias = use_bias
  
imports: import itertools
  import math
  from collections.abc import Callable, Sequence
  from typing import Optional, TypeVar, Union
  import jax
  import jax.lax
  import jax.numpy
  import jax.random
  import numpy
  from jaxtyping import Array, PRNGKeyArray
  from .._module import field, Module
  from ._misc import all_sequences
label_desc: Use image node_id 1 to create a new Conv object from inherited base classes: Module with example: obj = Conv(num_spatial_dims, in_channels, out_channels, kernel_size, stride, padding, dilation, groups, use_bias)