Customization of BlockDiagonalMask or BlockDiagonalCausalMask

❓ Questions and Help

I am trying to write a customization of BlockDiagonalMask with an additional variable group_lengths_list. Inside each block, tokens are separated into 3 groups (number of groups is a hyper-param), all tokens can see the tokens within the same group and those in the last group (the 3rd group here). For example, 2 blocks on the diagonal with group_lengths_list: [1, 1, 2] and [2, 2, 3] would be like

1 0 1 1   0 0 0 0 0 0 0 
0 1 1 1   0 0 0 0 0 0 0 
0 0 1 1   0 0 0 0 0 0 0
0 0 1 1   0 0 0 0 0 0 0 

0 0 0 0   1 1 0 0 1 1 1
0 0 0 0   1 1 0 0 1 1 1
0 0 0 0   0 0 1 1 1 1 1
0 0 0 0   0 0 1 1 1 1 1
0 0 0 0   0 0 0 0 1 1 1
0 0 0 0   0 0 0 0 1 1 1
0 0 0 0   0 0 0 0 1 1 1

I was trying to realize it by making a customization of BlockDiagonalMask or BlockDiagonalCausalMask, like BlockDiagonalCausalLocalAttentionMask, but I cannot pass group_lengths into _create_block_mask (or, identifying current block, e.g., block index).

Is there any workaround?