Python tensor matrix multiply

问题

I have the tensor

A = 
[[[a,b],
  [c,d]],
 [[e,f],
  [g,h]]]

and the matrix

B = 
[[1,2],
 [3,4]]

I need to get

C = 
[[a*1+e*2,b*1+f*2],
 [c*3+g*4,d*3+h*4]]

How can I do this using numpy in matrix form? I've looked into np.tensordot() but it doesn't seem to help in this case.

回答1:

You can try this:

>>> import numpy as np
>>> a = np.arange(1,9).reshape(2,2,2)
>>> a
array([[[1, 2],
        [3, 4]],

       [[5, 6],
        [7, 8]]])
>>> b = np.arange(1,5).reshape(2,2)
>>> b
array([[1, 2],
       [3, 4]])
>>> (a * b[None,:,:].T).sum(axis = 0)
array([[11, 14],
       [37, 44]])

Intermediate steps look like this:

>>> b[None,:,:]
array([[[1, 2],
        [3, 4]]])
>>> b[None,:,:].T
array([[[1],
        [3]],

       [[2],
        [4]]])

回答2:

OP's problem can be reformulated in a standard format using tensor notation and the so called Einstein summation convention

A_{k i j} B_{i k} ⇒ C_{i j}

Numpy has a convenient utility function to perform the kind of tensor operations that can be descibed using the Einstein summation convention, unsurprisingly named numpy.einsum, that allows to straightforwardly map the tensor notation to optimized C-level loops by the means of an instruction string that exactly reflects the tensor notation, 'kij, ik -> ij'

import numpy as np
a = np.arange(8).reshape(2,2,2)+1
b = np.arange(4).reshape(2,2)+1
c = np.einsum('kij, ik -> ij', a, b)
print(c)
# [[11 14]
#  [37 44]]

Merits of numpy.einsum

The source code documents the details of the operation performed.

np.einsumtypically is fast

In [12]: import numpy as np 
    ...:  
    ...: i, j, k = 100, 320, 140 # just three largish numbers
    ...: a = np.random.random((k,i,j)) 
    ...: b = np.random.random((i,k))                                                      

In [13]: %timeit np.einsum('kij,ik->ij', a, b)                                            
7.47 ms ± 82.8 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

In [14]: %timeit (a * b[None,:,:].T).sum(axis = 0)                                        
49.3 ms ± 6.77 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

来源：https://stackoverflow.com/questions/60357297/python-tensor-matrix-multiply

标签

python

arrays

numpy

linear-algebra

tensor