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