This is a very smart solution!
However, it assumes that my “M” is sorted by column. When “M” is
M = torch.Tensor([[1, 7, 4], [4, 5, 6], [7, 2, 6]])
and I = torch.LongTensor([0, 2, 1])
mask = M.index_select(0, I).diag().expand(3, 3) >= M
the results will be
mask
1  0  1
0  0  1
0  1  1