numpy.linalg.tensorsolve 中的向量化

Question

我正在尝试解方程：

斧头 = B

对于 x，其中 A 和 B 是形状为

(3, 3, 3, 3)

和

(3, 3)

的矩阵。这可以使用

numpy.linalg.tensorsolve()

:

来处理

numpy.linalg.tensorsolve(A, B, axes=(0, 2))

我现在正在尝试在 n 不同的实例中求解这个方程，这意味着我的矩阵 A 和 B 现在分别为

(10, 3, 3, 3, 3)

和

(10, 3, 3)

。

在这种情况下，我该如何使用tensorsolve来保持我的代码矢量化？我不知道要指定什么轴或者是否可能。

谢谢。

Answer 1

不幸的是，Numpy 本身并不支持此功能。 然而，看一下tensorsolve代码，事实证明它非常简单。 （它实际上只是在solve()之前调用reshape()。）

因此，可以使用以下代码：

import numpy as np


def btensorsolve(a, b, axes=None):
    # https://github.com/numpy/numpy/blob/
    # e7a123b2d3eca9897843791dd698c1803d9a39c2/numpy/linalg/_linalg.py#L291
    if axes is not None:
        allaxes = list(range(0, an))
        for k in axes:
            allaxes.remove(k)
            allaxes.insert(an, k)
        a = a.transpose(allaxes)

    # find right dimensions
    # a = [2 (dim1) 2 2 3 (dim2) 2 6]
    # b = [2 (dim1) 2 2 3 (dim2)]
    dim2 = b.ndim
    if a.shape[:dim2] != b.shape:
        raise ValueError(
            f"Shapes of a and b are incompatible: " f"{a.shape} and {b.shape}"
        )

    # the dimention of the linear system
    sol_dim = np.prod(a.shape[dim2:])
    sol_dim_ = 1
    for dim1 in range(dim2 - 1, -1, -1):
        sol_dim_ *= a.shape[dim1]
        if sol_dim_ == sol_dim:
            break
    else:
        raise ValueError("Incompatible dimensions")

    a_ = a.reshape(a.shape[:dim1] + (sol_dim, sol_dim))
    b_ = b.reshape(a.shape[:dim1] + (sol_dim, 1))
    x = np.linalg.solve(a_, b_)
    return x.reshape(a.shape[:dim1] + a.shape[dim2:])

from numpy.testing import assert_allclose

a = np.random.randn(2, 2, 2, 3, 2, 6)
b = np.random.randn(2, 2, 2, 3)
assert_allclose(np.einsum("...ijklm,...lm->...ijk", a, btensorsolve(a, b)), b)

Answer 2

您可以用纯Python以函数式的方式描述操作。

list(map(lambda x,y: np.linalg.tensorsolve(x,y), A, B))