Add support for Jax on Windows

[dwang55]

I failed to build jax from source with(similar with #253, but my system is win):
Traceback (most recent call last):
File "build/build.py", line 308, in
main()
File "build/build.py", line 304, in main
[":install_xla_in_source_tree", os.getcwd()])
File "build/build.py", line 50, in shell
output = subprocess.check_output(cmd)
File "C:\Users\wangd\AppData\Local\Programs\Python\Python36\lib\subprocess.py", line 336, in check_output
**kwargs).stdout
File "C:\Users\wangd\AppData\Local\Programs\Python\Python36\lib\subprocess.py", line 418, in run
output=stdout, stderr=stderr)
subprocess.CalledProcessError: Command '['C:\bazel\bazel.EXE', 'run', '--verbose_failures=true', '--config=mkl_open_source_only', ':install_xla_in_source_tree', 'C:\Users\wangd\jax\build']' returned non-zero exit status 2.

[mattjj]

Thanks for raising this.

JAX doesn’t support Windows right now, though it may be that the only thing standing in our way is getting the build to work. If Tensorflow can be built with XLA on Windows, then it should be possible to build JAX too, with some revisions to our own build scripts to fix issues like this.

Getting JAX on Windows is out of scope for us for now, just because we need to prioritize other kinds of progress. But contributions are welcome!

[mattjj]

On second thought, let’s leave this issue open with a “contributions welcome” tag so that it’s visible to others.

[dwang55]

Thanks for the response! I'm following the instruction from Eric Jang's MAML tutorial. Maybe it's helpful to illustrate it in README for beginners on machine learning like me.

[Prasitt]

7.2

[christopher-hesse]

Looks like XLA does not compile on Windows as of version v2.0.0-alpha and v1.13.1: tensorflow/tensorflow#25773

[christopher-hesse]

That issue has been fixed, so it's possible that jax compiles on windows now. In the meantime, I installed the CPU only version of jax using Windows Subsystem for Linux and it seems to work fine there.

[peterdsharpe]

Confirming @cclauss 's report - I was also able to successfully build the CPU-only version on Windows Subsystem for Linux. If you use PyCharm as your IDE, you can even link it to your Windows version of PyCharm.

[RobeeF]

Hello JAX Team,

Is developing a Windows version of JAX now on the agenda or do you still have other more urgent developments ?

Thanks in advance !

[mattjj]

Thanks for asking!

Windows support is still not on the agenda. We're maxed out on other things, and moreover no one on the JAX team is a Windows user, which only makes development harder.

That's why this issue remains as a "contributions welcome" item, especially from people who know how to develop on Windows!

[Artem-B]

There's been a fair amount of progress making XLA available on windows, which is a prerequisite for JAX. It will be available in TF 2.2 but the nightlies and 2.2rc builds should already have it enabled. I have no clue whether JAX will need additional work, but it would be mostly on JAX side. I would not be surprised if it would just build&work now.

[mattjj]

Thanks for pointing that out, @Artem-B !

JAX is almost pure Python and so I'm optimistic it'll work well once XLA is available.

[RobeeF]

Thanks a lot for your very quick responses!
I will stay tuned.

[ImNotMyself]

Hello JAX Team

TensorFlow 2.2.0 has been released, and XLA is marked available on Windows in its release note. May I know if there’s any plan or progress on Windows version of JAX ?

[agktogether]

@mattjj JAX is writting in python and c++ ,so it gotta not be problem for being compatible for windows
both them have your compiler and independent of os layer??
your architecture what is mage?

[mattjj]

I think the latest attempt to build on Windows is in #2582. Maybe the TF 2.2 release makes it even easier!

In any case, we don't have the bandwidth on the JAX core team to work on a Windows build (none of us have Windows machines either!), but contributions here are more than welcome.

[agktogether]

@mattjj ubuntu or Linux like systems are not perfect for GUI Computing and coding but windows although is very weak and stupid in developing, you can use other OS for generating software for another OS
it is possible...any way all of us is its architecture not me

[michaelosthege]

Unfortunately WSL2 is not the ultimate solution: Accessibility for many Python users is terrible and installing WSL2 installation requires admin permissions, which for good reasons are not given by default in many companies.

I'd offer to help, but I have no clue about compiler settings..

[enojoker]

It seems that the build script is the only thing stopping people from using JAX on windows now. Could you guys look into it? There are many use cases in the windows system.

[hawkinsp]

@cloudhan submitted PR #4843 which allows JAX to build and run on Windows. There are still a few things for us to do:

• not all tests pass (PRs welcome!)
• we will need to set up CI testing and include wheel builds in our releases.

But this is wonderful progress! Thanks @cloudhan !

[elisim]

Hi, any progress of supporting JAX for windows?
Alternatively, please consider changing the error msg to something for informative:

ModuleNotFoundError: jax requires jaxlib to be installed. See https://github.com/google/jax#installation for installation instructions.

[javrtg]

As a begginer myself with wsl2 and linux, is there a sort of step-by-step tutorial that shows how to get jax up and running?
Concretely, I am interested in using jax in a conda environment (idk if this is possible at the moment),

btw amazing work!

[tanujdhiman]

Hello everybody !!
Is it now working in windows ? with pip install jaxlib.
Thanks

[cloudhan]

@tanujdhiman just try https://github.com/cloudhan/jax-windows-builder

[sew2022]

Could not find a version that satisfies the requirement jaxlib (from versions: none)

[FurkanGozukara]

this is beyond unacceptable

single person able to provide windows version but you can't as google?

https://github.com/cloudhan/jax-windows-builder

[r3z8]

Thanks for raising this.

JAX doesn’t support Windows right now, though it may be that the only thing standing in our way is getting the build to work. If Tensorflow can be built with XLA on Windows, then it should be possible to build JAX too, with some revisions to our own build scripts to fix issues like this.

Getting JAX on Windows is out of scope for us for now, just because we need to prioritize other kinds of progress. But contributions are welcome!

Hi, Hope all is well for you
Is this still the case? I've been dealing with Jax not being supported on windows and I was wondering if the support is going to be around any time soon? or is it in your schedule for future?

I should add that I need the jaxlib=0.3.15 for windows while it is only available for linux at the moment.

Thanks in advance

[hawkinsp]

@r3z8 JAX does make Windows CPU releases these days. But we only started doing that during the 0.4 series and we won't be going back and retroactively making Windows releases for older JAX versions, I'm afraid.

I think your best bet, if you need on old JAX version on Windows, would be to use WSL2 and install the Linux version. But going forward we're doing our best to support Windows on CPU (with help from the community).

[FurkanGozukara]

@r3z8 JAX does make Windows CPU releases these days. But we only started doing that during the 0.4 series and we won't be going back and retroactively making Windows releases for older JAX versions, I'm afraid.

I think your best bet, if you need on old JAX version on Windows, would be to use WSL2 and install the Linux version. But going forward we're doing our best to support Windows on CPU (with help from the community).

this is great news

but what about Jax GPU for Windows?