jax.lax.conv#
- jax.lax.conv(lhs, rhs, window_strides, padding, 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 (
str
) – either the string ‘SAME’, the string ‘VALID’.precision (
Union
[None
,str
,Precision
,tuple
[str
,str
],tuple
[Precision
,Precision
]]) – Optional. EitherNone
, which means the default precision for the backend, aPrecision
enum value (Precision.DEFAULT
,Precision.HIGH
orPrecision.HIGHEST
) or a tuple of twoPrecision
enums indicating precision oflhs`
andrhs
.preferred_element_type (
Union
[str
,type
[Any
],dtype
,SupportsDType
,None
]) – Optional. EitherNone
, 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:
- Returns:
An array containing the convolution result.