Will this repo support `jax` framework?

[HiccupRL]

Thanks for your solid work! Since jax has been widely used, will this repo offer a jax version?

[ValerianRey]

Hi! We've been experimenting a bit with jax lately, but making a jax version would require quite a lot of work. In the long term, I would love to have a jax version, but it's not planned at the moment. There are two things that could accelerate things:

• If jax becomes even more popular, or people keep asking for a jax version / upvote this issue.
• If we get much better performance on jax (likely due to torch limitations)

If we do that, it will be in another repo and another library, and we would have to rename the github organization from TorchJD to something more independant of torch.

[PierreQuinton]

I would like to add that if someone wants JD in Jax and has the skills, then they can join the discord server to start working on it with us. In the current state, we are two people working on this project, and we have many more things to do than we can handle.