[emu]

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.

[marmaduke]

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

[emu]

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.

[marmaduke]

Very 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.