Patch Embedding with CLS Token

Turn a batch of images into an ordered sequence of vectors: one class token followed by one embedding for each square image patch. This is an input stage for a vision transformer, before any attention layers run.

Inputs and output

Implement patch_embed_with_cls(x, w_proj, cls_token, pos_embed). All inputs are floating-point tensors with the same dtype and device.

• x: (N, C, H, W), using channels-first image layout.
• w_proj: (C * P * P, d_model), the projection shared by all patches.
• cls_token: (d_model,), the shared class-token vector.
• pos_embed: (num_patches + 1, d_model), one position vector per sequence slot.
• Return (N, num_patches + 1, d_model), preserving the input dtype.

P is a positive square patch size, not a separate argument. Infer it from P * P = w_proj.shape[0] // C. Inputs are consistent: P divides both H and W, and num_patches = (H // P) * (W // P). Images need not be square.

Build the sequence

1. Split each image into non-overlapping P by P patches. Visit patches left to right across each patch row, then top to bottom.
2. Flatten each patch in channel-first order: all pixels of channel 0, then channel 1, and so on. Within a channel, visit rows top to bottom and pixels left to right. Each flattened patch has C * P * P values.
3. Multiply each flattened patch by w_proj. There is no projection bias.
4. Prepend cls_token to each image's sequence. Do not project the class token.
5. Add pos_embed to every sequence, including slot 0 for the class token.

Use reshaping, axis reordering and matrix multiplication, not a convolution layer. Leave all inputs unchanged. Do not add padding, normalization, an activation or attention.

For example, a one-channel image with rows [1, 2, 3, 4] and [5, 6, 7, 8] has two size-2 patches. Their flattened vectors, in order, are [1, 2, 5, 6] and [3, 4, 7, 8]. After projection, these occupy slots 1 and 2. Slot 0 is cls_token + pos_embed[0].

What does the class token contain?

At this stage, the class-token output is identical across the batch and does not depend on the image. Later attention layers can mix patch information into that position for a classification head to use. Those layers are outside this exercise. Position embeddings give each sequence slot its own supplied vector; they are shared across images rather than generated separately for each image.

The diagram below illustrates non-overlapping windows along one image axis with size and stride both 4. It does not show channel flattening, projection or the class token. Keep padding at 0 and dilation at 1 for this exercise.