Patch Embedding Token Count

Problem Statement

Before a Vision Transformer can attend to anything, the image has to become a sequence of tokens. Given an image size, a patch size and an embedding dimension, report the patch grid, the number of visual tokens, the final sequence length, and the parameter count of the patch-projection layer.

Background

A ViT cuts the image into a non-overlapping grid of patch_size x patch_size squares:

grid_h = image_h // patch_size
grid_w = image_w // patch_size
num_patches = grid_h * grid_w

Each patch is flattened to a vector of length patch_size * patch_size * in_channels and pushed through a single shared linear layer into embed_dim. That layer is the patch embedding, and it has

proj_params = patch_size * patch_size * in_channels * embed_dim + embed_dim

parameters (weights plus one bias per output channel). It does not depend on how many patches there are - the projection is applied identically to every patch.

Most ViTs then prepend a learned [CLS] token whose final hidden state is used as the pooled image representation, so the sequence the transformer actually sees is one longer than the patch count:

seq_len = num_patches + 1   (with CLS)

This is the number that matters for a VLM: it is exactly how many positions the image will occupy in the language model's context if the visual tokens are passed through unpooled. CLIP ViT-L/14 at 224x224 gives a 16x16 grid = 256 patches, 257 positions with the CLS token.

Your Task

Implement:

def patch_embedding_stats(image_h, image_w, patch_size, embed_dim, use_cls_token=True, in_channels=3):

Return a list [grid_h, grid_w, num_patches, seq_len, proj_params] of plain Python integers.

Input Format

• image_h, image_w - image height and width in pixels, each divisible by patch_size
• patch_size - side length of a square patch
• embed_dim - transformer hidden width
• use_cls_token - if True, add one [CLS] position to the sequence
• in_channels - channels in the input image (3 for RGB)

Output Format

A list of five integers, as described above.

Sample

print(patch_embedding_stats(224, 224, 14, 1024))

Output:

[16, 16, 256, 257, 603136]