How to zero a matrix differently on each row

1

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.

2

I think if you are dealing with indices having a variable shape for each dimension, creating a mask should work. The mask creation itself might also involve a loop, which might be cheaper than the indexing loop.

3

Hm, in the end I want to sum over a certain axis. Is there no way to combine the mask plus summing in such a way so as to avoid doing operations on the just zeroed elements?

4

Surprisingly the mask is slower!

mask = to.zeros((N, T,), dtype=to.bool, device=x.device)
idx = 0
for dur, count in zip(durs, counts):
    mask[idx: idx + count, :dur] = True
    idx += count
joined = (x * mask.unsqueeze(2).expand(-1, -1, x.size(2))).sum(dim=1)

is slower than

joined = to.zeros((N, self.repr_dim), dtype=to.float32, device=x.device)
idx = 0
for dur, count in zip(durs, counts):
    x_part = x[idx: idx + count, :dur]
    joined[idx: idx + count] = x_part.sum(dim=1)
    idx += count