Jax support

This would broadly involve:

1. Adding a jax directory that follows the same structure as the PyTorch/Tensorflow ones to our main module and appropriately add that layer: https://github.com/cvxgrp/cvxpylayers/tree/master/cvxpylayers
2. Filling in the details of the module and tests
3. Updating the README/image to include Jax
4. Adding some Jax examples: https://github.com/cvxgrp/cvxpylayers/tree/master/examples

\cc @gsp-27 who has expressed interest in helping with this – please let us know if you make any progress here!

[bamos]

I’ve been recently thinking about adding a jax/flax version but haven’t started.

On Feb 17, 2021, at 7:38 AM, bodono [e-mail] wrote:

Are there any updates on this?

— You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub, or unsubscribe.

[eserie]

That looks useful! Here we’d need to define new operations in each of the frameworks with custom derivatives and I don’t immediately see how to do that in the eagerpy docs

Yes sure, its true that eagerpy is more design to go in the other way around: write some code that uses a tensor framework to compute derivatives of it. So not sure it could help here. However, maybe the fact to have lightweighted unified tensor API may help to factorize a bit.

I’ll have a look at the Jax integration in #100!