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. https://github.com/google/jax/issues/10403
3. https://github.com/google/jax/pull/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:

• GPU-accelerated Hungarian algorithms for the Linear Assignment Problem