Numpy - Matrix multiplication to return ndarray, not sum

Question

All, I have an application that requires returning a numpy ndarray, rather than a simple sum, when multiplying two matrices; e.g.:

import numpy as np
x = np.array([[1, 1, 0], [0, 1, 1]])
y = np.array([[1, 0, 0, 1], [1, 0, 1, 0], [0, 0, 0, 0]])
w = x @ y
>>> array([[2, 0, 1, 1],
           [1, 0, 1, 0]])

However, the requirement is to return an ndarray (in this case..):

array([[[1,1,0], [0,0,0], [0,1,0], [1,0,0]],
       [[0,1,0], [0,0,0], [0,1,0], [0,0,0]]])

Note that the matrix multiplication operation may be repeated; the output will be used as the left-side matrix of ndarrays for the next matrix multiplication operation, which would yield a higher-order ndarray after the second matrix multiplication operation, etc..

Any way to achieve this? I've looked at overloading __add__, and __radd__ by subclassing np.ndarray as discussed here, but mostly got dimension incompatibility errors.

Ideas?

Update:

Addressing @Divakar's answer E.g., for chained operation, adding

z = np.array([[1, 1, 0], [0, 0, 0], [1, 0, 0], [0, 1, 0]])
s1 = x[...,None] * y
s2 = s1[...,None] * z

results in an undesired output.

I suspect the issue starts with s1, which in the case above returns s1.shape = (2,3,4). It should be (2,4,3) since [2x3][3x4] = [2x4], but we're not really summing here, just return an array of length 3.

Similarly, s2.shape should be (2,3,4,3), which [incidentally] it is, but with undesired output (it's not 'wrong', just not what we're looking for). To elaborate, s1*z should be [2x4][4x3] = [2x3] matrix. Each element of the matrix is itself an ndarray, of [4x3] since we have 4 rows in z to multiply the elements in s1, and each element in s1 is itself 3 elements long (again, we're not arithmetically adding elements, but return ndarrays with the extended dimension being the row count in the R-matrix of the operation.

Ultimately, the desired output would be:

s2 = array([[[[1, 1, 0],
              [0, 0, 0],
              [0, 1, 0],
              [0, 0, 0]],

              [[1, 1, 0],
               [0, 0, 0],
               [0, 0, 0],
               [1, 0, 0]],

              [[0, 0, 0],
               [0, 0, 0],
               [0, 0, 0],
               [0, 0, 0]]],


             [[[0, 1, 0],
               [0, 0, 0],
               [0, 1, 0],
               [0, 0, 0]],

              [[0, 1, 0],
               [0, 0, 0],
               [0, 0, 0],
               [0, 0, 0]],

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

Divakar Divakar · Accepted Answer · 2017-12-26T06:58:36

Extend them to 3D and leverage broadcasting -

x[:,None] * y.T

Or with np.einsum -

np.einsum('ij,jk->ikj',x,y)

Going by OP's comment and the quote from the question :

... matrix multiplication operation may be repeated; the output will be used as the left-side matrix of ndarrays for the next matrix multiplication operation, which would yield a higher-order ndarray after the second matrix multiplication operation, etc..

It seems, we need to do something along these lines -

s1 = x[...,None] * y
s2 = s1[...,None] * z # and so on.

Though, the order of the axes would be different in this case, but it seems to be the simplest way to extend the solution to a generic number of incoming 2D arrays.

Following the edits in the question, seems like you are placing the incoming arrays from the first axis onwards for element-wise multiplication. So, if I got that right, you can swap axes to get the correct order, like so -

s1c = (x[...,None] * y).swapaxes(1,-1)
s2c = (s1c.swapaxes(1,-1)[...,None] * z).swapaxes(1,-1) # and so on.

If you are only interested in the final output, swap axes only at the final stage and skip those in the intermediate ones.

Numpy - Matrix multiplication to return ndarray, not sum

1 Answers