top | item 37705838

(no title)

emu | 2 years ago

Not true!

We release Windows CPU wheels (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.

discuss

marmaduke|2 years ago

Oops so sorry. But this is recent isn't it? I thought it was actually due to XLA/Bazel not supporting it?

emu|2 years ago

Yes, 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.