ViT Class Token

Problem

Prepend the learnable ViT class token to a sequence of patch embeddings, returning a new sequence of length len(x) + 1.

The signature is vit_class_token(x, cls_token). x has shape (num_patches, d_model) and cls_token is a length d_model vector.

Hint

The class token is prepended like a BERT-style special token; its output is used for classification.

case 1

in: [[[1,2],[3,4]],[0,0]]

expected: [[0,0],[1,2],[3,4]]

case 2

in: [[[1]],[9]]

expected: [[9],[1]]

case 3

in: [[[0.5,-0.5],[1,2],[3,4]],[1,-1]]

expected: [[1,-1],[0.5,-0.5],[1,2],[3,4]]

case 4

in: [[[2,3]],[0,0]]

expected: [[0,0],[2,3]]

Starter codePython

def vit_class_token(x, cls_token):
    # Your code here
    pass

4 test cases · run your code against them in the editor