Moving from autograd to jax

mirapy-org / mirapy

MiraPy: A Python package for Deep Learning in Astronomy

https://mirapy.org/

MIT License

44 stars 10 forks source link

Moving from autograd to jax #35

Open ericmjl opened 5 years ago

ericmjl commented 5 years ago

The next generation autograd is now called jax, and is built by the same guys who built autograd + more, with a somewhat nicer API. One other attractive feature is the ability to JIT compile functions to CPU, GPU and TPU!

As such, I'd like to propose switching over from autograd to jax in requirements.txt. Naturally, this is just a suggestion; I simply thought I'd check out the repo having seen @mattjj star it.

cc: @mattjj, @dougalm

swapsha96 commented 5 years ago

Thank you @ericmjl for your recommendation! I have begun digging deeper into jax and I see it would be beneficial to switch at this initial development stage of MiraPy. 😄 It seems easy to switch from autograd so why not?!