能否指定PyTorch最小二乘解输出维度？求特定形状x张量

能否得到形状为1 × k × m的张量x？ #

可以，但不能直接用默认的torch.linalg.lstsq调用，需要先调整输入张量的维度结构：

你的场景中，torch.linalg.lstsq默认会把A的第一个维度（3）当作独立batch维度，对每个3对应的(n, k)切片和B的(n, m)切片单独计算最小二乘，因此输出会是3 × k × m的张量——每个batch对应一个解。

如果要得到共享的1 × k × m的x（即让所有3组Ax=B共享同一个最优解），需要把A的前两个维度合并，将问题转化为全局的最小二乘：

1. 重构A：将[3, n, k] reshape为[3*n, k]，把3组n×k的矩阵拼接成一个大的3n×k矩阵；
2. 重构B：将[3, n, m] reshape为[3*n, m]，同理拼接成3n×m的矩阵；
3. 调用torch.linalg.lstsq(A_reshaped, B_reshaped)，得到的解形状为[k, m]；
4. 最后给解添加一个维度，用x = solution.unsqueeze(0)得到[1, k, m]的张量。

是否存在指定PyTorch最小二乘解输出维度的方法？ #

PyTorch的torch.linalg.lstsq没有直接指定输出维度的参数，它的输出形状完全由输入A和B的维度决定：

• 输入A的形状为(*batch_dims, M, N)，B的形状为(*batch_dims, M, K)；
• 输出解的形状为(*batch_dims, N, K)，其中*batch_dims是A和B对齐的批量维度，最后两个维度是最小二乘的计算维度。

要得到特定形状的解，只能通过手动调整输入的维度结构（比如合并/拆分批量维度、转置等），再对输出做维度变换来实现。

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