基于形状为(N,d)的矩阵a、b，如何向量化构造(N,d,d)张量T？

嘿，这个需求其实用**数组广播（Broadcasting）**就能完美解决，比折腾tensordot简单多啦！核心思路是通过扩展两个矩阵的维度，让它们能在对应位置进行逐元素相乘，自动生成你要的(N,d,d)张量，全程不需要任何for循环。

具体实现（以NumPy为例） #

假设你已经有了形状为(N,d)的矩阵a和b，只需要给两个矩阵分别扩展一个维度，再直接相乘就行：

import numpy as np

# 示例数据：N=3，d=2
N, d = 3, 2
a = np.array([[1,2], [3,4], [5,6]])
b = np.array([[7,8], [9,10], [11,12]])

# 扩展维度后相乘
T = a[:, :, np.newaxis] * b[:, np.newaxis, :]

# 验证形状
print(T.shape)  # 输出 (3, 2, 2)，完全符合要求

代码解释 #

• a[:, :, np.newaxis]：把a从(N,d)的形状扩展为**(N,d,1)**，相当于给每个向量a[i]增加了一个长度为1的第三维度。
• b[:, np.newaxis, :]：把b从(N,d)的形状扩展为**(N,1,d)**，给每个向量b[i]增加了一个长度为1的第二维度。
• 当这两个扩展后的数组相乘时，NumPy的广播机制会自动把它们的形状对齐到(N,d,d)，每个位置T[i,p,q]自然就是a[i,p] * b[i,q]。

更清晰的写法（用np.expand_dims） #

如果你觉得np.newaxis不够直观，也可以用np.expand_dims来明确指定扩展的轴：

T = np.expand_dims(a, axis=2) * np.expand_dims(b, axis=1)

其他框架的实现（PyTorch/TensorFlow） #

思路完全一致，只是扩展维度的API略有不同：

PyTorch #

import torch

a = torch.rand(N, d)
b = torch.rand(N, d)
T = a.unsqueeze(2) * b.unsqueeze(1)

TensorFlow #

import tensorflow as tf

a = tf.random.normal((N, d))
b = tf.random.normal((N, d))
T = tf.expand_dims(a, 2) * tf.expand_dims(b, 1)

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