conv_transpose

Contents

conv_transpose#

class brainunit.lax.conv_transpose(x, y, strides, padding, rhs_dilation=None, dimension_numbers=None, transpose_kernel=False, precision=None, preferred_element_type=None, **kwargs)#

Convenience wrapper for calculating the N-d convolution “transpose”.

This function directly calculates a fractionally strided conv rather than indirectly calculating the gradient (transpose) of a forward convolution.

Parameters:

  • lhs – a rank n+2 dimensional input array.

  • rhs – a rank n+2 dimensional array of kernel weights.

  • strides (Sequence[int]) – sequence of n integers, sets fractional stride.

  • padding (str | Sequence[tuple[int, int]]) – ‘SAME’, ‘VALID’ will set as transpose of corresponding forward conv, or a sequence of n integer 2-tuples describing before-and-after padding for each n spatial dimension.

  • rhs_dilation (Sequence[int] | None) – None, or a sequence of n integers, giving the dilation factor to apply in each spatial dimension of rhs. RHS dilation is also known as atrous convolution.

  • dimension_numbers (tuple[str, str, str] | ConvDimensionNumbers | None) – tuple of dimension descriptors as in lax.conv_general_dilated. Defaults to tensorflow convention.

  • transpose_kernel (bool) – if True flips spatial axes and swaps the input/output channel axes of the kernel. This makes the output of this function identical to the gradient-derived functions like keras.layers.Conv2DTranspose applied to the same kernel. For typical use in neural nets this is completely pointless and just makes input/output channel specification confusing.

  • precision (None | str | Precision | tuple[str, str] | tuple[Precision, Precision] | DotAlgorithm | DotAlgorithmPreset) – Optional. Either None, which means the default precision for the backend, a Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST) or a tuple of two Precision enums indicating precision of lhs` and rhs.

  • preferred_element_type (str | type[Any] | dtype | SupportsDType | None) – Optional. Either None, which means the default accumulation type for the input types, or a datatype, indicating to accumulate results to and return a result with that datatype.

Return type:

saiunit.Quantity | Array

Returns:

Transposed N-d convolution, with output padding following the conventions of keras.layers.Conv2DTranspose.

Contents