add Hungarian solver from optax #2253

Job	Run time
Fast GPU tests Python 3.10 on ubuntu-22.04	34m 45s
Python 3.10 on ubuntu-latest	1h 19m 5s
Fast tests Python 3.10 jax-default	10m 13s
Python 3.11 on ubuntu-latest	1h 20m 57s
Fast tests Python 3.10 jax-latest	11m 32s
Python 3.12 on ubuntu-latest	1h 23m 29s
Python 3.13 on ubuntu-latest	57s
Python 3.9 on macos-13	1h 28m 12s
Python 3.10 on macos-14	1h 10m 43s
	7h 39m 53s

Provide feedback