Logit Lens Argmax

Problem

Apply the logit lens: for every layer's logits, return the argmax token id, breaking ties by the smaller index.

The signature is logit_lens_argmax(layer_logits). layer_logits is a (num_layers, vocab_size) matrix and the result has one id per layer.

Hint

The logit lens reveals what the residual stream would decode to at each layer.

case 1

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

expected: [2,0]

case 2

in: [[[0.5,0.5],[0.5,0.5]]]

expected: [0,0]

case 3

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

expected: [0,1,2]

case 4

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

expected: [0,1]

Starter codePython

def logit_lens_argmax(layer_logits):
    # Your code here
    pass

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