numpy.tensordot

numpy.tensordot(a, b, axes=2)

Returns the tensor dot product for (ndim >= 1) arrays along an axes.

The first element of the sequence determines the axis or axes in a to sum over, and the second element in axes argument sequence determines the axis or axes in b to sum over.

Parameters:

a : array_like

Input array.

b : array_like

Input array.

axes : shape tuple

Axes to be summed over.

See also

dot

Notes

r_{xxx, yyy} = sum_k a_{xxx,k} b_{k,yyy}

When there is more than one axis to sum over, the corresponding arguments to axes should be sequences of the same length with the first axis to sum over given first in both sequences, the second axis second, and so forth.

If the axes argument is an integer, N, then the last N dimensions of a and first N dimensions of b are summed over.

Examples

>>> a = np.arange(60.).reshape(3,4,5)
>>> b = np.arange(24.).reshape(4,3,2)
>>> c = np.tensordot(a,b, axes=([1,0],[0,1]))
>>> c.shape
(5, 2)
>>> c
array([[ 4400.,  4730.],
       [ 4532.,  4874.],
       [ 4664.,  5018.],
       [ 4796.,  5162.],
       [ 4928.,  5306.]])

>>> # A slower but equivalent way of computing the same...
>>> c = np.zeros((5,2))
>>> for i in range(5):
...   for j in range(2):
...     for k in range(3):
...       for n in range(4):
...         c[i,j] += a[k,n,i] * b[n,k,j]