📘

Patchify an Image into Flattened Vectors

MediumVision-Language Models 📚 VLM 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.

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

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

👑

Continue the full explanation

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

🔒

Editor locked

The code editor is locked for Pro problems. It is only available for free problems. Please upgrade to gain access to the code editor for all problems.

solution.py

Python 3.13

Test Results

0/0

Run code to see test results.