Add missing projections from jaxopt

[carlosgmartin]

Related: #977

The following projections are present in jaxopt but missing in optax:

• sparse_simplex
• hyperplane (now available here)
• halfspace (now available here)
• affine_set
• polyhedron
• box_section
• transport
• birkhoff

I can work on some of these.

[aymuos15]

I am interested in making contributions (and learning as I have only recently started becoming adept with jax in general) to this.

So could someone potentially outline the skeleton for the best way to approach this?

Thanks!

[carlosgmartin]

@aymuos15 To see how Optax currently implements its projection functions, see the projections module.

To see how JAXopt implements the ones in the above list, you can click on one of them (e.g. projection_sparse_simplex) and then click on the button labeled [source], which will take you to the source code.

Note that I'm already adding a few at #1351, so you can work on some of the remaining ones, if you'd like.

[aymuos15]

Thank you so much @carlosgmartin .

I will get on with this my time morning.

[aymuos15]

Two questions.

I have started working on transport and birkhoff. - https://github.com/aymuos15/optax/tree/more_projections

1. Do I create a separate PR, or PR into your fork?
2. Dealing with non existent dependencies present in jaxopt but not optax
   • Affine_set is dependent on EqualityConstrainedQP
   • polyhedron is dependent on OSQP
   • box is dependent on Bisection

Do I do the implementations of these or how do I carry forward? Thanks.

[carlosgmartin]

I suggest creating a separate PR on top of the current head and then merging after #1351 is merged (or waiting until #1351 is merged and then creating a PR on top of the new head).

You're right that some of these projections are dependent on JAXopt-specific solvers that aren't (yet) present in Optax. I'd suggest holding off on those for now, until Optax decides what to do about those solvers (see #977). Feel free to also chime into that discussion.

[aymuos15]

Perfect! Thanks a lot. Ill go through the solver issue and see if I can add anything to it.

I will just wait for your pr to be merged then. I think its much easier/cleaner that way.

[carlosgmartin]

@aymuos15 #1351 and #1439 have finally been merged.

Are you still interested in taking this on?

[aymuos15]

Yes definitely. Thank you for the message.

I have kept subscribed to the two. I am away from my machine till 3rd Jan and will complete this on the 4th if that's okay.

[carlosgmartin]

Of course!

[aymuos15]

Thank you!

[anirgalgali]

Hi @carlosgmartin ,

I would like to contribute to this. Have been using the JAX ecosystem for over a year now and would like to contribute in anyway possible.

Are there any from the original list that haven't yet been implemented? I see that sparse_simplex, transport and birkhoff where added recently?

Thanks!

[carlosgmartin]

Hi @anirgalgali. @aymuos15 is working on those at #1545, but the others are open.