core/nn/vision/vit_backbone library

Vision Transformer (ViT) backbone.

Takes a batch of pre-patchified images (each patch flattened to a vector) and produces per-token contextualized features via the standard ViT recipe:

  1. Patch projection — Linear(patchSize * patchSize * numChannels, embedDim) maps each patch pixel vector to an embedding.
  2. CLS token — a single learnable [1, embedDim] vector is prepended to the sequence. Downstream heads read the CLS row of the encoded output as an image-level summary.
  3. Positional embeddings — a learnable [numPatches + 1, embedDim] table is added to the (CLS + patches) sequence.
  4. Transformer encoder — the standard pre-LN block stack from TransformerEncoder. No causal mask (vision is bidirectional).

This is the "learn from images" analog of a language model: swap the token embedding for a patch projection, keep the transformer stack, tack a task-specific head on top of the CLS output.

Currently 2D (single image per call): input shape [numPatches, patchSize * patchSize * numChannels], output shape [numPatches + 1, embedDim]. Batched image support can be added by switching the concat to axis 1 with [B, ...] inputs, once a batched concat kernel lands.

Classes

ViTBackbone

Functions

vitClsFeature(Tensor encoded) → Tensor
Take the CLS row (row 0) of a ViT encoder output. Convenience for heads that want a [1, embedDim] image-level feature vector.