Skip to content

← Flax step 13 of 100

Medium Primitives

Implement Conv2D with Stride and Padding

Jump to the task

Move a filter across an image

Conv2D applies the same learned filter at different image positions. At each position, it sums products over the filter's height, width, and input channels to produce one value per output channel. As in Conv1D, the operation here is cross-correlation: do not reverse the kernel.

You will create the parameters yourself and call jax.lax.conv_general_dilated. This makes the array axes, stride, and padding explicit. The stride controls how far the window moves between output positions; padding lets windows include zeros beyond the image.

Array layouts

This task uses channels-last arrays:

Array Shape
Input image (H, W, C_in)
Input with batch axis (1, H, W, C_in)
Kernel (kernel_h, kernel_w, C_in, features)
Bias (features,)
Output after removing batch axis (H_out, W_out, features)

The layout strings ("NHWC", "HWIO", "NHWC") describe the input, kernel, and output in that order. N means batch, H and W mean spatial dimensions, and I and O distinguish the kernel's input and output channels. JAX can use other layouts; these are the conventions selected for this exercise, not a universal hardware performance rule.

Calculate SAME padding

Use the same positive stride s for height and width, with no dilation. For either spatial dimension, let n be the input size and k the kernel size:

  1. out = (n + s - 1) // s, which is ceil(n / s).
  2. total_pad = max((out - 1) * s + k - n, 0).
  3. pad_before = total_pad // 2 and pad_after = total_pad - pad_before.

Compute this separately for height and width. When padding is odd, the extra zero goes on the bottom or right. SAME preserves spatial size at stride one; at larger strides, the output is smaller. It does not always mean adding k // 2 zeros to both sides.

A worked strided example

Take an image with one channel and rows [1, 2, 3, 4], [5, 6, 7, 8], and [9, 10, 11, 12]. Use kernel [[1, 2, 3], [4, 5, 6]] and stride two. The output has shape (ceil(3/2), ceil(4/2), 1), or (2, 2, 1). Padding adds one zero row below and one zero column to the right.

The first output is 1*1 + 2*2 + 3*3 + 5*4 + 6*5 + 7*6 = 106. The full output, before bias, is [[106, 79], [62, 35]]:

import jax
import jax.numpy as jnp
import numpy as np

x = jnp.arange(1., 13.).reshape(1, 3, 4, 1)
kernel = jnp.arange(1., 7.).reshape(2, 3, 1, 1)
y = jax.lax.conv_general_dilated(
    x, kernel, window_strides=(2, 2), padding="SAME",
    dimension_numbers=("NHWC", "HWIO", "NHWC"),
)
np.testing.assert_allclose(y[0, :, :, 0], [[106., 79.], [62., 35.]])

The interactive diagram below shows a separate one-axis example: input length seven, width-three kernel, stride two, and one zero on each side give four output positions. Apply the same reasoning to both image axes. Its symmetric padding control cannot show the asymmetric padding in the worked example above.

Loading visualization…

Problem

Implement my_conv2d_forward(seed, x, features, kernel_h, kernel_w, stride):

  1. Define a compact Flax module MyConv2D with the four integer fields features, kernel_h, kernel_w, and stride.
  2. Create "kernel" with shape (kernel_h, kernel_w, x.shape[-1], features) using nn.initializers.lecun_normal().
  3. Create "bias" with shape (features,) using nn.initializers.zeros.
  4. Add a batch axis to the input and call the convolution primitive with window_strides=(stride, stride), padding="SAME", and dimension_numbers=("NHWC", "HWIO", "NHWC"). Use no dilation or channel grouping.
  5. Remove only the batch axis with y[0] and add bias across all spatial positions.
  6. Initialize this module once on x using jax.random.PRNGKey(int(seed)), apply it with those variables, and flatten the result in row-major order.

Instantiate the custom module directly. Keep the parameter names and creation order above; splitting the key or adding a parent module changes seeded initialization.

Inputs: x is a float array (H, W, C_in) with positive dimensions. The four module arguments are positive integer-valued scalars, possibly supplied as floats; convert them to Python integers. seed is an integer-valued scalar. Output: a one-dimensional array of length ceil(H / stride) * ceil(W / stride) * features.

A rectangular image or kernel helps catch swapped height/width axes. With multiple channels, also check the kernel's input/output channel ordering. Avoid unrestricted squeeze(), which can remove a valid length-one spatial or channel axis. Keep all scalar setup arguments static if JIT-compiling the wrapper.