Patchify an Image into Fixed-Size Patches

The problem, from first principles

A transformer operates on a sequence of tokens — that's the entire premise of everything in 05-transformers-llm's attention and transformer-block questions. An image isn't naturally a sequence at all; it's a 2D (or 3D, with channels) grid. The Vision Transformer's founding idea is disarmingly simple: chop the image into a grid of fixed-size square patches, treat each patch as if it were one "token," and feed the resulting sequence of patches through an ordinary transformer, completely unmodified from how it processes word tokens.

From theory to code

Theory says: split the image into a regular grid of non-overlapping patch_size x patch_size squares (covering every channel), then flatten each individual patch (all its channels and pixels) down into one single vector — turning a (C, H, W) image into a (num_patches, C*patch_size*patch_size) sequence of flattened-patch vectors.

Implement patchify(image, patch_size) against that reasoning.

Constraints

• image: shape (C, H, W), with both H and W divisible by patch_size.
• patch_size: side length of each square patch.
• Returns shape (num_patches, C * patch_size * patch_size), where num_patches = (H // patch_size) * (W // patch_size).
• Patches are ordered row-major over the patch grid: left-to-right within a row of patches, then top-to-bottom across rows.
• image is never modified in place.

Hints

Open one at a time. Each gives away a little more than the last.