Jasper Zhou
Jasper Zhou

Reputation: 527

how to properly use np.subtract's broadcast

I have 2 matrices, the shape of the first one is (2,64) and the shape of the second one is (2,256,64), now I want to do np.subtract between this 2 matrices, because np.subtract(matrix1, matrix2) cannot broadcast automatically, what I did is below

step_1 = np.subtract(matrix1[0], matrix2[0]).shape  ## shape is (256,64)
step_2 = np.subtract(matrix1[1], matrix2[1]).shape  ## shape is (256,64)
res = np.array([step_1, step_2]) ## shape is (2,256,64)

or

res = np.array([np.subtract(matrix1[i], matrix2[i]) for i in range(2)]) ## shape is (2,256,64)

Can I do something like this only using np.subtract(By setting some kinds of parameters) in a single step to get the same answer, (or using other technique like np.swapaxes)?

Upvotes: 1

Views: 146

Answers (1)

Mad Physicist
Mad Physicist

Reputation: 114230

You can introduce a new axis using None (which is an alias for np.newaxis) to line them up directly:

matrix_1[:, None, :] - matrix2

Unless you want to use some of the features of an explicit call to np.subtract, the minus operator (-) is cleaner.

Another alternative is to get the same view with np.expand_dims:

np.expand_dims(matrix1, 1) - matrix2

You can also reshape:

matrix1.reshape(matrix1.shape[0], 1, *matrix1.shape[1:]) - matrix2

The solution you propose with swapaxes is a bit drastic, but will work:

(matrix1 - matrix2.swapaxes(0, 1)).swapaxes(0, 1)

The final swapaxes on the result is necessary to get the original shape back. You can achieve similar results with transpose:

(matrix1 - matrix2.transpose(1, 0, 2)).swapaxes(1, 0, 2)

Upvotes: 1

Related Questions