Open ryancoe opened 1 year ago
@ryancoe @cmichelenstrofer it appears a Windows-compatible pip install is now available for JAX as of v0.4.13. That's one of the major implementation hurdles out of the way!
We should try it at some point. But it does require changing the source a bit, so it won't be a small task.
@michaelcdevin - can you quickly see how well the Jax Windows install works as a first step?
pip install jax
works without a hitch on Windows. I tested some of the basic jax.numpy
and jax.grad
operations and all seems to work as expected.
NumPy v2.0 was released four days ago with various breaking changes. Since autograd
is no longer maintained, this makes it so autograd
is incompatible with current and future releases of NumPy.
It looks like jax
was proactive about maintaining compatibility with NumPy v2.0. I added in a NumPy version limitation in 902c17b as a stopgap so WecOptTool doesn't break, but switching from autograd
to jax
is a higher priority now so we don't fall behind in NumPy versions.
@cmichelenstrofer @ryancoe
Ever since originally adopting
autograd
, we've been concerned that most of the development energy fromautograd
has moved toJAX
. In addition to continued development,JAX
also has more complete functionality (e.g., forfft
).We did not use
JAX
initially because support for MS Windows is not great - users can either compile on their own, use a third-party binary, or use Windows Subsystem for Linux (WSL) (https://github.com/google/jax#installation).Given that more direct
JAX
support for MS Windows does not imminent, we there are two major hurdles preventing us from transitioning:import autograd.numpy as np
toimport jax.numpy as np
, but apparently there is more to it than that