Add an assignment problem solver

[carlosgmartin]

Feature request: Add a GPU/TPU-friendly solver for the assignment problem. For context, see:

1. scipy.optimize.linear_sum_assignment
2. Implement scipy.optimize.linear_sum_assignment jax-ml/jax#10403
3. Add linear_sum_assignment jax-ml/jax#16974

The last page contains the following comment:

There is a TPU-friendly implementation of the Hungarian algorithm here: https://github.com/google-research/scenic/blob/main/scenic/model_lib/matchers/hungarian_cover.py

Potentially relevant literature:

• A distributed auction algorithm for the assignment problem (2008)
• Bipartite Graph Matching Computation on GPU (2009)
• An anytime assignment algorithm: From local task swapping to global optimality (2013)
• On implementing 2D rectangular assignment algorithms (2016)
• GPU-accelerated Hungarian algorithms for the Linear Assignment Problem (2016)
• Fast block distributed CUDA implementation of the Hungarian algorithm (2019)
• Development of an Accelerating Hungarian Method for Assignment Problems (2021)
• Algorithm 1015: A Fast Scalable Solver for the Dense Linear (Sum) Assignment Problem (2021)
• New auction algorithms for the assignment problem and extensions (2024)

[carlosgmartin]

@vroulet What would be the best place to add a solver function? optax.contrib.solve_linear_assignment_problem? I can create a draft PR with my current implementation.

[vroulet]

Hello @carlosgmartin,
That's a good question. The contrib folder is more for algorithms that did not pass the test of time (fancy new optimizers for deep learning typically). A linear assignment problem has passed the test of time so I think it should not be in the contrib folder. I think the best would be to follow the structure in scipy, i.e. simply making a new folder called assignment (maybe one day someone will also want to add quadratic_assignment).
The most important part here would be to have

1. a clear documentation with doctests,
2. a notebook example to understand how to use it and in which context.
   Given your past record of contributions, I think you'll be able to produce such a clean contribution.
   Then the main issue as pointed out by Jake in Add linear_sum_assignment jax-ml/jax#16974 is the burden of maintenance. Would you commit to help maintaining that function? (open-source means potentially forever maintenance just for you to be aware)

[carlosgmartin]

The linear assignment problem is an important and standard problem in optimization, as you've noted, so I think it's worth adding to optax. To that end, I've started a PR at #1083 and welcome any feedback.

[carlosgmartin]

Closed by #1083.