Let’s say I have a=to.randn(4, 4).
What I want to do is zero the rows starting an index that change for each row. Example manual implementation would be
a[:2, 1:] = 0.
a[2:, 2:] = 0.

Is there a way to do it without having to do a loop? What would be given for example is a list of the lengths to keep for each row for example, in this case [1, 1, 2, 2], then the question is how to turn that into a mask.