MLP-Mixer Token Mixing

Problem

Apply the token-mixing MLP of an MLP-Mixer: mix information across tokens with W @ tokens, where W is (num_tokens, num_tokens).

The signature is mlp_mixer_token_mixing(tokens, W). tokens is (num_tokens, channels) and the output has the same shape.

Hint

Mixer alternates token mixing (across positions) and channel mixing (within positions).

case 1

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

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

case 2

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

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

case 3

in: [[[1,0],[0,1]],[[1,1],[1,0]]]

expected: [[1,1],[1,0]]

case 4

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

expected: [[2,4,6]]

Starter codePython

def mlp_mixer_token_mixing(tokens, W):
    # Your code here
    pass

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