NNX Implement Conv1D

Apply one learned filter along a sequence

A 1-D convolution applies the same kernel at each sequence position. Each output feature combines nearby input positions and all input channels, then adds a learned bias. Sharing the kernel lets the layer recognize a local pattern wherever it appears.

This exercise uses the neural-network convention of cross-correlation: the kernel is used in its stored order, without reversing its width axis. Use stride 1, no dilation, and zero padding chosen by padding="SAME".

Shapes and axis labels

The input is one sequence of shape (L, C_in). Add a batch dimension to pass (1, L, C_in) to jax.lax.conv_general_dilated. The kernel has shape (W, C_in, C_out), and the bias has shape (C_out,).

Pass dimension_numbers=("NWC", "WIO", "NWC") to describe these layouts:

• NWC means batch, sequence position, channel for input and output.
• WIO means kernel position, input channel, output channel for weights.

These labels tell JAX how to interpret the arrays. They are not a reason to transpose the kernel arbitrarily: exchanging input and output axes can cause a shape error, or silently change the operation when the sizes happen to match.

SAME padding, including even kernels

At stride 1, SAME padding preserves sequence length. The total padding is W - 1 zeros: (W - 1) // 2 on the left, with the rest on the right. For width 3, add one zero on each side. For width 2, add none on the left and one on the right. The output shape is always (L, C_out) here.

The kernel weights multiply the padded input in order. For input [1, 2, 3, 4] and width-2 weights [10, 1], the outputs before bias are [12, 23, 34, 40]. Reversing the kernel would give different values.

A complete example

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

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

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

model = MyConv1D(1, 1, 2, nnx.Rngs(0))
model.kernel[...] = jnp.array([10., 1.], dtype=jnp.float32).reshape(2, 1, 1)
model.bias[...] = jnp.array([2.], dtype=jnp.float32)
x = jnp.array([1., 2., 3., 4.], dtype=jnp.float32).reshape(4, 1)
print(model(x).reshape(-1))  # [14., 25., 36., 42.]

The known weights are only for the example. For the exercise, keep the random kernel generated by the constructor. Call the LeCun-normal initializer with one key from rngs.params() and the full kernel shape. It uses a scaled truncated-normal distribution with fan-in W * C_in; plain standard-normal sampling is not the same recipe. The zero bias consumes no random key. Both arrays are nnx.Param objects.

Explore the window

The visualization starts with width 3 and one padded position per side, matching stride-1 SAME padding. Its padding control is symmetric. It cannot represent the asymmetric padding needed for even kernel widths; use the width-2 example above for that case. Changing its controls may also move away from the exercise's fixed stride and padding rule.