google-deepmind
Feature Request: Auction algorithm for linear assignment
Feature request
Add an implementation of the auction algorithm for the linear assignment problem to Optax, alongside the existing Hungarian algorithm.
Motivation
The auction algorithm is a simple, well-studied alternative to the Hungarian method that:
- scales well to large problems,
- is naturally parallelisable and fits JAX’s accelerator model,
- is widely used in optimisation, vision and control.
Having it in Optax would let users choose between two standard LAP solvers within a single JAX-native API.
References
-
Bertsekas, D. P. (1979). “A distributed algorithm for the assignment problem.” IEEE Trans. Autom. Control.
https://web.mit.edu/dimitrib/www/Bertsekas_Auction_Distributed_1979.pdf -
Bertsekas, D. P. (1981). “A new algorithm for the assignment problem.” Mathematical Programming.
https://link.springer.com/article/10.1007/BF02186476
(These papers form the standard reference for the auction algorithm.)
Heyo @rajasekharporeddy and @surajyadav-research I’d love to take up this feature request if it’s available.
I’ve worked with Python + optimization algorithms before, and I’ve started going through the Bertsekas papers to understand how the auction algorithm could fit into Optax’s current LAP API. The parallel-friendly nature of the algorithm seems like a great match for JAX, so I’m excited about exploring it.
If there’s no objection, I’m happy to:
put together a small design / API outline for feedback, implement the auction solver with JAX primitives, add tests + benchmarks, and update the docs so users understand when to choose auction vs Hungarian.
If this sounds good, feel free to assign the issue to me — I’m ready to get started. Thanks!
Hi @SanchitKS12 , I was also working on this. We can collaborate on it.
Heyo @surajyadav-research Thanks for replying — that sounds great I’d be happy to collaborate on this.
To make sure we’re aligned, could you share how far you’ve progressed and what direction you had in mind for the implementation? I’ve already started reviewing the Bertsekas papers and thinking about how we can match the existing Optax LAP API, so I’d love to sync our approaches.
Also, what’s the best way for us to coordinate — GitHub discussion here, a draft design doc, or maybe Discord? I’m fine with any.
Looking forward to working together on this!
Hi @SanchitKS12 , I’m currently working on the algorithm implementation. I’ve taken some reference from the existing Hungarian algorithm implementation and I’m comparing the results and runtime of both. The final code isn’t complete yet, but I’m exploring ways to parallelize certain parts so it can be faster than the Hungarian algorithm.
For communication, we can connect on Discord. Excited to collaborate on this.
Thanks for the update @surajyadav-research that sounds really impressive, especially the part about comparing runtime performance and exploring parallelization. I’d love to collaborate on this. To avoid overlapping work, maybe you can continue focusing on the core implementation and benchmarking, and I can start working on the tests, docs, usage examples and API integration once your structure is stable. We can sync regularly so the work flows smoothly and then eventually work on a PR together.
We can connect on Discord to coordinate quickly — my Discord name is sanchitks12. We can discuss details there, and post major progress updates here and in the Google DeepMind Discord so the maintainers stay in the loop. Really excited to work together on this and push the auction algorithm into Optax!
Thanks alot for the collaboration !!
Hi @SanchitKS12 I’m almost done with the code for the auction algorithm. The running time and optimal cost look good especially on GPU/TPU, where the runtime on large matrices is almost 6–10× faster. I’m currently checking the convergence rules and a few small edge cases. Other than that, everything is final.
I couldn’t find the Google DeepMind Discord. Could you please share the link here? My Discord ID is surajyadav_research.
Great collaborating with you!
Great to hear the update @surajyadav-research
the performance improvement on GPU/TPU sounds promising. Thanks for sharing the progress status.
About Discord: it looks like there isn’t an official public Discord for Optax / Google DeepMind. Most project communication seems to happen through GitHub issues, discussions, and PR threads. So we can collaborate via Discord DMs and keep project-relevant updates here on GitHub.
I’ll DM you shortly — my Discord is sanchitks12. Once you wrap up the remaining convergence and edge-case checks, I’ll move ahead with the next parts:
Integrating the solver into the existing Optax API alongside Hungarian.
Writing unit tests + benchmarks for CPU & accelerator performance.
Preparing documentation + usage examples.
helping with PR structure & review updates once we open it
Glad to collaborate — looking forward to pushing the implementation through to merge.