Is there a JAX util function for pairwise distance?

[yuhanfu]

Hi,

I'm new to JAX and wonder if there is an equivalent function for
tensorflow_addons.losses.metric_learning.pairwise_distance to compute the pairwise distance matrix. I know it's not hard to implement the (squared) euclidean distance in JAX, but wanted to ask in case it already exists.

Thanks!

[jakevdp]

No, there's no special function for this (though it's been discussed in #9235)

In case someone comes across this in a search, here's how you might implement simple pairwise euclidean distances in JAX using broadcasted vector operations, which is probably the most efficient way to compute this robustly in JAX, and is probably how any future pairwise API would be implemented:

def pairwise_euclidean(x, y):
  assert x.ndim == y.ndim == 2
  return jnp.sqrt(((x[:, None, :] - y[None, :, :]) ** 2).sum(-1))

(I say "robustly" because there are ways you can factor-out terms to compute the same result in fewer operations, but it is numerically unstable in edge cases)

[mattjj]

A vmap version might look like

def pairwise_distances(dist: Callable, xs, ys):
  return jax.vmap(lambda x: jax.vmap(lambda y: dist(x, y))(xs))(ys)

[YouJiacheng]

@mattjj the parameter name seems confusing - xs is fed into y, ys is fed into x, I guess this is a typo. (Or actually a downside of vmap version, i.e. the parameter declaring order and argument passing order are reverse)

[mattjj]

Good catch @YouJiacheng ! Just a mistake on my part. Thanks for noticing it.

Since naming is the hardest problem in computer science, maybe we should just write

pairwise_distances = jax.vmap(jax.vmap(dist, (0, None)), (None, 0))

[yuhanfu]

No, there's no special function for this (though it's been discussed in #9235)

In case someone comes across this in a search, here's how you might implement simple pairwise euclidean distances in JAX using broadcasted vector operations, which is probably the most efficient way to compute this robustly in JAX, and is probably how any future pairwise API would be implemented:

def pairwise_euclidean(x, y):
  assert x.ndim == y.ndim == 2
  return jnp.sqrt(((x[:, None, :] - y[None, :, :]) ** 2).sum(-1))

(I say "robustly" because there are ways you can factor-out terms to compute the same result in fewer operations, but it is numerically unstable in edge cases)

What are the edge cases where some other implementations would produce numerically unstable results?

I implemented it as

def pairwise_euclidean(x, y):
  x_norm = jnp.expand_dims(jnp.sum(x ** 2, axis=1), 1)
  y_norm = jnp.expand_dims(jnp.sum(y ** 2, axis=1), 0)
  cross = jnp.matmul(x, y.T)
  return jnp.sqrt(x_norm - 2 * cross + y_norm)

[jakevdp]

Basically you'll run into issues any time the distance is small compared to the absolute value of the coordinates. For example, using your function:

x = jnp.array([[0.0, 1.0]], dtype='float32')
y = jnp.array([[1.0, 0.0]], dtype='float32')

print(pairwise_euclidean(x, y))
# [[1.4142135]]

print(pairwise_euclidean(x + 10, y + 10))
# [[1.4142135]]

print(pairwise_euclidean(x + 4000, y + 4000))  # <-- fails in float32
# [[0.]]

This isn't an issue for the direct (x - y) ** 2 computation.