Compute x**k with x, k being arrays of arbitrary dimensionality

Question

I have two numpy arrays: One array x with shape (n, a0, a1, ...) and one array k with shape (n, b0, b1, ...). I would like to compute and array of exponentials such that the output has dimension (a0, a1, ..., b0, b1, ...) and

out[i0, i1, ..., j0, j1, ...] == prod(x[:, i0, i1, ...] ** k[:, j0, j1, ...])

If there is only one a_i and one b_j, broadcasting does the trick via

import numpy

x = numpy.random.rand(2, 31)
k = numpy.random.randint(1, 10, size=(2, 101))

out = numpy.prod(x[..., None]**k[:, None], axis=0)

If x has a few dimensions more, more Nones have to be added:

x = numpy.random.rand(2, 31, 32, 33)
k = numpy.random.randint(1, 10, size=(2, 101))

out = numpy.prod(x[..., None]**k[:, None, None, None], axis=0)

If x has a few dimensions more, more Nones have to be added at other places:

x = numpy.random.rand(2, 31)
k = numpy.random.randint(1, 10, size=(2, 51, 51))

out = numpy.prod(x[..., None, None]**k[:, None], axis=0)

How to make the computation of out generic with respect to the dimensionality of the input arrays?

Answer 1

Here's one using reshaping of the two arrays so that they are broadcastable against each other and then performing those operations and prod reduction along the first axis -

k0_shp = [k.shape[0]] + [1]*(x.ndim-1) + list(k.shape[1:])
x0_shp = list(x.shape) + [1]*(k.ndim-1)
out = (x.reshape(x0_shp) ** k.reshape(k0_shp)).prod(0)

Here's another way to reshape both inputs to 3D allowing one singleton dim per input and such that they are broadcastable against each other, perform prod reduction to get 2D array, then reshape back to multi-dim array -

s = x.shape[1:] + k.shape[1:] # output shape
out = (x.reshape(x.shape[0],-1,1)**k.reshape(k.shape[0],1,-1)).prod(0).reshape(s)

It must be noted that reshaping merely creates a view into the input array and as such is virtually free both memory-wise and performance-wise.

Answer 2

Without understanding fully the math of what you're doing, it seems that you need a constant number of None's for the number of dimensions of each x and k.

does something like this work?

out = numpy.prod(x[[...]+[None]*(k.ndim-1)]**k[[slice(None)]+[None]*(x.ndim-1)])

Here are the slices separately so they're a bit easier to read:

x[  [...] + [None]*(k.ndim-1)  ]

k[  [slice(None)] + [None]*(x.ndim-1)  ]

Compatibility Note:

[...] seems to only be valid in python 3.x If you are using 2.7 (I haven't tested lower) substitute [Ellipsis] instead:

x[  [Ellipsis] + [None]*(k.ndim-1)  ]