class_docstr: Convolution Module wrapping lax.conv_transpose.
  
  Example usage::
  
    >>> import flax.linen as nn
    >>> import jax, jax.numpy as jnp
  
    >>> # valid padding
    >>> layer = nn.ConvTranspose(features=4, kernel_size=(3,), padding='VALID')
    >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 8, 3)))
    >>> jax.tree_map(jnp.shape, variables)
    {'params': {'bias': (4,), 'kernel': (3, 3, 4)}}
    >>> out.shape
    (1, 10, 4)
    >>> # circular padding with stride 2
    >>> layer = nn.ConvTranspose(features=4, kernel_size=(6, 6), strides=(2, 2), padding='CIRCULAR', transpose_kernel=True)
    >>> out, variables = layer.init_with_output(jax.random.key(0), jnp.ones((1, 15, 15, 3)))
    >>> jax.tree_map(jnp.shape, variables)
    {'params': {'bias': (4,), 'kernel': (6, 6, 4, 3)}}
    >>> out.shape
    (1, 30, 30, 4)
    >>> # apply lower triangle mask
    >>> mask = jnp.tril(jnp.ones((3, 3, 4)))
    >>> layer = nn.ConvTranspose(features=4, kernel_size=(3,), mask=mask, padding='VALID')
    >>> variables = layer.init(jax.random.key(0), jnp.ones((1, 8, 3)))
  
  Attributes:
    features: number of convolution filters.
    kernel_size: shape of the convolutional kernel. For 1D convolution,
      the kernel size can be passed as an integer, which will be interpreted
      as a tuple of the single integer. For all other cases, it must be a
      sequence of integers.
    strides: a sequence of `n` integers, representing the inter-window strides.
    padding: either the string `'SAME'`, the string `'VALID'`, the string
      `'CIRCULAR'` (periodic boundary conditions), or a sequence of `n` `(low,
      high)` integer pairs that give the padding to apply before and after each
      spatial dimension. A single int is interpreted as applying the same padding
      in all dims and assign a single int in a sequence causes the same padding
      to be used on both sides.
    kernel_dilation: ``None``, or a sequence of ``n`` integers, giving the
      dilation factor to apply in each spatial dimension of the convolution
      kernel. Convolution with kernel dilation is also known as 'atrous
      convolution'.
    use_bias: whether to add a bias to the output (default: True).
    mask: Optional mask for the weights during masked convolution. The mask must
          be the same shape as the convolution weight matrix.
    dtype: the dtype of the computation (default: infer from input and params).
    param_dtype: the dtype passed to parameter initializers (default: float32).
    precision: numerical precision of the computation see ``jax.lax.Precision``
      for details.
    kernel_init: initializer for the convolutional kernel.
    bias_init: initializer for the bias.
    transpose_kernel: if True flips spatial axes and swaps the input/output
      channel axes of the kernel.
docstr: Applies a transposed convolution to the inputs.
  
  Behaviour mirrors of ``jax.lax.conv_transpose``.
  
  Args:
    inputs: input data with dimensions (*batch_dims, spatial_dims...,
      features). This is the channels-last convention, i.e. NHWC for a 2d
      convolution and NDHWC for a 3D convolution. Note: this is different from
      the input convention used by ``lax.conv_general_dilated``, which puts the
      spatial dimensions last.
      Note: If the input has more than 1 batch dimension, all batch dimensions
      are flattened into a single dimension for the convolution and restored
      before returning.  In some cases directly vmap'ing the layer may yield
      better performance than this default flattening approach.  If the input
      lacks a batch dimension it will be added for the convolution and removed
      n return, an allowance made to enable writing single-example code.
  
  Returns:
    The convolved data.
