Dot product along third axis

Question

I'm trying to take a tensor dot product in numpy using tensordot, but I'm not sure how I should reshape my arrays to achieve my computation. (I'm still new to the mathematics of tensors, in general.)

I have

arr = np.array([[[1, 1, 1],
                [0, 0, 0],
                [2, 2, 2]],

               [[0, 0, 0],
                [4, 4, 4],
                [0, 0, 0]]])

w = [1, 1, 1]

And I want to take a dot product along axis=2, such that I have the matrix

array([[3, 0, 6],
       [0, 12, 0]])

What's the proper numpy syntax for this? np.tensordot(arr, [1, 1, 1], axes=2) seems to raise a ValueError.

Answer 1

The reduction is along axis=2 for arr and axis=0 for w. Thus, with np.tensordot, the solution would be -

np.tensordot(arr,w,axes=([2],[0]))

Alternatively, one can also use np.einsum -

np.einsum('ijk,k->ij',arr,w)

np.matmul also works

np.matmul(arr, w)

Runtime test -

In [52]: arr = np.random.rand(200,300,300)

In [53]: w = np.random.rand(300)

In [54]: %timeit np.tensordot(arr,w,axes=([2],[0]))
100 loops, best of 3: 8.75 ms per loop

In [55]: %timeit np.einsum('ijk,k->ij',arr,w)
100 loops, best of 3: 9.78 ms per loop

In [56]: %timeit np.matmul(arr, w)
100 loops, best of 3: 9.72 ms per loop

hlin117 tested on Macbook Pro OS X El Capitan, numpy version 1.10.4.