Convert Numpy vector subtraction to Pytorch tensor subtraction

I'm trying to use this code (from here) but in Pytorch (it's an N-body simulation):

mass = 20.0*np.ones((500,1))/500  # total mass of particles is 20
pos  = np.random.randn(500,3)
G = 1.0

# positions r = [x,y,z] for all particles
x = pos[:,0:1]
y = pos[:,1:2]
z = pos[:,2:3]

# matrix that stores all pairwise particle separations: r_j - r_i
dx = x.T - x
dy = y.T - y
dz = z.T - z

inv_r3 = (dx**2 + dy**2 + dz**2)
inv_r3[inv_r3>0] = inv_r3[inv_r3>0]**(-1.5)

ax = G * (dx * inv_r3) @ mass
ay = G * (dy * inv_r3) @ mass
az = G * (dz * inv_r3) @ mass

# pack together the acceleration components
a = np.hstack((ax,ay,az))

I know I can break it down per dimension in pytorch:

dx = torch.tensor(pos[:,0:1]).T - torch.tensor(pos[:,0:1])

The issue is that my tensor is of much larger size than 3 dimension (e.g., torch.rand(500,1000) instead of np.random.randn(500,3)) so breaking it as done here (e.g., x = pos[:,0:1]) is not very practical. Is there a way to have the same code but with a Pytorch tensor of large dimensions without splitting it per dimension?

Let's first deal with the differential terms.

differentials = pos.T[:,None,:] - torch.swapaxes(pos.T[:,None,:], 1, 2)

The first part transposes pos and adds an extra axis in the middle. Doing that effectively separates the columns from each other. Then, the axis swap does what the transpose did in the 1D array, but only applies it to the desired axes.

You could also save the pos.T[:,None,:] to an intermediate variable if you so desire.

For inv_r3, that is just summing together the squares of the differentials, which are indexed by axis 0.

inv_r3 = torch.sum(differentials**2, axis=0)

The remaining lines are unchanged, though the acceleration terms are already together, it just needs to be squeezed and transposed.

All together:

import torch

N = 500
dims = 3
mass = 20.0*torch.ones((N,1))/N
pos = torch.randn(N,dims)
G = 1.0

differentials = pos.T[:,None,:] - torch.swapaxes(pos.T[:,None,:], 1, 2)
inv_r3 = torch.sum(differentials**2, axis=0)
inv_r3[inv_r3>0] = inv_r3[inv_r3>0]**(-1.5)
acceleration = (G*differentials*inv_r3@mass).squeeze(axis=-1).T