将Numpy向量减法转PyTorch张量减法：高维张量适配问询

我尝试将这份N体模拟的Numpy代码转换为PyTorch实现：

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))

我知道可以按维度拆分实现PyTorch版本：

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

但我的张量维度远大于3（例如torch.rand(500,1000)），按维度拆分的方式并不实用。请问是否有无需按维度拆分即可适配高维张量的实现方法？

可以利用PyTorch的广播机制和张量维度变换实现通用的高维版本，不需要逐维度拆分：

import torch

# 初始化高维数据（示例：500个粒子，1000维位置）
mass = 20.0 * torch.ones((500, 1)) / 500  # 总质量20
pos = torch.randn(500, 1000)
G = 1.0

# 计算所有粒子对的位置差: r_j - r_i，形状变为(500,500,1000)
dx = pos.unsqueeze(1) - pos.unsqueeze(0)

# 计算距离平方和，形状(500,500)
r_sq = torch.sum(dx ** 2, dim=-1)

# 计算1/r³，避免除以0（对角线是粒子自身，置0）
inv_r3 = torch.where(r_sq > 0, r_sq ** (-1.5), torch.zeros_like(r_sq))

# 计算每个维度的加速度，最终形状(500,1000)
acc = G * torch.matmul(dx * inv_r3.unsqueeze(-1), mass)

关键步骤说明： #

• pos.unsqueeze(1) - pos.unsqueeze(0)：通过维度扩展，将位置张量从(500, D)分别转为(500,1,D)和(1,500,D)，广播后直接得到所有粒子对的位置差(500,500,D)，其中D是任意位置维度数。
• torch.sum(dx**2, dim=-1)：对最后一个维度（位置维度）求和，得到所有粒子对的距离平方。
• torch.where：安全处理对角线的0值（粒子自身到自身的距离），避免计算0的-1.5次幂出现NaN。
• inv_r3.unsqueeze(-1)：将inv_r3从(500,500)扩展为(500,500,1)，和dx(500,500,D)广播相乘后，再与mass(500,1)做矩阵乘法，最终得到每个粒子在D个维度上的加速度。

这个实现完全通用，不管位置维度是3还是1000都能直接运行，无需修改代码。

内容的提问来源于stack exchange，提问作者Penguin