JAX implementation

[simeon-ned]

Hello,

Firstly, I would like to extend my gratitude for the impressive work you have presented in your recent paper and the accompanying repository.

I watched the accompanying video from your ICRA presentation, where you mentioned potential JAX support for your solver. I browsed through the repository in search of this feature but was unable to locate any related implementations. Could you please confirm whether JAX support is currently available? If it is not yet implemented, are there any plans to integrate this feature in the near future?

The addition of JAX support could be greatly beneficial to the community, particularly considering the recent updates to Mujoco (now featuring MJX, which integrates with JAX).

If you are considering this feature but need additional resources or contributions, I would be eager to assist in any way possible.

[johnzhang3]

thanks for your interest!

we played around with JAX implementations but never had the bandwidth to do a release. community support is always welcomed in that area! I also recently found out about the JAX implementation of osqp that might be worth checking out for more general use cases.