Skip to content

← JAX Numerical Computing step 13 of 25

Medium Primitives

1-D Conv via lax.conv_general_dilated

Jump to the task

Slide a weighted window along a sequence

At each position, multiply a window of input values by the kernel and sum the products. Move the window one place and repeat. This is the operation commonly called convolution in neural networks.

lax.conv_general_dilated uses the kernel in its given order, without flipping it. Mathematically, this is cross-correlation. A signal-processing convolution flips the kernel, so an asymmetric kernel can produce a different result.

Worked example

import jax.numpy as jnp
from jax import lax

x = jnp.array([1.0, 2.0, 4.0, 8.0])
kernel = jnp.array([1.0, -1.0])
out = lax.conv_general_dilated(
    x.reshape(1, 1, -1), kernel.reshape(1, 1, -1),
    window_strides=(1,), padding='VALID',
    dimension_numbers=('NCH', 'OIH', 'NCH')
)
result = out.reshape(-1)
# [-1.0, -2.0, -4.0]

The first output is 1*1 + 2*(-1) = -1. The next window gives 2*1 + 4*(-1) = -2. There are three positions where the two-entry kernel fits fully inside the four-entry input. No zeros are added at the edges.

Problem

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

  • Inputs are one-dimensional real float32 arrays of lengths L and K, with 1 <= K <= L.
  • Use stride (1,), padding 'VALID', and no dilation.
  • Use one batch, one input channel, and one output channel.
  • Return a one-dimensional array of length L - K + 1.
  • Keep the kernel in its given order; do not reverse it or add a bias.

A length-one kernel preserves the input length. A kernel as long as the input produces a single output. VALID does not always mean a shorter result.

Give each axis a meaning

This API uses a batch axis, channel axes, and spatial axes. For one spatial axis, both operands need three dimensions:

  • Input (1, 1, L) uses NCH: batch, input channels, length.
  • Kernel (1, 1, K) uses OIH: output channels, input channels, length.
  • Output uses NCH and has shape (1, 1, L - K + 1).

The letter H is just the shared label for the spatial axis here. The three layout strings describe input, kernel and output in that order. Pass them with the dimension_numbers= keyword. The argument is optional in the API, but spelling it out makes the chosen layout clear; passing it as the fifth positional argument would instead set input dilation.

Other layouts are supported when their strings match the actual arrays. For this task, add the singleton axes, apply the operation, then remove the batch and channel axes from the result.

Explore the window

The diagram below shows a separate seven-entry input and three-entry kernel. It illustrates which values a window covers; the code example above uses different lengths and weights.

Loading visualization…