class_docstr: General N-dimensional transposed 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 transposed convolutional kernel.
  - `stride`: The stride used on the equivalent [`equinox.nn.Conv`][].
  - `padding`: The amount of padding used on the equivalent [`equinox.nn.Conv`][].
  - `output_padding`: Additional padding for the output shape.
  - `dilation`: The spacing between kernel points.
  - `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
      `ConvTranspose` 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 transposed convolution.
  - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
      initialisation. (Keyword only argument.)
  
  !!! info
  
      All of `kernel_size`, `stride`, `padding`, `output_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.
  
  !!! tip
  
      Transposed convolutions are often used to go in the "opposite direction" to
      a normal convolution. That is, from something with the shape of the output
      of a convolution to something with the shape of the input to a convolution.
      Moreover, to do so with the same "connectivity", i.e. which inputs can
      affect which outputs.
  
      Relative to an [`equinox.nn.Conv`][] layer, this can be accomplished by
      switching the values of `in_channels` and `out_channels`, whilst keeping
      `kernel_size`, `stride`, `padding`, `dilation`, and `groups` the same.
  
      When `stride > 1` then [`equinox.nn.Conv`][] maps multiple input shapes to the
      same output shape. `output_padding` is provided to resolve this ambiguity,
      by adding a little extra padding to just the bottom/right edges of the
      input.
  
      See [these animations](https://github.com/vdumoulin/conv_arithmetic/blob/af6f818b0bb396c26da79899554682a8a499101d/README.md#transposed-convolution-animations)
      and [this report](https://arxiv.org/abs/1603.07285) for a nice reference.
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,
      output_padding: Union[int, Sequence[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 transposed convolutional kernel.
      - `stride`: The stride used on the equivalent [`equinox.nn.Conv`][].
      - `padding`: The amount of padding used on the equivalent [`equinox.nn.Conv`][].
      - `output_padding`: Additional padding for the output shape.
      - `dilation`: The spacing between kernel points.
      - `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
          `ConvTranspose` 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 transposed convolution.
      - `key`: A `jax.random.PRNGKey` used to provide randomness for parameter
          initialisation. (Keyword only argument.)
  
      !!! info
  
          All of `kernel_size`, `stride`, `padding`, `output_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.
  
      !!! tip
  
          Transposed convolutions are often used to go in the "opposite direction" to
          a normal convolution. That is, from something with the shape of the output
          of a convolution to something with the shape of the input to a convolution.
          Moreover, to do so with the same "connectivity", i.e. which inputs can
          affect which outputs.
  
          Relative to an [`equinox.nn.Conv`][] layer, this can be accomplished by
          switching the values of `in_channels` and `out_channels`, whilst keeping
          `kernel_size`, `stride`, `padding`, `dilation`, and `groups` the same.
  
          When `stride > 1` then [`equinox.nn.Conv`][] maps multiple input shapes to the
          same output shape. `output_padding` is provided to resolve this ambiguity,
          by adding a little extra padding to just the bottom/right edges of the
          input.
  
          See [these animations](https://github.com/vdumoulin/conv_arithmetic/blob/af6f818b0bb396c26da79899554682a8a499101d/README.md#transposed-convolution-animations)
          and [this report](https://arxiv.org/abs/1603.07285) for a nice reference.
      """  # noqa: E501
  
      wkey, bkey = jrandom.split(key, 2)
  
      parse = _ntuple(num_spatial_dims)
      kernel_size = parse(kernel_size)
      stride = parse(stride)
      output_padding = parse(output_padding)
      dilation = parse(dilation)
  
      for s, o in zip(stride, output_padding):
          if output_padding >= stride:
              raise ValueError(
                  "Must have `output_padding < stride` (elementwise)."
              )
  
      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.output_padding = output_padding
      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 ConvTranspose object from inherited base classes: Module with example: obj = ConvTranspose(num_spatial_dims, in_channels, out_channels, kernel_size, stride, padding, output_padding, dilation, groups, use_bias)