Speculative Tree Attention Mask

Problem

Build the attention mask for a speculative decoding tree. parents[i] gives the parent index of node i, or -1 for the root; node i attends to itself and all of its ancestors.

The signature is speculative_tree_mask(parents). Return the 0/1 mask matrix.

Hint

Each draft branch only sees its own ancestors, so the tree can be verified in one forward pass.

Test cases

case 1

in: [[-1]]

expected: [[1]]

case 2

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

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

case 3

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

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

case 4

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

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

Starter codePython

def speculative_tree_mask(parents):
    # Your code here
    pass

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