mlx.core.conv2d#

conv2d(input: array, weight: array, /, stride: int | Tuple[int, int] = 1, padding: int | Tuple[int, int] = 0, dilation: int | Tuple[int, int] = 1, groups: int = 1, *, stream: None | Stream | Device = None) → array#

2D convolution over an input with several channels

Note: Only the default groups=1 is currently supported.

Parameters:

  • input (array) – input array of shape (N, H, W, C_in)

  • weight (array) – weight array of shape (C_out, H, W, C_in)

  • stride (int or tuple(int), optional) – tuple of size 2 with kernel strides. All spatial dimensions get the same stride if only one number is specified. Default: 1.

  • padding (int or tuple(int), optional) – tuple of size 2 with symmetric input padding. All spatial dimensions get the same padding if only one number is specified. Default: 0.

  • dilation (int or tuple(int), optional) – tuple of size 2 with kernel dilation. All spatial dimensions get the same dilation if only one number is specified. Default: 1

  • groups (int, optional) – input feature groups. Default: 1.

Returns:

The convolved array.

Return type:

array