本文介绍在 numpy 中根据输入向量维度（一维 `(3,)` 或二维 `(3, n)`）动态生成对应形状的单位矩阵数组：单个 `3×3` 矩阵或形状为 `(3, 3, n)` 的批量单位矩阵，强调简洁、向量化与内存效率。
在科学计算和物理建模中，常需对一批三维向量（如力、速度、旋转向量）并行执行线性变换，此时往往需要与之匹配的批量单位矩阵（即每个样本对应一个独立的 3×3 单位矩阵）。手动用 np.zeros 初始化再逐对角线赋值（如 unit[0,0,:] = 1）不仅冗长，还易出错且无法充分利用 NumPy 的广播与向量化能力。
推荐使用 np.repeat 结合维度重塑实现高效构建。核心思路是：先生成标准单位矩阵 np.eye(3)，再沿新轴“复制” n 次，最终得到形状为 (3, 3, n) 的三维数组：
import numpy as np
def func(x):
d = 3 # 假设为三维空间；也可泛化为 x.shape[0]
if x.ndim == 1:
return np.eye(d) # shape: (3, 3)
elif x.ndim == 2 and x.shape[0] == d:
n = x.shape[1]
identity_2d = np.eye(d) # shape: (3, 3)
# 方法1：展平后重复再重塑（注意顺序）
unit = np.repeat(identity_2d, n).reshape(d, d, n)
# 方法2（更直观）：升维后沿最后一轴重复
unit = np.repeat(identity_2d[:, :, None], n, axis=-1) # shape: (3, 3, n)
return unit
else:
raise ValueError("Input must be 1D (3,) or 2D (3, n)")✅ 优势说明：
- np.repeat(..., axis=-1) 显式指定扩展轴，语义清晰，避免 reshape 引发的索引混淆；
- 全程零显式循环，纯向量化操作，性能优异；
- 内存连续性良好（尤其方法2），利于后续计算加速。
? 进阶推荐：若项目已引入 einops，可进一步提升可读性与表达力：
import einops as eo # 在 func 中替换为： unit = eo.repeat(np.eye(3), 'i j -> i j n', n=x.shape[1])
该写法以声明式语法明确表达“将 (i,j) 矩阵沿新维度 n 复制”，逻辑直白，且天然支持任意维度扩展（如 (3,3,n,m)），是处理复杂张量重排的工业级实践方案。
⚠️ 注意事项：
- 确保输入 x 的第一维恒为 3（即 x.shape[0] == 3），否则需动态提取 d = x.shape[0] 并调用 np.eye(d)；
- 避免使用 np.tile——它会复制整个数组块，可能导致冗余内存占用；
- 若后续需参与 @ 矩阵乘法（如 unit @ x），注意 x 形状应为 (3, n)，而 unit @ x 不直接兼容；此时建议转为 np.einsum('ijk,jk->ik', unit, x) 或使用 np.matmul 配合 np.moveaxis 调整维度。
综上，结合 np.repeat（基础可靠）或 einops.repeat（高阶清晰），即可优雅、高效地实现单位矩阵的标量/批量统一构造。
