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