PyTorch如何将(A,B)张量元素映射为数组得到(A,B,N)张量并消除for循环

你之前的写法触发维度不匹配报错，是因为PyTorch的广播机制默认从最后一个维度开始匹配对齐：

• 你生成的x是形状为(N,)的1维张量
• 原张量my_old_tensor是形状为(A,B)的2维张量
  两者最后一个维度的大小分别是N和B，数值不匹配，无法直接广播运算。

只需要给x新增两个前置的单维度，把它的形状调整为(1, 1, N)，就能和(A,B)的原张量自动广播为(A,B,N)的结果，完全消除Python层的for循环：

import torch
A, B, N = 3, 4, 5
my_old_tensor = torch.ones((A,B), dtype=torch.float32)

# 方法1：用索引None新增维度，代码更简洁
x = torch.arange(N, dtype=torch.float32)[None, None, :]
my_new_tensor = (x - my_old_tensor) / 2

# 方法2：用unsqueeze接口新增维度，可读性更强
# x = torch.arange(N, dtype=torch.float32).unsqueeze(0).unsqueeze(0)
# my_new_tensor = (x - my_old_tensor) / 2

你可以用以下代码验证结果和你原来的for循环实现完全一致：

# 原循环实现
my_new_tensor_old = torch.zeros((A, B, N), dtype=torch.float32)
for val in range(N):
    my_new_tensor_old[:,:,val] = (val - my_old_tensor)/2

# 验证结果一致
print(torch.allclose(my_new_tensor, my_new_tensor_old)) # 输出True

该实现是完全向量化的运算，所有计算都在PyTorch的张量内核执行，没有Python循环的开销，CPU和GPU环境下都能获得最优性能。

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