Create 3D scatter matrix

1

I want to create a tensor having specific size like the below,

# a: (B, N, 2) = (2, 3, 2)
a = torch.LongTensor([[[2, 1], [2, 2], [3, 3]], [[2, 0], [3, 2], [1, 2]]])
K = 5

# The expected tensor has the size of (B, K, K) = (2, 5, 5)
# For the first batch,
result[0] = torch.LongTensor([[0, 0, 0, 0, 0],
                              [0, 0, 0, 0, 0],
                              [0, 1, 1, 0, 0],
                              [0, 0, 0, 1, 0],
                              [0, 0, 0, 0, 0]])

The result tensor can be a LongTensor or BoolTensor whatever. I need to use it as a mask.
I want to avoid a loop operation since N might be scalable.

Thanks,

2

Directly indexing the result tensor should work:

a = torch.LongTensor([[[2, 1], [2, 2], [3, 3]], [[2, 0], [3, 2], [1, 2]]])
K = 5

out = torch.zeros(a.size(0), K, K)
out[torch.arange(out.size(0)).unsqueeze(1), a[:, :, 0], a[:, :, 1]] = 1.
print(out)
> tensor([[[0., 0., 0., 0., 0.],
           [0., 0., 0., 0., 0.],
           [0., 1., 1., 0., 0.],
           [0., 0., 0., 1., 0.],
           [0., 0., 0., 0., 0.]],

          [[0., 0., 0., 0., 0.],
           [0., 0., 1., 0., 0.],
           [1., 0., 0., 0., 0.],
           [0., 0., 1., 0., 0.],
           [0., 0., 0., 0., 0.]]])