Switches XLA Python bindings to use pybind11.
Update XRT support to point to newer XRT client.
Update minimum bazel version to 0.24.0.
Fix missing backend argument to XLA Compile() calls.
Updates XLA to 00afc7bb81.
The new XLA release removes the use of protocol buffers from the XLA client. Fixes#349.
Add backward compatibility shims to jaxlib to allow older jax releases to still work on an up to date jaxlib.
The new XLA release also incorporates a fix that avoids a host-device copy for every iteration of a `lax.fori_loop()` on GPU. Fixes#402.
Add a new jaxlib.__version__ field, change jax/jaxlib compatibility logic to check for it.
Add a new build option --enable_mkl_dnn that enables MKLDNN contraction kernels in XLA. This leads to significant performance improvements for XLA's dot operator. Enable MKL-DNN by default.
Update XLA version to include MKL-DNN build fix.
Also add a new --enable_march_native build option that turns on -march=native. This is unlikely to have a significant performance impact since XLA JIT-compiles most of its code. Leaving this off by default because it also generates code unlikely to run across a wide selection of architectures and so is unsuitable for building pip wheels.
Significant changes:
* Mac OS X support.
* build script is in Python, not shell.
* build configuration is passed via flags, not environment variables.
* build script configures TF itself, and does not require explicitly checking out the TF git repository and running its configure script. Changes the TF dependency in the Bazel workspace to be an http_archive(), rather than a local checkout of TF.
* rather than trying to guess the path for Bazel-generated XLA artifacts, use a sh_binary() to perform installation of the built artifacts in to the JAX source tree. Bazel's runfiles mechanism is the supported route to find build artifacts.
* downloads Bazel in Python and checks its SHA256 before running it, rather than running an untrusted binary from the internet.
* intentionally does not delete the Bazel cache or Bazel after building.
Example of new build interaction:
Building without CUDA on Mac or Linux:
$ cd jax
$ python3 build.py (or python2 build.py if you want a Python 2 build)
_ _ __ __
| | / \ \ \/ /
_ | |/ _ \ \ /
| |_| / ___ \ / \
\___/_/ \_\/_/\_\
Starting local Bazel server and connecting to it...
Bazel binary path: /Users/xyz/bin/bazel
Python binary path: /Library/Frameworks/Python.framework/Versions/3.7/bin/python3
CUDA enabled: no
Building XLA and installing it in the JAX source tree...
...
Example of building with CUDA enabled on Linux:
$ python3 build.py --enable_cuda --cudnn_path=/usr/lib/x86_64-linux-gnu/
... as before, except ...
CUDA enabled: yes
CUDA toolkit path: /usr/local/cuda
CUDNN library path: /usr/lib/x86_64-linux-gnu/
...
PiperOrigin-RevId: 222868835
