Is there a way to use numpy.outer on only a subset of dimensions?

Question

I have an array of arrays, like this:

import numpy as np
a = np.array([[1, 2, 3], [4, 5, 6]])

I would like to calculate the pairwise differences between the array elements, something like:

[[[ 0,  0,  0], [-3, -3, -3]],
 [[ 3,  3,  3], [ 0,  0,  0]]

My first thought was to use np.subtract.outer(a, a), but that doesn't do what I want - it goes one layer too deep into the arrays. I can see that the numbers I need are in the output of np.subtract.outer(a, a), but the array that I'm actually working with is very large, and I don't have enough memory to be able to allocate for the result.

Thankyou!

Answer 1

You can simply use broadcasting to solve this.

a[:, None, :] - a[None, :, :]

Gives you what you want.