Fill value to matrix based on index

I have a index matrix and a matrix filled by zero .
For example :

index = Tensor([[ 0,  1],
                [ 1,  2],
                [ 4,  1]])
a = Tensor([[ 0.,  0.,  0.,  0.,  0.,  0.],
            [ 0.,  0.,  0.,  0.,  0.,  0.],
            [ 0.,  0.,  0.,  0.,  0.,  0.]])

Therefore, I would like to fill 1 to the zero matrix based on the index matrix .
Like this :

a = Tensor([[ 1,  1,  0,  0.,  0.,  0.],
            [ 0.,  1,  1,  0.,  0.,  0.],
            [ 0.,  1,  0.,  0.,  1,  0.]])

Is there any easy way ?

You should use index_fill_() (doc link).

for i, ind in enumerate(index):
    a[i].index_fill_(0, ind, 1)
1 Like

Here is another approach without a for loop:

a[torch.arange(a.size(0)).unsqueeze(1), index] = 1.
3 Likes

@ptrblck if index does not have the same number of rows as a, I think a more generic approach is needed. For instance

index = torch.tensor([
    [0, 1],
    [1, 2]])
a = torch.zeros(3, 6)
a[torch.arange(a.size(0)).unsqueeze(1), index] = 1  # IndexError
a[index[:, 0], index[:, 1]] = 1                     # Fixed

Maybe someone knows a more compact way to write the second approach, or one that extends easily to more than 2 dimensions? I have always been puzzled why a[index] = 1 doesn’t work (rather than selecting the multi-dimensional indices, it simply makes all elements of a equal to 1)

Your operation won’t yield the same results.
The original code snippet gives:

index = torch.tensor([
    [0, 1],
    [1, 2],
    [4, 1]])
a = torch.zeros(3, 6)
a[torch.arange(a.size(0)).unsqueeze(1), index] = 1

tensor([[1., 1., 0., 0., 0., 0.],
        [0., 1., 1., 0., 0., 0.],
        [0., 1., 0., 0., 1., 0.]])

while your approach will fail with an IndexError as it’ll try to index dim0 with 4 so they are not compatible.

index = torch.tensor([
    [0, 1],
    [1, 2],
    [4, 1]])
a = torch.zeros(3, 6)
a[index[:, 0], index[:, 1]] = 1

# > IndexError: index 4 is out of bounds for dimension 0 with size 3

I had redefined index so that [4, 1] was not included. My intended point was that rather than assuming each row of a will have 2 elements equal to 1, there is a more general case where any subset of the rows and elements could be chosen. For that setup, I found it more convenient to treat index as a list of (x, y) coordinates i.e.

index = torch.tensor([
    [0, 1],  # First row, second column                      
    [1, 2],  # Second row, third column
             # Nothing in the 3rd row
])

Sure, but your code snippet even with the missing [4, 1] row wouldn’t return the same result, so it’s not a replacement:

index = torch.tensor([
    [0, 1],
    [1, 2]])
a = torch.zeros(3, 6)
a[torch.arange(index.size(0)).unsqueeze(1), index] = 1
print(a)
> tensor([[1., 1., 0., 0., 0., 0.],
          [0., 1., 1., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0.]])

b = torch.zeros(3, 6)
b[index[:, 0], index[:, 1]] = 1
print(b)
> tensor([[0., 1., 0., 0., 0., 0.],
          [0., 0., 1., 0., 0., 0.],
          [0., 0., 0., 0., 0., 0.]])