# How to zero a matrix differently on each row

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.

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.

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?

``````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
``````