jax.lax.conv
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 (
Any
) β a rank n+2 dimensional input array.rhs (
Any
) β 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 (
Optional
[Any
]) β 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.