Patch Embedding kernel

[superposition]

Description

Convert images to patch token sequences for Vision Transformers, JEPA, MAE.

Operation

# Input: (B, C, H, W) image
# Output: (B, num_patches, embed_dim)
patches = patch_embed(image, patch_size=16, embed_dim=768)

Implementation

• Fused kernel: reshape + linear projection
• Support variable patch sizes (14, 16, 32)
• Optional: include CLS token prepend

Tests

• Output shape matches expected (B, H*W/P², D)
• Matches nn.Conv2d with kernel_size=stride=patch_size
• Gradient flows correctly
• gradcheck passes with float64
• Works with non-square images

Used In

ViT, JEPA, MAE, BEiT, DINOv2