Sometimes array processing involves operations that include both, first broadcasting, then reducing along a specific axis. By using the natural broadcasting syntax, this will create intermediary arrays that contain all dimensions, the final ones as well as the ones that will be reduced, which might lead to significantly increased memory usage (beyond what is necessary from the algorithmic point of view).
Consider the following example function will computes the MSE between two vectors:
def mse(x, y):
return np.mean((x - y)**2, axis=-1)
Now, if I have a collection of vectors and want to compute the pairwise mse between all vectors, I can do that in the following way:
N, d = 100, 3
vectors = np.random.uniform(size=(N, d))
pairwise = mse(vectors[None, :, :], vectors[:, None, :])
However this creates an intermediary array x - y of shape (N, N, d) which might consume too much memory, especially if d is large. From the algorithmic point of view, this is not necessary and only (N, N) memory is required:
pairwise = np.empty(shape=(len(vectors),)*2, dtype=vectors.dtype)
for i in range(len(vectors)):
pairwise[i] = mse(vectors[i:i+1], vectors)
This however means that the concise syntax mse(vectors[None, :, :], vectors[:, None, :]) cannot be used anymore.
So my question is, is there any way to instruct Numpy to perform the broadcasting lazily and only evaluate it when a certain reduce operation is encountered (e.g. np.mean)?