jax.lax.conv_with_general_padding#

jax.lax.conv_with_general_padding(lhs, rhs, window_strides, padding, lhs_dilation, rhs_dilation, precision=None, preferred_element_type=None)[source]#

Convenience wrapper around conv_general_dilated.

Parameters:

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

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

  • window_strides (Sequence[int]) – a sequence of n integers, representing the inter-window strides.

  • padding (Union[str, Sequence[tuple[int, int]]]) – either the string ‘SAME’, the string ‘VALID’, or a sequence of n (low, high) integer pairs that give the padding to apply before and after each spatial dimension.

  • lhs_dilation (Optional[Sequence[int]]) – None, or a sequence of n integers, giving the dilation factor to apply in each spatial dimension of lhs. LHS dilation is also known as transposed convolution.

  • rhs_dilation (Optional[Sequence[int]]) – 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.

  • precision (Union[None, str, Precision, tuple[str, str], tuple[Precision, Precision]]) – 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 (Union[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:

Array

Returns:

An array containing the convolution result.