Search code examples
python-3.xpytorchsparse-matrix

Using Pytorch how to define a tensor with indices and corresponding values


Problem

I have a list of indices and a list of values like so:

i = torch.tensor([[2, 2, 1], [2, 0, 2]])
v = torch.tensor([1, 2, 3])

I want to define a (3x3 for the example) matrix which contains the values v at the indices i (1 at position (2,2), 2 at position (2, 0) and 3 at position (1,2)):

tensor([[0, 0, 0],
        [0, 0, 3],
        [2, 0, 1]])

What I have tried

I can do it using a trick, with torch.sparse and .to_dense() but I feel that it's not the "pytorchic" way to do that nor the most efficient:

f = torch.sparse.FloatTensor(indices, values, torch.Size([3, 3]))
print(f.to_dense())

Any idea for a better solution ? Ideally I would appreciate a solution at least as fast than the one provided above. Of course this was just an example, no particular structure in tensors i and v are assumed (neither for the dimension).


Solution

  • There is an alternative, as below:

    import torch
    
    i = torch.tensor([[2, 2, 1], [2, 0, 2]])
    v = torch.tensor([1, 2, 3], dtype=torch.float)   # enforcing same data-type
    
    target = torch.zeros([3,3], dtype=torch.float)   # enforcing same data-type
    target.index_put_(tuple([k for k in i]), v)
    
    print(target)
    

    The target tensor will be as follows:

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

    This medium.com blog article provides a comprehensive list of all index functions for PyTorch Tensors.