4 ms·
Not true! We release Windows CPU wheels (https://pypi.org/project/jaxlib/#files https://pypi.org/project/jaxlib/#files). So JAX on CPU works great on Windows.
by emu 3y ago
Not true!
We release Windows CPU wheels (https://pypi.org/project/jaxlib/#files https://pypi.org/project/jaxlib/#files). So JAX on CPU works great on Windows.
We don't release Windows GPU wheels at the moment, but that's because we're a small team and none of us use Windows personally. We welcome contributions!
(I verified that the Windows CUDA GPU support built as recently as two weeks ago, but I don't have the ability to test that it works.)
We recommend WSL2 because that's just using our existing Linux CUDA release.
- marmaduke 3y agoOops so sorry. But this is recent isn't it? I thought it was actually due to XLA/Bazel not supporting it?
- emu 3y agoYes, we made this more formally supported recently. We felt that Windows CPU support was important so everyone can run JAX, even if it's not always the most-accelerated version of JAX. And we got some great PRs from the community that helped fix a few open issues.
- marmaduke 3y agoVery nice! I just installed, I hope to eventually contribute down the line, especially in terms of custom operators. They weren't even document until recently, and there's still quite some work to add them.