Skip to content

← JAX Numerical Computing step 15 of 25

Medium Primitives

Conv Padding & Strides

Jump to the task

Decide what happens at the edges

The padding argument to lax.conv_general_dilated controls how the output length relates to the input length. This exercise uses stride 1 and no dilation, with input length L and kernel length K:

  • 'VALID': no padding is added. When the kernel fits, output length is L - K + 1. A length-1 kernel preserves the length; larger kernels shorten it.
  • 'SAME': zeros are added so output length equals input length at stride 1. If the total required padding is odd, the extra zero goes at the end. For example, a length-2 kernel needs 0 zeros before and 1 after.

VALID uses only windows that fit entirely inside the input. SAME includes windows that extend beyond an edge, treating the missing values as zeros. Neither makes edge behavior disappear; they define different calculations.

Worked example

import jax.numpy as jnp
from jax import lax

x = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
kernel = jnp.array([1.0, 1.0, 1.0])

x_3d = x.reshape(1, 1, -1)
k_3d = kernel.reshape(1, 1, -1)

# VALID: output length 5 - 3 + 1 = 3
out_valid = lax.conv_general_dilated(
    x_3d, k_3d, (1,), 'VALID', dimension_numbers=('NCH', 'OIH', 'NCH')
).reshape(-1)
# → [6.0, 9.0, 12.0]

# SAME: output length = input length = 5
out_same = lax.conv_general_dilated(
    x_3d, k_3d, (1,), 'SAME', dimension_numbers=('NCH', 'OIH', 'NCH')
).reshape(-1)
# → [3.0, 6.0, 9.0, 12.0, 9.0]

At the left edge, SAME's first window is [0, 1, 2], giving 3. At the right edge it is [4, 5, 0], giving 9. The middle three windows match VALID. As in the preceding lessons, the kernel is not reversed.

How stride and dilation change the size

  • The simple formula assumes stride 1 and no dilation. In general, with no input dilation, the effective kernel length is D * (K - 1) + 1. For VALID, output length is max(0, floor((L - effective_kernel) / S) + 1).
  • A kernel need not shrink the input. At stride 1, a length-1 VALID kernel preserves its length. An oversized kernel has no valid positions; this JAX operation can return an empty output rather than raising an error.
  • SAME does not always mean length L. With stride S > 1 and no input dilation, output length is ceil(L / S). The padding split depends on the input, effective kernel and stride, not just whether K is even.
  • Explicit zero padding is valid. For this 1-D operation, padding=[(0, 0)] is equivalent to 'VALID'. Pass a sequence of (before, after) pairs, not the bare integer 0.

Problem

Implement conv_with_padding(x, kernel, padding_mode) that applies 1-D convolution using lax.conv_general_dilated with the padding selected by padding_mode.

  • x and kernel are nonempty one-dimensional real float32 arrays.
  • padding_mode is a scalar: 0.0 means 'VALID', 1.0 means 'SAME'. These are the only values used in this exercise.
  • Use stride (1,), no dilation and no bias. Keep the kernel in its original order.
  • Add singleton batch/channel axes and use dimension_numbers=('NCH', 'OIH', 'NCH').
  • Return a one-dimensional array. At stride one, its length is max(0, L - K + 1) for VALID and L for SAME.

The name introduces strides, but this function fixes the stride at one. The general formulas above help explain what would change with other strides.

Choosing the mode inside JIT

In an eager call, a concrete scalar can select the Python padding string. Converting a traced value with int(padding_mode) does not make it static. To JIT this function, pass the mode as a Python scalar and mark it static: jax.jit(conv_with_padding, static_argnames=('padding_mode',)). Each mode can then have its own compiled output shape.

A dynamic lax.cond is not a general replacement here: its branches must have matching shapes, while VALID and SAME usually return different lengths. The exercise runner may supply a scalar JAX value for eager execution; use a Python value when passing a static argument to your own JIT wrapper.

Explore the window

The diagram starts with the worked example: input 5, kernel 3, stride 1, no padding. Set padding to 1 to match its 'SAME' case. The padding slider adds the same amount on both sides; it does not automatically compute asymmetric 'SAME' padding for every kernel and stride.

Loading visualization…