关于numpy.linalg.tensorsolve的技术咨询：求附带公式的补充说明

核心功能与数学定义 #

tensorsolve 是NumPy中用于求解张量形式线性方程组的函数，是普通矩阵求解函数solve的高维扩展——它能批量处理嵌套在高维张量中的多个线性方程组，无需手动循环展开。

数学公式表述 #

假设我们要解的张量方程为：
$$\mathcal{A} \star \mathcal{X} = \mathcal{B}$$
其中：

• $\mathcal{A}$ 是形状为 (..., N, N) 的张量（最后两个维度为N×N的方阵）；
• $\mathcal{B}$ 是形状为 (..., N) 的张量；
• $\mathcal{X}$ 是待求解的张量，形状与$\mathcal{B}$一致；
• $\star$ 表示对$\mathcal{A}$的最后一个轴和$\mathcal{X}$的最后一个轴进行张量收缩（求和操作），具体展开为：
  $$\sum_{j=0}^{N-1} \mathcal{A}{i_1 i_2 ... i_k i j} \cdot \mathcal{X}{i_1 i_2 ... i_k j} = \mathcal{B}_{i_1 i_2 ... i_k i}$$

若通过axes参数指定自定义收缩轴（长度为k的元组），则$\mathcal{A}$的最后k个维度需构成k阶“方阵”（每个维度大小均为N），$\mathcal{B}$的最后k个维度需与$\mathcal{A}$的前k个指定轴匹配，此时收缩操作会对$\mathcal{A}$的指定轴与$\mathcal{X}$的对应轴求和，最终满足方程。

实际使用示例 #

基础矩阵求解（与solve等价） #

import numpy as np

# 定义2x2矩阵A和向量b
A = np.array([[1, 2], [3, 4]])
b = np.array([5, 11])

# 求解Ax=b
x = np.linalg.tensorsolve(A, b)
print(x)  # 输出: [1. 2.]

批量求解高维张量中的方程组 #

# 定义包含2个2x2矩阵的3维张量A_batch，以及对应的2个向量b_batch
A_batch = np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
b_batch = np.array([[5, 11], [17, 31]])

# 批量求解两组Ax=b
x_batch = np.linalg.tensorsolve(A_batch, b_batch)
print(x_batch)  
# 输出: [[1. 2.]
#        [1. 2.]]

核心原理简化理解 #

tensorsolve 的底层逻辑可拆解为三步：

1. 将高维张量$\mathcal{A}$和$\mathcal{B}$的非求解轴（除指定收缩轴外的维度）“扁平化”为批量维度；
2. 对每个批量中的N×N矩阵和N维向量，调用普通矩阵求解逻辑（与solve一致）；
3. 将求解结果重新塑形为原张量的形状，返回高维解$\mathcal{X}$。

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