一维numpy数组的张量幂

Question

我想实现一个函数，它采用我所说的一维 numpy 数组的“张量幂”：

def tensor_pow(x: np.ndarray, n: int) -> np.ndarray:
    
    # code goes here

    return y

对于任何形状为

x

的 numpy 数组

(Nx)

，输出

y

应该是形状为

(Nx, ..., Nx)

的 numpy 数组，其中存在

n

的

Nx

副本，其条目定义为

y[i, j, ..., k] = x[i] * x[j] * ... * x[k]

。一个简单的例子是：

y = tensor_pow(np.arange(3), 2)  # an array of shape (3, 3)
y[1, 2] == 2                     # returns True

有什么简单的方法可以实现这一点吗？

Answer 1

我认为你无法避免

n

上的显式循环。如果张量变大，您可能需要一些利用排列对称性的更复杂的结构。

def tensor_pow(x, n):
    y = x
    for i in range(1, n):
        y = y[..., None] * x 
    return y