Top-2 Gating Indices

Problem

Return the indices of the two largest routing logits, breaking ties by the smaller index, sorted ascending.

The signature is top2_gating(logits). Return the two indices.

Hint

Top-2 MoE routing activates two experts per token to combine specialization and smoothness.

case 1

in: [[1,3,2]]

expected: [1,2]

case 2

in: [[5,5,5]]

expected: [0,1]

case 3

in: [[0,-1,2]]

expected: [0,2]

case 4

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

expected: [1,3]

Starter codePython

def top2_gating(logits):
    # Your code here
    pass

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