Tensor assigment when doubly indexed

1

I’m having trouble understanding why this piece of code does not work as expected:

a = torch.rand(1000)
idx1 = torch.arange(500, 600)
idx2 = torch.rand(100) < 0.5

a[idx1][idx2] = 1.0
print(a[idx1][idx2])

I expect the result to be all ones, but the tensor is unchanged. If instead, I index only once like so:

a[idx1] = 1.0
print(a[idx1])

Then, as expected I get all ones. Why does the double index not work and how can I achieve this behavior?

2

What you are loking for is:

a[idx1[idx2]] = 1.0

(Mask idx1 indexes tensor using boolean idx2 tensor before slicing tensor a)

3

Maybe my example wasn’t broad enough. Your solution doesn’t work if the idx shapes don’t match. Say idx1 is a boolean array of shape 1000, and idx2 is of shape idx1.sum(). What would you do then?

Something like torch.where(idx1)[0][idx2] could work, but that seems very inefficient.