Skip to content

← Flax step 38 of 100

Medium Primitives

ViT Patch Embedding

Jump to the task

Turn image patches into tokens

Patch embedding divides an image into non-overlapping square patches and applies the same learned affine projection to each patch. A patch containing P * P * C pixel values becomes a vector of width D. These vectors can form the input sequence for a vision transformer.

This layer only produces patch embeddings. It does not add a class token, positional embeddings, attention, or a classification head.

A convolution is a shared patch projection

For a channels-last image (H, W, C), use a convolution with kernel size (P, P), stride (P, P), and VALID padding. Each window starts exactly one patch-width after the previous window, so patches neither overlap nor leave gaps when both spatial dimensions are divisible by P.

The kernel has shape (P, P, C, D) and the bias has shape (D,). Flattening the kernel's first three axes gives the matrix of shape (P * P * C, D) that projects each flattened patch. Include the bias: the default Flax convolution is an affine map, not just a matrix product.

The convolution output is (H // P, W // P, D). Reshape its spatial axes into a token axis, giving (num_patches, D) where num_patches = (H // P) * (W // P). This reshape orders patches from left to right across the top row, then across each subsequent row. Within a patch, the projection uses row, column, then channel order.

Check patch order with known values

For this 4-by-4 single-channel image and P=2, the four patches begin at the top-left, top-right, bottom-left, and bottom-right corners. A kernel of ones with one output channel sums each patch:

import jax.numpy as jnp
import flax.linen as nn
import numpy as np

image = jnp.arange(1., 17.).reshape(4, 4, 1)
conv = nn.Conv(1, kernel_size=(2, 2), strides=(2, 2), padding="VALID")
variables = {"params": {
    "kernel": jnp.ones((2, 2, 1, 1)), "bias": jnp.zeros(1)
}}
tokens = conv.apply(variables, image).reshape(4, 1)
np.testing.assert_array_equal(tokens[:, 0], [14., 22., 46., 54.])
# First patch: 1 + 2 + 5 + 6 = 14.

These fixed weights make the ordering easy to inspect. The graded function must use the random initialization specified below.

Shape and boundary checks

The task requires positive image dimensions divisible by P. VALID padding adds no zeros. Outside that contract, complete windows are still used and leftover bottom/right pixels may be omitted; this function is not asked to resize or pad an image.

For this exact kernel-equals-stride configuration on divisible images, SAME padding also needs no zeros and gives the same windows. Use VALID as specified. Smaller strides would overlap patches; larger strides would skip positions. Those are different operations from this task.

The diagram shows one spatial axis of length eight split into two width-four patches. A two-dimensional image applies this window layout along both height and width; each grid cell produces one D-dimensional token. The diagram does not show learned projection weights.

Loading visualization…

Problem

Implement patch_embed_forward(seed, image, patch_size, d_model):

  1. Convert seed, patch_size, and d_model to Python integers.
  2. Create a standalone nn.Conv with features=d_model, kernel_size=(patch_size, patch_size), strides=(patch_size, patch_size), and padding="VALID".
  3. Keep its default initializer and bias. Initialize it directly on the unbatched image with jax.random.PRNGKey(seed), without splitting the key or wrapping the layer in another module.
  4. Apply it with the initialized variables, reshape to (num_patches, d_model), and flatten in row-major order for return.

Inputs: image is a float array (H, W, C) with positive dimensions. patch_size is positive and divides both H and W. d_model is positive. These widths and seed are integer-valued scalars that may arrive as floats. Output: a one-dimensional array of length (H // patch_size) * (W // patch_size) * d_model.

The parameter count is P * P * C * D + D, independent of the number of patches. Identical patches receive identical embeddings because the projection is shared and no positional encoding is added here. Keep all scalar setup arguments static when JIT-compiling the wrapper.