Patchify an Image into Flattened Vectors

mediumPROPatch Embedding & Visual TokensVision-Language ModelsVLM 1: Vision-Language Models

Problem Statement

Turn an image tensor into the flat matrix a patch-embedding Linear expects: one row per patch, each row the patch's pixels flattened in (channel, row, col) order.

Background

A ViT patch embedding is a linear projection applied to each non-overlapping patch. Concretely, an image of shape (C, H, W) split into p x p patches produces (H//p) * (W//p) patches, each flattened to a vector of length C * p * p. Patches are visited in row-major order (left to right, then top to bottom), and within a patch the flatten order is channel-major: all of channel 0's p*p values, then channel 1's, and so on.

Your Task

Implement:

def patchify(image, p):

image: nested list of shape (C, H, W).
Return a nested list of shape (num_patches, Cpp).

Input Format

image: (C, H, W) nested list of numbers; H and W are divisible by p.
p (int): patch edge.

Output Format

A nested list, num_patches rows in row-major patch order.

Sample

image = [[[1, 2], [3, 4]]]   # C=1, 2x2
print(patchify(image, 1))

Output:

[[1], [2], [3], [4]]

Example

Input

image = [[[1, 2], [3, 4]]]
print(patchify(image, 1))

Output

[[1], [2], [3], [4]]

Reasoning

The input image has shape $(C, H, W) = (1, 2, 2)$ with values $\begin{bmatrix} 1 & 2 \\ 3 & 4 \end{bmatrix}$ , and the patch size is $p=1$ . This divides the image into $\frac{H}{p} \times \frac{W}{p} = 2 \times 2 = 4$ non-overlapping patches, each of size $1 \times 1$ .
Patches are extracted in row-major order (top-to-bottom, left-to-right). The first patch is at position $(0,0)$ , containing the single pixel value $1$ .
Since $p=1$ and $C=1$ , each patch flattens to a vector of length $C \cdot p \cdot p = 1 \cdot 1 \cdot 1 = 1$ . The first patch yields the vector $[1]$ .
The second patch is at position $(0,1)$ , containing the pixel value $2$ , which flattens to the vector $[2]$ .
The third patch is at position $(1,0)$ , containing the pixel value $3$ , which flattens to the vector $[3]$ .
The fourth and final patch is at position $(1,1)$ , containing the pixel value $4$ , which flattens to the vector $[4]$ .
The final output is [[1], [2], [3], [4]]

Constraints

1 <= C <= 8, H, W divisible by p, up to 64.
Patch order is row-major; within a patch, channel-major then row then col.
Return ints/floats matching the input values.

Learn this in

Multimodal LLMsLLM Study Plan Vision EncodersVLM Study Plan

More from VLM 1: Vision-Language Models

Visual Token Pruning by Attention Masshard CHAIR-i and CHAIR-s Hallucination Ratesmedium Merge a LoRA Update into the Base Weightmedium QLoRA Memory Footprint Estimatemedium POPE Precision, Recall, and F1medium

Concept primerPREVIEW

1. Background Knowledge

In Vision Transformers (ViT), an image is not processed as a whole grid but as a sequence of visual tokens. The first step is patch embedding: the image tensor of shape $(C, H, W)$ is divided into non-overlapping square patches of size $p \times p$ . Each patch is then flattened into a 1D vector and projected via a linear layer. The number of patches (tokens) is $\frac{H}{p} \times \frac{W}{p}$ , and each token has dimension $C \cdot p \cdot p$ .

The flattening order is critical. Within a single patch, the order is channel-major: all $p \times p$ spatial values of channel 0 come first, followed by channel 1, and so on. This mirrors how a Linear layer expects its input: a contiguous vector where the first $p^2$ elements belong to channel 0, the next $p^2$ to channel 1, etc. The patch traversal order across the image is row-major: left-to-right along the width, then top-to-bottom along the height.

Think of the image as a 3D block. You are slicing it into $p \times p$ cubes along the spatial dimensions, then unrolling each cube into a vector. The result is a 2D matrix where each row is one patch's flattened content, ready for a linear projection.

2. Algorithm Approach

This is a nested-loop extraction problem. You iterate over patch positions in row-major order, and for each patch, you extract and flatten its pixels in channel-major order.

The key insight is separating two levels of iteration:

Outer loop: iterate over patch row indices $i$ and patch column indices $j$ .
Inner loop: for each patch, iterate over channels $c$ , then local row $r$ , then local column $col$ , collecting values in that order.

You can also think of it as: for each patch position, gather a sub-tensor of shape $(C, p, p)$ and flatten it. The flatten order is determined by the loop nesting: channel outermost, then row, then column.

3. Step-by-Step Strategy

PRO

Continue the full explanation

You're reading the free preview. Unlock the complete walkthrough, the code editor, test runner and reference solution with Premium.