Get multiplication table generalized for n dimensions with numpy

Let a = np.arange(1, 4).

To get the 2 dimensional multiplication table for a, I do:

>>> a * a[:, None]      
>>> array([[1, 2, 3],
           [2, 4, 6],
           [3, 6, 9]])

For 3 dimensions, I can do the following:

>>> a * a[:, None] * a[:, None, None] 
>>> array([[[ 1,  2,  3],
            [ 2,  4,  6],
            [ 3,  6,  9]],

           [[ 2,  4,  6],
            [ 4,  8, 12],
            [ 6, 12, 18]],

           [[ 3,  6,  9],
            [ 6, 12, 18],
            [ 9, 18, 27]]])

How could I write a function that takes a numpy array a and a number of dimensions n as input and ouputs the n dimensional multiplication table for a?

Answer

This should do what you need:

import itertools
a = np.arange(1, 4)
n = 3

def f(x, y):
    return np.expand_dims(x, len(x.shape))*y

l = list(itertools.accumulate(np.repeat(np.atleast_2d(a), n, axis=0), f))[-1]

Just change n to be whatever dimension you need