ViT Patch Embedding

Background

Vision Transformer (ViT) treats an image as a sequence of fixed-size patches. Each patch is flattened, linearly projected to the embedding dimension, and fed into a standard transformer encoder. The patching + projection step is the only thing different from a text transformer's tokeniser.

Problem statement

Implement patch_embed(image, patch_size, projection):

1. Verify $H$ and $W$ are divisible by patch_size.
2. Reshape $(H, W, C)$ → $(n_H, n_W, \text{ps}, \text{ps}, C)$ → flatten each patch to $(\text{ps}^2 \cdot C)$.
3. Project: flat @ projection → $(n_H \cdot n_W, \text{embed\_dim})$.

Input

• image — np.ndarray shape $(H, W, C)$.
• patch_size — int >= 1, must divide both $H$ and $W$.
• projection — np.ndarray shape $(\text{ps}^2 \cdot C, \text{embed\_dim})$, the linear projection matrix.

Output

Returns np.ndarray shape $(n_{\text{patches}}, \text{embed\_dim})$ where $n_{\text{patches}} = (H/\text{ps}) \cdot (W/\text{ps})$.

Examples

Example 1 — 8x8 image, 2x2 patches, 3 channels, embed 16

Input:  image=(8, 8, 3); patch_size=2; projection=(12, 16)
Output: (16, 16)

Example 2 — single-patch image of ones, all-ones projection

Input:  image=ones((4,4,1)); patch_size=4; projection=ones((16, 8))
Output: shape (1, 8), all values = 16  (sum of 16 ones)

Example 3 — non-divisible patch size

Input:  image=(5, 5, 3); patch_size=4
Output: ValueError

Constraints

• $H$ and $W$ must be divisible by patch_size.
• Patches are emitted in row-major order: top-left to bottom-right.
• Returns np.ndarray shape $(n_{\text{patches}}, \text{embed\_dim})$.

Notes

• CLS token + positional embeddings. Production ViTs prepend a learned [CLS] row and add learned positional embeddings before the transformer. Out of scope for this patching kernel.
• Why patches not pixels. Pixel-level attention is $O(H^2 W^2)$ — astronomical at 224x224. Patching cuts the sequence by patch_size^2.