Is layered broadcasting available in NumPy?

Question

I wonder if there is a built-in operation which would free my code from Python-loops.

The problem is this: I have two matrices A and B. A has N rows and B has N columns. I would like to multiply every i row from A with corresponding i column from B (using NumPy broadcasting). The resulting matrix would form i layer in the output. So my result would be 3-dimensional array.

Is such operation available in NumPy?

Answer 1

Yes, in it's simplest form you just add "zero" dimensions so the NumPy broadcasts along the rows of A and columns of B:

>>> import numpy as np

>>> A = np.arange(12).reshape(3, 4) # 3 row, 4 colums
>>> B = np.arange(15).reshape(5, 3) # 5 rows, 3 columns
>>> res = A[None, ...] * B[..., None]
>>> res
array([[[  0,   0,   0,   0],
        [  4,   5,   6,   7],
        [ 16,  18,  20,  22]],

       [[  0,   3,   6,   9],
        [ 16,  20,  24,  28],
        [ 40,  45,  50,  55]],

       [[  0,   6,  12,  18],
        [ 28,  35,  42,  49],
        [ 64,  72,  80,  88]],

       [[  0,   9,  18,  27],
        [ 40,  50,  60,  70],
        [ 88,  99, 110, 121]],

       [[  0,  12,  24,  36],
        [ 52,  65,  78,  91],
        [112, 126, 140, 154]]])

The result has a shape of (5, 3, 4) and you can easily move the axis around if you want a different shape. For example using np.moveaxis:

>>> np.moveaxis(res, (0, 1, 2), (2, 0, 1))  # 0 -> 2 ; 1 -> 0, 2 -> 1
array([[[  0,   0,   0,   0,   0],
        [  0,   3,   6,   9,  12],
        [  0,   6,  12,  18,  24],
        [  0,   9,  18,  27,  36]],

       [[  4,  16,  28,  40,  52],
        [  5,  20,  35,  50,  65],
        [  6,  24,  42,  60,  78],
        [  7,  28,  49,  70,  91]],

       [[ 16,  40,  64,  88, 112],
        [ 18,  45,  72,  99, 126],
        [ 20,  50,  80, 110, 140],
        [ 22,  55,  88, 121, 154]]])

With a shape of (3, 4, 5).

Answer 2

One way to express your requirement directly is by using np.einsum():

>>> A = np.arange(12).reshape(3, 4)
>>> B = np.arange(15).reshape(5, 3)
>>> np.einsum('...i,j...->...ij', A, B)
array([[[  0,   0,   0,   0,   0],
        [  0,   3,   6,   9,  12],
        [  0,   6,  12,  18,  24],
        [  0,   9,  18,  27,  36]],

       [[  4,  16,  28,  40,  52],
        [  5,  20,  35,  50,  65],
        [  6,  24,  42,  60,  78],
        [  7,  28,  49,  70,  91]],

       [[ 16,  40,  64,  88, 112],
        [ 18,  45,  72,  99, 126],
        [ 20,  50,  80, 110, 140],
        [ 22,  55,  88, 121, 154]]])

This uses the Einstein summation convention.

For further discussion, see chapter 3 of Vectors, Pure and Applied: A General Introduction to Linear Algebra by T. W. Körner. In it, the author cites an amusing passage from Einstein's letter to a friend:

"I have made a great discovery in mathematics; I have suppressed the summation sign every time that the summation must be made over an index which occurs twice..."