Positional Embedding Count

Problem

Count the trainable positional embeddings a ViT needs.

Each of the num_patches patch tokens gets one positional vector, and the CLS token gets one more when include_cls is True. Return (num_patches + include_cls) * dim.

Test cases

case 1

in: [196,768,true]

expected: 151296

case 2

in: [49,384,false]

expected: 18816

case 3

in: [0,768,true]

expected: 768

case 4

in: [256,192,false]

expected: 49152

Starter codePython

def pos_embed_count(num_patches, dim, include_cls):
    # Your code here
    pass

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