code: def __call__(self, inputs: Array) -> Array:
      """Applies a transposed convolution to the inputs.
  
      Behaviour mirrors of ``jax.lax.conv_transpose``.
  
      Args:
        inputs: input data with dimensions (*batch_dims, spatial_dims...,
          features). This is the channels-last convention, i.e. NHWC for a 2d
          convolution and NDHWC for a 3D convolution. Note: this is different from
          the input convention used by ``lax.conv_general_dilated``, which puts the
          spatial dimensions last.
          Note: If the input has more than 1 batch dimension, all batch dimensions
          are flattened into a single dimension for the convolution and restored
          before returning.  In some cases directly vmap'ing the layer may yield
          better performance than this default flattening approach.  If the input
          lacks a batch dimension it will be added for the convolution and removed
          n return, an allowance made to enable writing single-example code.
  
      Returns:
        The convolved data.
      """
      kernel_size: Tuple[int, ...]
      if isinstance(self.kernel_size, int):
          kernel_size = (self.kernel_size,)
      else:
          kernel_size = tuple(self.kernel_size)
  
      # Combine all input batch dimensions into a single leading batch axis.
      num_batch_dimensions = inputs.ndim - (len(kernel_size) + 1)
      if num_batch_dimensions != 1:
          input_batch_shape = inputs.shape[:num_batch_dimensions]
          total_batch_size = int(np.prod(input_batch_shape))
          flat_input_shape = (total_batch_size,) + inputs.shape[
              num_batch_dimensions:
          ]
          inputs = jnp.reshape(inputs, flat_input_shape)
  
      strides: Tuple[int, ...]
      if self.strides is None:
          strides = (1,) * (inputs.ndim - 2)
      else:
          strides = tuple(self.strides)
  
      in_features = jnp.shape(inputs)[-1]
      if self.transpose_kernel:
          kernel_shape = kernel_size + (self.features, in_features)
      else:
          kernel_shape = kernel_size + (in_features, self.features)
  
      if self.mask is not None and self.mask.shape != kernel_shape:
          raise ValueError(
              "Mask needs to have the same shape as weights. "
              f"Shapes are: {self.mask.shape}, {kernel_shape}"
          )
  
      kernel = self.param(
          "kernel", self.kernel_init, kernel_shape, self.param_dtype
      )
  
      if self.mask is not None:
          kernel *= self.mask
  
      padding_lax = canonicalize_padding(self.padding, len(kernel_size))
      if padding_lax == "CIRCULAR":
          padding_lax = "VALID"
  
      if self.use_bias:
          bias = self.param(
              "bias", self.bias_init, (self.features,), self.param_dtype
          )
      else:
          bias = None
  
      inputs, kernel, bias = promote_dtype(
          inputs, kernel, bias, dtype=self.dtype
      )
  
      y = lax.conv_transpose(
          inputs,
          kernel,
          strides,
          padding_lax,
          rhs_dilation=self.kernel_dilation,
          transpose_kernel=self.transpose_kernel,
          precision=self.precision,
      )
  
      if self.padding == "CIRCULAR":
          # For circular padding, we need to identify the size of the final output
          # ("period") along each spatial dimension, pad each dimension to an
          # integer number of periods, and wrap the array periodically around each
          # dimension. Padding should be done in such a way that the start of the
          # original input data inside the padded array is located at integer
          # number of periods - otherwise the result would be circularly shifted.
  
          # Compute period along each spatial dimension - it's input size scaled
          # by the stride.
          scaled_x_dims = [
              x_dim * stride
              for x_dim, stride in zip(jnp.shape(inputs)[1:-1], strides)
          ]
          # Compute difference between the current size of y and the final output
          # size, and complement this difference to 2 * period - that gives how
          # much we need to pad.
          size_diffs = [
              -(y_dim - x_dim) % (2 * x_dim)
              for y_dim, x_dim in zip(y.shape[1:-1], scaled_x_dims)
          ]
          if self.transpose_kernel:
              # If the kernel is transposed, the "+1" is put on the right to
              # mirror the regular convolution. If the same kernel parameters are used
              # as for Conv, this layer then computes the proper transpose convolution.
              total_pad = [
                  (size_diff // 2, (size_diff + 1) // 2)
                  for size_diff in size_diffs
              ]
          else:
              # Divide the padding equally between left and right. The choice to put
              # "+1" on the left (and not on the right) represents a convention for
              # aligning even-sized kernels.
              total_pad = [
                  ((size_diff + 1) // 2, size_diff // 2)
                  for size_diff in size_diffs
              ]
          y = jnp.pad(y, [(0, 0)] + total_pad + [(0, 0)])
          # Wrap the result periodically around each spatial dimension,
          # one by one.
          for i in range(1, y.ndim - 1):
              y = y.reshape(
                  y.shape[:i]
                  + (-1, scaled_x_dims[i - 1])
                  + y.shape[i + 1 :]
              )
              y = y.sum(axis=i)
  
      if self.use_bias:
          y += jnp.reshape(bias, (1,) * (y.ndim - 1) + (-1,))
  
      if num_batch_dimensions != 1:
          output_shape = input_batch_shape + y.shape[1:]
          y = jnp.reshape(y, output_shape)
  
      return y
  
imports: import dataclasses
  from typing import Any, Callable, Iterable, List, Optional, Sequence, Tuple, Union
  import jax
  import jax.numpy
  import numpy
  from jax import eval_shape, lax
  from jax.core import ShapedArray
  from flax.core import meta
  from flax.linen import initializers
  from flax.linen.dtypes import promote_dtype
  from flax.linen.module import Module, compact
label_desc: Use image node_id 1 for calling the ConvTranspose obj's underlying member method code with example usage: obj.__call__(inputs) and returns: y