如何实现一维numpy数组的tensor power（张量幂）函数？

这里有几种简便的实现方式，都能满足你的需求：

方法1：用reduce结合multiply.outer #

利用np.multiply.outer的外积特性，通过functools.reduce累积计算n次外积，这是最简洁的实现方式之一：

import numpy as np
from functools import reduce

def tensor_pow(x: np.ndarray, n: int) -> np.ndarray:
    if n < 1:
        raise ValueError("n必须是正整数")
    return reduce(np.multiply.outer, [x] * n)

验证示例 #

y = tensor_pow(np.arange(3), 2)
print(y.shape)  # 输出 (3, 3)
print(y[1, 2] == 2)  # 输出 True

原理：np.multiply.outer(a, b)会生成a和b的外积数组，形状为a.shape + b.shape。多次调用后，就能得到n个维度的张量，每个元素对应位置的x元素乘积。

方法2：利用广播机制 #

通过将x变形为不同维度的视图，利用NumPy的广播特性实现元素级相乘：

import numpy as np

def tensor_pow(x: np.ndarray, n: int) -> np.ndarray:
    if n < 1:
        raise ValueError("n必须是正整数")
    # 为每个维度生成对应的广播视图
    views = []
    for i in range(n):
        # 构造形状：前i个维度为1，中间是x的长度，后n-i-1个维度为1
        shape = [1] * i + [-1] + [1] * (n - i - 1)
        views.append(x.reshape(shape))
    
    result = 1
    for v in views:
        result *= v
    return result

验证示例 #

y = tensor_pow(np.arange(3), 3)
print(y.shape)  # 输出 (3, 3, 3)
print(y[0, 1, 2] == 0 * 1 * 2)  # 输出 True

原理：每个视图在对应维度上保留x的原始长度，其余维度为1，广播机制会自动扩展这些维度，最终相乘得到所有元素的乘积组合。

方法3：用np.einsum实现 #

通过构造爱因斯坦求和的下标字符串，直接指定张量的乘积规则：

import numpy as np

def tensor_pow(x: np.ndarray, n: int) -> np.ndarray:
    if n < 1:
        raise ValueError("n必须是正整数")
    # 生成输入和输出的下标字符串，比如n=2时是"i,j->ij"
    input_indices = ','.join([chr(ord('a') + i) for i in range(n)])
    output_indices = ''.join([chr(ord('a') + i) for i in range(n)])
    return np.einsum(f"{input_indices}->{output_indices}", *([x] * n))

验证示例 #

y = tensor_pow(np.array([2, 3]), 2)
print(y)
# 输出：
# [[4 6]
#  [6 9]]

原理：np.einsum通过下标字符串定义张量的运算规则，这里我们让每个输入数组对应独立的下标，输出保留所有下标，本质就是计算所有元素的乘积组合。

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