Skip to content

← Flax step 12 of 100

Medium Primitives

Implement Conv1D from Scratch

Jump to the task

Slide a learned filter along a sequence

A one-dimensional convolution combines nearby positions and input channels to produce output features. The same kernel is reused at every position. Here you will build a Flax layer with self.param and call JAX's convolution primitive directly, rather than using nn.Conv.

The operation used here is cross-correlation: it uses kernel positions in their stored order, without reversing them. This is the convention normally used by neural-network convolution layers. A mathematical convolution reverses the filter, so jnp.convolve is not a direct replacement for this multi-channel operation.

Keep the axes explicit

Array Shape Meaning
Input (L, C_in) Sequence positions, input channels
Batched input (1, L, C_in) One sequence for the primitive
Kernel (K, C_in, C_out) Window positions, input channels, output channels
Bias (C_out,) One offset per output channel
Layer output (L, C_out) Positions, output channels

Pass dimension_numbers=("NWC", "WIO", "NWC") so the primitive interprets the arrays this way. N is batch, W is the spatial axis, C is the input/output array's channel axis, and I/O identify the kernel's input/output channel axes. JAX supports other layouts too; these strings specify the choice for this exercise.

SAME padding at stride one

With stride one and no dilation, SAME padding adds K - 1 zeros in total. Put (K - 1) // 2 on the left and the rest on the right. The output length is then L, including when the kernel is wider than the input. For an even kernel width, the right side gets one more zero.

For example, with input [1, 2, 3, 4] and kernel [1, 2, 3], the padded input is [0, 1, 2, 3, 4, 0]. The first output is 0 * 1 + 1 * 2 + 2 * 3 = 8. The complete result is [8, 14, 20, 11] before any bias:

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

x = jnp.array([1., 2., 3., 4.]).reshape(1, 4, 1)
kernel = jnp.array([1., 2., 3.]).reshape(3, 1, 1)
y = jax.lax.conv_general_dilated(
    x, kernel, window_strides=(1,), padding="SAME",
    dimension_numbers=("NWC", "WIO", "NWC"),
)
np.testing.assert_allclose(y[0, :, 0], [8., 14., 20., 11.])

The asymmetric filter makes its orientation visible. A filter of all ones would give the same result if reversed and would not catch that mistake. With multiple input channels, sum the contribution from every channel as well as every window position.

The diagram starts at this example's width, stride, and padding: four input positions, a width-three kernel, one zero at each end, and four output positions. Its padding control is symmetric, so it does not depict the asymmetric SAME padding needed by even-width kernels.

Loading visualization…

Problem

Implement my_conv1d_forward(seed, x, features, kernel_size):

  1. Define a compact Flax module MyConv1D with integer fields features and kernel_size.
  2. Create a parameter named "kernel" of shape (kernel_size, x.shape[-1], features) with nn.initializers.lecun_normal().
  3. Create a parameter named "bias" of shape (features,) with nn.initializers.zeros.
  4. Add a batch axis to x, then call jax.lax.conv_general_dilated with stride (1,), padding "SAME", and layouts ("NWC", "WIO", "NWC"). Use no dilation or channel grouping.
  5. Remove only the added batch axis and add the bias across positions.
  6. Initialize this module once with jax.random.PRNGKey(int(seed)) and x, apply it with those variables, and flatten the result in row-major order.

Instantiate MyConv1D directly. Keep the parameter names and creation order above so seeded initialization matches. Do not wrap it in another module or split the initialization key.

Inputs: x is a float array (L, C_in) with positive dimensions. features and kernel_size are positive integer-valued scalars; convert them to Python integers. seed is an integer-valued scalar. The scalar inputs may arrive as floats. Output: a one-dimensional array of length L * features.

Use y[0] to remove the batch axis. An unrestricted squeeze() could also remove a length-one spatial or channel axis. Keep seed, features, and kernel_size static if JIT-compiling the wrapper.

Swapping the kernel's input and output axes can produce a shape error, or an incorrect result when the channel counts happen to match. Test numerical values as well as shapes. See the JAX convolution API for the layout and padding arguments.