Search code examples
pythonpython-3.xnumpynumpy-einsum

Multiplying elementwise over final axis of two arrays


Given a 3d array and a 2d array,

a = np.arange(10*4*3).reshape((10,4,3))
b = np.arange(30).reshape((10,3))

How can I run elementwise-multiplication across the final axis of each, resulting in c where c has the shape .shape as a? I.e.

c[0] = a[0] * b[0]
c[1] = a[1] * b[1]
# ...
c[i] = a[i] * b[i]

Solution

  • Without any sum-reduction involved, a simple broadcasting would be really efficient after extending b to 3D with np.newaxis/None -

    a*b[:,None,:] # or simply a*b[:,None]
    

    Runtime test -

    In [531]: a = np.arange(10*4*3).reshape((10,4,3))
         ...: b = np.arange(30).reshape((10,3))
         ...: 
    
    In [532]: %timeit np.einsum('ijk,ik->ijk', a, b) #@Brad Solomon's soln
         ...: %timeit a*b[:,None]
         ...: 
    100000 loops, best of 3: 1.79 µs per loop
    1000000 loops, best of 3: 1.66 µs per loop
    
    In [525]: a = np.random.rand(100,100,100)
    
    In [526]: b = np.random.rand(100,100)
    
    In [527]: %timeit np.einsum('ijk,ik->ijk', a, b)
         ...: %timeit a*b[:,None]
         ...: 
    1000 loops, best of 3: 1.53 ms per loop
    1000 loops, best of 3: 1.08 ms per loop
    
    In [528]: a = np.random.rand(400,400,400)
    
    In [529]: b = np.random.rand(400,400)
    
    In [530]: %timeit np.einsum('ijk,ik->ijk', a, b)
         ...: %timeit a*b[:,None]
         ...: 
    10 loops, best of 3: 128 ms per loop
    10 loops, best of 3: 94.8 ms per loop