class_docstr: General n-dimensional transposed convolution (aka. deconvolution).
class_docstr_tok: ["General","n-dimensional","transposed","convolution","(","aka",".","deconvolution",")","."]
docstr: Computes the transposed convolution of the input.
  
  Args:
    inputs: An array of shape ``[spatial_dims, C]`` and rank-N+1 if unbatched,
      or an array of shape ``[N, spatial_dims, C]`` and rank-N+2 if batched.
    precision: Optional :class:`jax.lax.Precision` to pass to
      :func:`jax.lax.conv_transpose`.
  
  Returns:
    An array of shape ``[spatial_dims, output_channels]`` and rank-N+1 if
      unbatched, or an array of shape ``[N, spatial_dims, output_channels]``
      and rank-N+2 if batched.
docstr_tok: ["Computes","the","transposed","convolution","of","the","input",".","Args",":","inputs",":","An","array","of","shape","``","[","spatial_dims",",","C","]","``","and","rank-N+1","if","unbatched",",","or","an","array","of","shape","``","[","N",",","spatial_dims",",","C","]","``","and","rank-N+2","if","batched",".","precision",":","Optional",":","class",":","`","jax.lax.Precision","`","to","pass","to",":","func",":","`","jax.lax.conv_transpose","`",".","Returns",":","An","array","of","shape","``","[","spatial_dims",",","output_channels","]","``","and","rank-N+1","if","unbatched",",","or","an","array","of","shape","``","[","N",",","spatial_dims",",","output_channels","]","``","and","rank-N+2","if","batched","."]
code: def __call__(
      self,
      inputs: jax.Array,
      *,
      precision: Optional[lax.Precision] = None,
  ) -> jax.Array:
      """Computes the transposed convolution of the input.
  
      Args:
        inputs: An array of shape ``[spatial_dims, C]`` and rank-N+1 if unbatched,
          or an array of shape ``[N, spatial_dims, C]`` and rank-N+2 if batched.
        precision: Optional :class:`jax.lax.Precision` to pass to
          :func:`jax.lax.conv_transpose`.
  
      Returns:
        An array of shape ``[spatial_dims, output_channels]`` and rank-N+1 if
          unbatched, or an array of shape ``[N, spatial_dims, output_channels]``
          and rank-N+2 if batched.
      """
      unbatched_rank = self.num_spatial_dims + 1
      allowed_ranks = [unbatched_rank, unbatched_rank + 1]
      if inputs.ndim not in allowed_ranks:
          raise ValueError(
              f"Input to ConvNDTranspose needs to have rank in "
              f"{allowed_ranks}, but input has shape {inputs.shape}."
          )
  
      unbatched = inputs.ndim == unbatched_rank
      if unbatched:
          inputs = jnp.expand_dims(inputs, axis=0)
  
      input_channels = inputs.shape[self.channel_index]
      w_shape = self.kernel_shape + (
          self.output_channels,
          input_channels,
      )
  
      if self.mask is not None and self.mask.shape != w_shape:
          raise ValueError(
              "Mask needs to have the same shape as weights. "
              f"Shapes are: {self.mask.shape}, {w_shape}"
          )
  
      w_init = self.w_init
      if w_init is None:
          fan_in_shape = self.kernel_shape + (input_channels,)
          stddev = 1.0 / np.sqrt(np.prod(fan_in_shape))
          w_init = hk.initializers.TruncatedNormal(stddev=stddev)
      w = hk.get_parameter("w", w_shape, inputs.dtype, init=w_init)
  
      if self.mask is not None:
          w = w * self.mask
  
      padding = self.padding
      if self.output_shape is not None:
          input_shape = (
              inputs.shape[2:]
              if self.channel_index == 1
              else inputs.shape[1:-1]
          )
          padding = tuple(
              map(
                  lambda i, o, k, s: compute_adjusted_padding(
                      i, o, k, s, self.padding
                  ),
                  input_shape,
                  self.output_shape,
                  self.kernel_shape,
                  self.stride,
              )
          )
  
      out = lax.conv_transpose(
          inputs,
          w,
          strides=self.stride,
          padding=padding,
          dimension_numbers=self.dimension_numbers,
          precision=precision,
      )
  
      if self.with_bias:
          if self.channel_index == -1:
              bias_shape = (self.output_channels,)
          else:
              bias_shape = (self.output_channels,) + (
                  1,
              ) * self.num_spatial_dims
          b = hk.get_parameter(
              "b", bias_shape, inputs.dtype, init=self.b_init
          )
          b = jnp.broadcast_to(b, out.shape)
          out = out + b
  
      if unbatched:
          out = jnp.squeeze(out, axis=0)
      return out
  
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 2 for calling the ConvNDTranspose obj's underlying member method code with example usage: obj.__call__(inputs) and returns: out