Attention Head Dimension

Problem

Split the transformer embedding dimension across attention heads.

Return the per-head dimension dim // num_heads. You may assume dim is divisible by num_heads.

Hint

Multi-head attention keeps the total projection width equal to dim.

case 1

in: [768,12]

expected: 64

case 2

in: [384,6]

expected: 64

case 3

in: [512,8]

expected: 64

case 4

in: [768,16]

expected: 48

Starter codePython

def head_dim(dim, num_heads):
    # Your code here
    pass

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