Add Hungarian algorithm for the linear assignment problem.

[carlosgmartin]

#954

[fabianp]

thanks for starting this @carlosgmartin !

At a high level this looks good to me, but it would be important to have an example of this method in the gallery (https://optax.readthedocs.io/en/latest/gallery.html) to showcase its usage

[vroulet]

Thanks for integrating this
Here are a few comments

[carlosgmartin]

@fabianp @vroulet Done.

[vroulet]

Thanks again! A few minor remaining comments

[carlosgmartin]

@vroulet Updated.

[vroulet]

Great work, looks good to me

[carlosgmartin]

Perhaps it's worth mentioning somewhere in the docs that this can be used to compute the Wasserstein distance between empirical distributions:

Code

import argparse

import jax
import optax
from jax import numpy as jnp, random
from matplotlib import collections, pyplot as plt, rcParams


def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--seed", type=int, default=0)
    p.add_argument("--points", type=int, default=200)
    p.add_argument("--power", type=float, default=2.0)
    p.add_argument("--markersize", type=float, default=4.0)
    p.add_argument("--linewidth", type=float, default=1.0)
    return p.parse_args()


def get_wasserstein_distance(x, y, power=2):
    displacements = x[:, None] - y[None, :]
    pow_distances = (jnp.abs(displacements) ** power).sum(-1)
    i, j = optax.assignment.hungarian_algorithm(pow_distances)
    distance = pow_distances[i, j].mean() ** (1 / power)
    return distance, (i, j)


def estimate_wasserstein_distance(key, sample_fn_1, sample_fn_2, n_samples):
    keys = random.split(key)
    x = jax.vmap(sample_fn_1)(random.split(keys[0], n_samples))
    y = jax.vmap(sample_fn_2)(random.split(keys[1], n_samples))
    distance, (i, j) = get_wasserstein_distance(x, y)
    return distance


def main():
    args = parse_args()

    key = random.key(args.seed)

    keys = random.split(key)
    x = random.normal(keys[0], (args.points, 2))
    y = random.normal(keys[1], (args.points, 2)) + jnp.array([0.2, 0.0])

    distance, (i, j) = get_wasserstein_distance(x, y, args.power)

    fig, ax = plt.subplots(constrained_layout=True)

    ax.scatter(*x.T, s=args.markersize**2, label="facility", edgecolor="none")
    ax.scatter(*y.T, s=args.markersize**2, label="client", edgecolor="none")

    data = jnp.stack((x[i], y[j]), 1)
    lc = collections.LineCollection(
        list(data),
        linewidth=args.linewidth,
        color="lightgrey",
        zorder=-10,
        label="assignment",
    )
    ax.add_collection(lc)

    ax.set(title=f"Wasserstein {args.power}-distance: {distance:g}")

    ax.legend()

    rcParams["savefig.dpi"] = 300
    plt.show()


if __name__ == "__main__":
    main()

This in turn can be used to estimate the Wasserstein distance between arbitrary distributions, by taking the Wasserstein distance between large batches of samples from each.

[vroulet]

Yes, that's clearly worth it :)
Ideally we would like your effort to be adopted so such comments are good.
Also, I was wondering, your implementation is neat, but aren't there other implementations potentially more tpu friendly? (not very important, but one may mention it if it's the case? I'm not an expert on that, packages in optimal transport may have some clever tricks).

[carlosgmartin]

@vroulet Right, there are faster specialized algorithms for approximating the Wasserstein distance (see OTT-JAX).

Perhaps I can append the example above to the LAP gallery entry?

[JTT94]

Not a Contribution, but how does this compare to hungarian_tpu_matcher? from
https://github.com/google-research/scenic/blob/main/scenic/model_lib/matchers/hungarian_jax.py

[carlosgmartin]

@JTT94 I ran a comparison. Interestingly, it looks like that implementation is both faster and yields a smaller jaxpr than the current optax implementation.

Example:

running on a batch of cost matrices of size (100000, 20, 21)
lax.map runtime for optax version: 47.6521
lax.map runtime for scenic version: 7.84279
vmap runtime for optax version: 37.7756
vmap runtime for scenic version: 17.987

This was on a CPU (Apple M3 Max). I encourage someone to try the comparison on a GPU.

@fabianp @vroulet Thoughts? If this implementation is indeed better, I could create a PR that replaces the current optax implementation with it.

[vroulet]

Hello @carlosgmartin,
Yes that would be great. Your implementation can be kept for test if you want. I think your implementation is the standard one so it's good to have but ideally we want the fastest implementation on accelerators.

[carlosgmartin]

@vroulet Submitted a PR: #1140.