I am having a system of atoms from which I want to create a graph. Let’s also say for clarity that I am only interested in distances between atoms as feature at the moment.

Moreover, let’s assume we have a function which from input coordinates spits out distances for **all** atoms to all other atoms within **constant radius** (in other words it takes only neighboring atoms which are within a sphere of this radius). This is very important because it means that it is very likely that different atoms will have **different** number of neighbors

Now… that’s when I encountered a problem

So, as I just mentioned I have a function which returns me a 1-D tensor consisting of distances for all atoms (i.e. torch.Size([#DistancesFromAtomsToAllOtherAtoms]), from which I want to build 2-D tensor of the dimensions torch.Size([#NumAtoms, #Neighbors] but the second dimensions “Neighbors” has to be **padded**.

For example let’s say we are looking at three atoms and the first has 20 neighbors, second 23 and third 21, then 1-D tensor will have the shape of torch.Size([64]), while the 2-D tensor needs to have shape of torch.Size([3,23]) which asks for 69 elements and I get error when I just try to index new tensor with some mask (i.e output[mask] = distances), where mask is boolean tensor with the dimensions torch.Size([3,23])

For this specific example Pytorch error would be:

RuntimeError: The size of tensor a (23) must match the size of tensor b (64) at non-singleton dimension 1

Finally, it’s worth noting that the same process with analogous functions, and with arrays with same dimensions as tensors above, work perfectly in numpy. It just puts zero in resulting tensor whenever the value of mask array is False, and takes value from source array when the value of mask array is True.

Hope someone could suggest me a way to do the same indexing process in Pytorch, while **retaining** the possibility to do **backprop pass** on this operation.

Thanks in advance,

Ivan