jax.lax.conv

jax.lax.conv(lhs, rhs, window_strides, padding, precision=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 (Optional[Any]) – Optional. Either None, which means the default precision for the backend, or a lax.Precision enum value (Precision.DEFAULT, Precision.HIGH or Precision.HIGHEST).

Return type

Any

Returns

An array containing the convolution result.