NNX Implement Conv2D

Slide a learned kernel across an image

A 2-D convolution combines a local patch of an image into output features. The same kernel is used at every output position. Each output channel sums contributions across the patch's height, width and input channels, then adds a learned bias.

As in most neural-network convolution layers, this operation is cross-correlation: use the kernel in its stored order, without flipping its spatial axes. This exercise uses zero padding, no dilation, and the same positive stride in both spatial directions.

Input, kernel and output shapes

The input is one image shaped (H, W, C_in). Add a batch axis to make (1, H, W, C_in) for the JAX convolution primitive. Store the kernel as (kernel_h, kernel_w, C_in, C_out) and the bias as (C_out,).

Use dimension_numbers=("NHWC", "HWIO", "NHWC") to describe those arrays. N is batch, H and W are spatial axes, and I and O are the kernel's input and output channel axes. The output channel axis is last. After removing the batch axis, the result is (H_out, W_out, C_out).

These are logical axis choices for this exercise. They are not claims about the fastest memory layout on every accelerator. Confusing the two channel axes can cause a shape error or, if their lengths match, a numerically wrong result with a plausible shape.

How SAME padding works

For an input dimension D, kernel width K along that dimension, and stride s, compute:

• Output size: O = (D + s - 1) // s, which is ceil(D / s).
• Total zero padding: P = max((O - 1) * s + K - D, 0).
• Padding before the input: P // 2; padding after it: P - P // 2.

Apply this rule separately to height and width. An odd total puts the extra zero on the bottom or right. SAME preserves spatial size only at stride 1. For a 3-by-3 image and stride 2, the output is 2-by-2. Some input/kernel/stride combinations need no padding at all.

A complete numerical example

import jax
import jax.numpy as jnp
from flax import nnx

class MyConv2D(nnx.Module):
    def __init__(self, in_features, out_features, kernel_h, kernel_w, stride, rngs):
        key = rngs.params()
        initializer = jax.nn.initializers.lecun_normal()
        self.kernel = nnx.Param(initializer(
            key, (kernel_h, kernel_w, in_features, out_features), jnp.float32
        ))
        self.bias = nnx.Param(jnp.zeros((out_features,), dtype=jnp.float32))
        self.stride = stride

    def __call__(self, x):
        output = jax.lax.conv_general_dilated(
            x[None, ...], self.kernel[...],
            window_strides=(self.stride, self.stride), padding="SAME",
            dimension_numbers=("NHWC", "HWIO", "NHWC")
        )
        return output[0] + self.bias[...]

model = MyConv2D(1, 1, 2, 2, 2, nnx.Rngs(0))
model.kernel[...] = jnp.array([[1., 2.], [3., 4.]], dtype=jnp.float32).reshape(2, 2, 1, 1)
model.bias[...] = jnp.array([1.], dtype=jnp.float32)
x = jnp.arange(1, 10, dtype=jnp.float32).reshape(3, 3, 1)
print(model(x)[..., 0])  # [[38., 22.], [24., 10.]]

There is one padded row at the bottom and one padded column on the right. The top-left output is 1*1 + 2*2 + 4*3 + 5*4 + 1 = 38. The bottom-right window contains only the input value 9 and three padded zeros, giving 9*1 + 1 = 10.

The example sets known weights to expose the arithmetic. In the exercise, keep the constructor's random kernel. LeCun normal uses a scaled truncated-normal distribution with fan-in kernel_h * kernel_w * C_in. Draw one parameter key and initialize the complete HWIO shape in float32; initialize the bias to float32 zeros without another random draw.

Kernel and bias are nnx.Param objects. Stride is a fixed Python integer attribute, not a trainable parameter. A forward pass leaves all three unchanged and consumes no RNG keys.

Explore one spatial direction

The diagram below illustrates a single spatial axis, not the full 2-D kernel. Its width-3, stride-1, padding-1 starting point matches SAME on that axis. The control applies equal padding to both ends; asymmetric SAME padding, as in the worked example, is not represented by that control. Changing the settings does not automatically recompute SAME padding.