Add sparse simplex, transport, and Birkhoff projections from JAXopt

[aymuos15]

I have set it up so one can feed the solver for Transport/Birkhoff projections via make_solver argument. @carlosgmartin for this PR is this enough (as the tests are namesake for these) or should I only port the sparse_simplex for now?

Continuation of #1280, #1351 and #1439

[carlosgmartin]

Thanks!

I'd ask one of the project maintainers, perhaps @rdyro, about the desired API.

[rdyro]

Thanks!

[vroulet]

It looks pretty good to me. Just to be sure @aymuos15 :

• what changes did you do compared to jaxopt? None?
• did you take all tests?
• did you check if you kept optax formatting of e.g. docstrings? I think references are not well formatted (compare to optimizers). You can build the docs locally to see if the docs look good.

Thanks again!

[aymuos15]

I did just push a fix for the docstrings. Thank you very much for pointing that out.

@vroulet For simplex, I have made sure we are matching everything. For transport and birkhoff, they are usable with an outside solver.

[vroulet]

Ok, would it be possible to make two separate PRs? One for the sparse simplex projection and one for the transport/birkhoff projections.
This one can be the sparse simplex.

The problem with the transport/birkhoff is that (i) it requires a solver, (ii) the api for the solver is not documented and it takes a specific form like solver.run(...).params. Such api choices need to be documented. If we add these projections I would actually use a default solver rather than letting the user select. Right now, there is no "solver" in optax that could fit the requirements

[aymuos15]

Okay, that's a very fair argument. Thank you. I will change this to just the sparse simplex. Instead of raising another PR right away, should I first raise a separate issue to discuss what the default solver should be (or continue the convo in one of the other issues)?

[vroulet]

@aymuos15 you may raise the issue and open a discussion with maybe @carlosgmartin if he's interested.

For now, just prune this PR from the Birkhoff and transport projections. I'm happy to merge after that.

Thanks again!