class_docstr: General N-dimensional convolutional.
class_docstr_tok: ["General","N-dimensional","convolutional","."]
docstr: Initializes the module.
  
  Args:
    num_spatial_dims: The number of spatial dimensions of the input.
    output_channels: Number of output channels.
    kernel_shape: The shape of the kernel. Either an integer or a sequence of
      length ``num_spatial_dims``.
    stride: Optional stride for the kernel. Either an integer or a sequence of
      length ``num_spatial_dims``. Defaults to 1.
    rate: Optional kernel dilation rate. Either an integer or a sequence of
      length ``num_spatial_dims``. 1 corresponds to standard ND convolution,
      ``rate > 1`` corresponds to dilated convolution. Defaults to 1.
    padding: Optional padding algorithm. Either ``VALID`` or ``SAME`` or a
      sequence of n ``(low, high)`` integer pairs that give the padding to
      apply before and after each spatial dimension. or a callable or sequence
      of callables of size ``num_spatial_dims``. Any callables must take a
      single integer argument equal to the effective kernel size and return a
      sequence of two integers representing the padding before and after. See
      ``haiku.pad.*`` for more details and example functions. Defaults to
      ``SAME``. See:
      https://www.tensorflow.org/xla/operation_semantics#conv_convolution.
    with_bias: Whether to add a bias. By default, true.
    w_init: Optional weight initialization. By default, truncated normal.
    b_init: Optional bias initialization. By default, zeros.
    data_format: The data format of the input.  Can be either
      ``channels_first``, ``channels_last``, ``N...C`` or ``NC...``. By
      default, ``channels_last``. See :func:`get_channel_index`.
    mask: Optional mask of the weights.
    feature_group_count: Optional number of groups in group convolution.
      Default value of 1 corresponds to normal dense convolution. If a higher
      value is used, convolutions are applied separately to that many groups,
      then stacked together. This reduces the number of parameters
      and possibly the compute for a given ``output_channels``. See:
      https://www.tensorflow.org/xla/operation_semantics#conv_convolution.
    name: The name of the module.
