Cdist using Einops?

[ydalmia]

Hi,

Let me start by saying this is a super cool package! Out of curiosity, I was wondering how one could use Einops to get the squared euclidian distance matrix between two lists of d-dimensional points. For example, if A had 1000 points, each with 32 dimensions, (so shape (1000, 32) and B had 500 points, ie shape (500, 32), then the result would be (5000, 500) with the squared euclidian distance between each point in A and each point in B.

I know that scipy cdist can perform this calculation, but is there a way to use Einops?

[arogozhnikov]

Hi @ydalmia, there are a number of ways to accomplish this, the simplest one is:

import numpy as np
from einops import rearrange, reduce

x = np.random.randn(1000, 32)
y = np.random.randn(500, 32)


diff = rearrange(x, 'i c -> i 1 c') - rearrange(y, 'j c -> 1 j c') 
sq_distances = reduce(diff ** 2, 'i j c -> i j', 'sum')

Here is how you can verify result:

from scipy.spatial.distance import cdist
cdist_result = cdist(x, y, metric='sqeuclidean')
assert np.allclose(sq_distances, cdist_result)

Code in einops will work for pytorch, tensorflow, jax, etc - just use tensors from corresponding frameworks

[ydalmia]

Ah I see, thanks for the response. Wouldn’t this approach use large amounts of memory, as a result of allocating an array that is 5000 x 500 x 32 bytes?

I think it’s probably hard to formulate this using einops in a memory-efficient way, but I was just curious, and asked anyways! Thanks again for your response

[arogozhnikov]

the way it is written, it will require allocation of a large amount of memory, that's correct.

You can either split one of matrices in chunks, or use trick from linear algebra, |a - b|^2 = |a|^2 + |b|^2 - 2 <a, b>
all of those are very computable (e.g. last one is conveniently expressed with einsum). That's kinda generic recommendations.

Thanks again for your response

welcome!