Search code examples
pythonarraysnumpytensorbroadcast

Tensor power of 1D numpy array

I would like to implement a function that takes what I call the "tensor power" of an one-dimensional numpy array:

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

    return y

For any numpy array x of shape (Nx), the output y should be a numpy array of shape (Nx, ..., Nx) where there are n copies of Nx, and its entries are defined as y[i, j, ..., k] = x[i] * x[j] * ... * x[k]. A simple example would be:

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

Is there any easy way to achieve this?

Solution

  • I think you cannot avoid an explicit loop over n. If the tensor becomes large, you may want some more complex structure that makes use of the permutation symmetries.

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