docstr_tok: ["Initializes","the","module",".","Args",":","num_spatial_dims",":","The","number","of","spatial","dimensions","of","the","input",".","output_channels",":","Number","of","output","channels",".","kernel_shape",":","The","shape","of","the","kernel",".","Either","an","integer","or","a","sequence","of","length","``","num_spatial_dims","``",".","stride",":","Optional","stride","for","the","kernel",".","Either","an","integer","or","a","sequence","of","length","``","num_spatial_dims","``",".","Defaults","to","1.","rate",":","Optional","kernel","dilation","rate",".","Either","an","integer","or","a","sequence","of","length","``","num_spatial_dims","``",".","1","corresponds","to","standard","ND","convolution",",","``","rate",">","1","``","corresponds","to","dilated","convolution",".","Defaults","to","1.","padding",":","Optional","padding","algorithm",".","Either","``","VALID","``","or","``","SAME","``","or","a","sequence","of","n","``","(","low",",","high",")","``","integer","pairs","that","give","the","padding","to","apply","before","and","after","each","spatial","dimension",".","or","a","callable","or","sequence","of","callables","of","size","``","num_spatial_dims","``",".","Any","callables","must","take","a","single","integer","argument","equal","to","the","effective","kernel","size","and","return","a","sequence","of","two","integers","representing","the","padding","before","and","after",".","See","``","haiku.pad",".","*","``","for","more","details","and","example","functions",".","Defaults","to","``","SAME","``",".","See",":","https",":","\/\/www.tensorflow.org\/xla\/operation_semantics","#","conv_convolution",".","with_bias",":","Whether","to","add","a","bias",".","By","default",",","true",".","w_init",":","Optional","weight","initialization",".","By","default",",","truncated","normal",".","b_init",":","Optional","bias","initialization",".","By","default",",","zeros",".","data_format",":","The","data","format","of","the","input",".","Can","be","either","``","channels_first","``",",","``","channels_last","``",",","``","N","...","C","``","or","``","NC","...","``",".","By","default",",","``","channels_last","``",".","See",":","func",":","`","get_channel_index","`",".","mask",":","Optional","mask","of","the","weights",".","feature_group_count",":","Optional","number","of","groups","in","group","convolution",".","Default","value","of","1","corresponds","to","normal","dense","convolution",".","If","a","higher","value","is","used",",","convolutions","are","applied","separately","to","that","many","groups",",","then","stacked","together",".","This","reduces","the","number","of","parameters","and","possibly","the","compute","for","a","given","``","output_channels","``",".","See",":","https",":","\/\/www.tensorflow.org\/xla\/operation_semantics","#","conv_convolution",".","name",":","The","name","of","the","module","."]
code: def __init__(
      self,
      num_spatial_dims: int,
      output_channels: int,
      kernel_shape: Union[int, Sequence[int]],
      stride: Union[int, Sequence[int]] = 1,
      rate: Union[int, Sequence[int]] = 1,
      padding: Union[
          str,
          Sequence[tuple[int, int]],
          hk.pad.PadFn,
          Sequence[hk.pad.PadFn],
      ] = "SAME",
      with_bias: bool = True,
      w_init: Optional[hk.initializers.Initializer] = None,
      b_init: Optional[hk.initializers.Initializer] = None,
      data_format: str = "channels_last",
      mask: Optional[jax.Array] = None,
      feature_group_count: int = 1,
      name: Optional[str] = None,
  ):
      """Initializes the module.
  
      Args:
        num_spatial_dims: The number of spatial dimensions of the input.
        output_channels: Number of output channels.
        kernel_shape: The shape of the kernel. Either an integer or a sequence of
          length ``num_spatial_dims``.
        stride: Optional stride for the kernel. Either an integer or a sequence of
          length ``num_spatial_dims``. Defaults to 1.
        rate: Optional kernel dilation rate. Either an integer or a sequence of
          length ``num_spatial_dims``. 1 corresponds to standard ND convolution,
          ``rate > 1`` corresponds to dilated convolution. Defaults to 1.
        padding: Optional padding algorithm. Either ``VALID`` or ``SAME`` or a
          sequence of n ``(low, high)`` integer pairs that give the padding to
          apply before and after each spatial dimension. or a callable or sequence
          of callables of size ``num_spatial_dims``. Any callables must take a
          single integer argument equal to the effective kernel size and return a
          sequence of two integers representing the padding before and after. See
          ``haiku.pad.*`` for more details and example functions. Defaults to
          ``SAME``. See:
          https://www.tensorflow.org/xla/operation_semantics#conv_convolution.
        with_bias: Whether to add a bias. By default, true.
        w_init: Optional weight initialization. By default, truncated normal.
        b_init: Optional bias initialization. By default, zeros.
        data_format: The data format of the input.  Can be either
          ``channels_first``, ``channels_last``, ``N...C`` or ``NC...``. By
          default, ``channels_last``. See :func:`get_channel_index`.
        mask: Optional mask of the weights.
        feature_group_count: Optional number of groups in group convolution.
          Default value of 1 corresponds to normal dense convolution. If a higher
          value is used, convolutions are applied separately to that many groups,
          then stacked together. This reduces the number of parameters
          and possibly the compute for a given ``output_channels``. See:
          https://www.tensorflow.org/xla/operation_semantics#conv_convolution.
        name: The name of the module.
      """
      super().__init__(name=name)
      if num_spatial_dims <= 0:
          raise ValueError(
              "We only support convolution operations for `num_spatial_dims` "
              f"greater than 0, received num_spatial_dims={num_spatial_dims}."
          )
  
      self.num_spatial_dims = num_spatial_dims
      self.output_channels = output_channels
      self.kernel_shape = utils.replicate(
          kernel_shape, num_spatial_dims, "kernel_shape"
      )
      self.with_bias = with_bias
      self.stride = utils.replicate(stride, num_spatial_dims, "strides")
      self.w_init = w_init
      self.b_init = b_init or jnp.zeros
      self.mask = mask
      self.feature_group_count = feature_group_count
      self.lhs_dilation = utils.replicate(
          1, num_spatial_dims, "lhs_dilation"
      )
      self.kernel_dilation = utils.replicate(
          rate, num_spatial_dims, "kernel_dilation"
      )
      self.data_format = data_format
      self.channel_index = hk.get_channel_index(data_format)
      self.dimension_numbers = to_dimension_numbers(
          num_spatial_dims,
          channels_last=(self.channel_index == -1),
          transpose=False,
      )
  
      if isinstance(padding, str):
          self.padding = padding.upper()
      elif hk.pad.is_padfn(padding):
          self.padding = hk.pad.create_from_padfn(
              padding=padding,
              kernel=self.kernel_shape,
              rate=self.kernel_dilation,
              n=self.num_spatial_dims,
          )
      else:
          self.padding = hk.pad.create_from_tuple(
              padding, self.num_spatial_dims
          )
  
imports: from collections.abc import Sequence
  from typing import Optional, Union
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import pad
  from haiku._src import utils
  import jax
  from jax import lax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 to create a new ConvND object from inherited base classes: hk with example: obj = ConvND(num_spatial_dims, output_channels, kernel_shape, stride, rate, padding, with_bias, w_init, b_init, data_format, mask, feature_group_count, name)