Skip to content

← Flax step 14 of 100

Medium Primitives

Implement Transposed Convolution

Jump to the task

Learn an upsampling operation

A transposed convolution can turn a short sequence of features into a longer one. Unlike repeating or interpolating values with a fixed rule, it uses a learned kernel to combine contributions at output positions. This exercise builds a one-dimensional layer with JAX's lax.conv_transpose and Flax parameters.

For the SAME padding convention used here, an input of length L produces length L * stride. Stride two doubles the length; stride one preserves it. The number of output channels is set separately by features.

What is transposed?

A convolution without bias can be written as a matrix multiplication y = A @ x after flattening the arrays. Its linear transpose maps a vector back with A.T. This is not generally an inverse: A.T @ A need not be the identity matrix. Upsampling therefore does not recover information discarded by a downsampling operation.

This matrix view does not mean that two independently initialized layers share weights. It also does not specify a library's stored kernel orientation or boundary alignment. Those conventions must match before you compare an API call with the transpose of a particular forward convolution.

JAX's kernel convention for this task

Keep transpose_kernel=False, the default. With layouts ("NWC", "WIO", "NWC"), the input has shape (1, L, C_in) and the kernel has shape (kernel_size, C_in, features). The kernel is used without an additional spatial flip or channel swap.

Setting transpose_kernel=True flips spatial axes and swaps the kernel's input/output channel axes. That option is useful when matching other kernel conventions, but it changes the interpretation of the supplied parameter array and is not part of this exercise.

The keyword is strides=(stride,), whereas conv_general_dilated uses window_strides. Do not interchange the keyword names. See the JAX API reference for these options.

Check orientation with an asymmetric kernel

For input [1, 2], kernel [1, 2, 3], stride two, and SAME padding, the default JAX call below returns [3, 2, 7, 4]. The third value combines contributions: 1 * 1 + 2 * 3 = 7.

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

x = jnp.array([1., 2.]).reshape(1, 2, 1)
kernel = jnp.array([1., 2., 3.]).reshape(3, 1, 1)
y = jax.lax.conv_transpose(
    x, kernel, strides=(2,), padding="SAME",
    dimension_numbers=("NWC", "WIO", "NWC"),
    transpose_kernel=False,
)
np.testing.assert_allclose(y.reshape(-1), [3., 2., 7., 4.])
assert y.shape == (1, 4, 1)

For this example, an equivalent calculation inserts a zero between the input values to form [1, 0, 2], pads two zeros on the left and one on the right, and slides the stored kernel over the result at stride one. These padding amounts are specific to this kernel and stride; let conv_transpose determine them for the general task. A kernel of all ones would hide an accidental reversal, so it is a weaker orientation test.

Problem

Implement my_transposed_conv_forward(seed, x, features, kernel_size, stride):

  1. Define a compact Flax module MyTransposedConv with integer fields features, kernel_size, and stride.
  2. Create a parameter named "kernel" with shape (kernel_size, x.shape[-1], features) using nn.initializers.lecun_normal().
  3. Create "bias" with shape (features,) using nn.initializers.zeros.
  4. Add a batch axis to x. Call jax.lax.conv_transpose with strides=(stride,), padding="SAME", and layouts ("NWC", "WIO", "NWC"). Leave transpose_kernel=False and use no kernel dilation.
  5. Remove only the batch axis with y[0] and add bias across positions.
  6. Initialize this module directly on x with jax.random.PRNGKey(int(seed)), apply it with those variables, and flatten the result in row-major order.

Keep the parameter names and creation order above. Do not split the initialization key or wrap this layer in another module.

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

Bias is added after the transposed convolution, including at positions whose kernel contribution is zero. There is no activation or normalization. Keep all scalar setup arguments static if JIT-compiling the wrapper, and avoid an unrestricted squeeze() that could remove spatial or channel axes of length one.