Skip to content

← JAX Numerical Computing step 14 of 25

Hard Primitives

2-D Conv via lax.conv_general_dilated

Jump to the task

Slide a weighted patch across a matrix

The two-dimensional version multiplies a rectangular input patch by a kernel and sums the products. Move the patch across columns, then down rows, to build the output matrix.

As in the 1-D exercise, lax.conv_general_dilated uses the kernel without flipping either spatial axis. This is cross-correlation, conventionally called convolution in neural networks.

Worked example

import jax.numpy as jnp
from jax import lax

x = jnp.array([[1.0, 2.0, 3.0],
               [4.0, 5.0, 6.0]])
kernel = jnp.array([[1.0, 0.0],
                    [0.0, -1.0]])
out = lax.conv_general_dilated(
    x.reshape(1, 1, *x.shape), kernel.reshape(1, 1, *kernel.shape),
    window_strides=(1, 1), padding='VALID',
    dimension_numbers=('NCHW', 'OIHW', 'NCHW')
)
result = out[0, 0]
# [[-4.0, -4.0]], shape (1, 2)

The first patch contributes 1*1 + 2*0 + 4*0 + 5*(-1) = -4. Moving right gives 2 - 6 = -4. The two-row kernel fits at only one row position, so the output keeps a row axis of size one.

Problem

Implement conv2d_basic(x, kernel) with lax.conv_general_dilated:

  • x is a real float32 array of shape (H, W).
  • kernel is a real float32 array of shape (Kh, Kw), with 1 <= Kh <= H and 1 <= Kw <= W.
  • Use stride (1, 1), 'VALID' padding and no dilation.
  • Use one batch, one input channel and one output channel.
  • Return shape (H - Kh + 1, W - Kw + 1), keeping both spatial axes.

Do not reverse the kernel, add bias, or collapse the result to one dimension. A 1-by-1 kernel preserves the input shape; a kernel matching the full input produces shape (1, 1), not a scalar.

Map the four axes

The chosen layout is channels first:

  • Input NCHW: batch, input channels, height, width.
  • Kernel OIHW: output channels, input channels, height, width.
  • Output NCHW: batch, output channels, height, width.

Add the singleton batch and channel axes to create (1, 1, H, W) and (1, 1, Kh, Kw). Supply the three strings through dimension_numbers=. After the operation, select out[0, 0] or reshape using its two spatial sizes. Avoid an unrestricted squeeze(), which could remove a spatial axis of size one.

Four-dimensional operands are required for two spatial dimensions, but NCHW is not the only supported layout. Channels-last arrays can use ('NHWC', 'HWIO', 'NHWC'); the strings must describe the arrays actually passed. A mismatch may cause an error or compute along unintended axes.

Explore one axis at a time

The diagram below is one-dimensional. It illustrates window placement along a single axis, not the full two-dimensional patch calculation. The same stride and padding reasoning applies independently to height and width in this exercise.

Loading visualization…