We do this to see if this reduces the incidence of errors fetching
the tf-nightly package. These tests are being run when we import
the code in Google.
* Add jax.image.resize.
This is a port of `tf.image.resize()` and the `ScaleAndTranslate` operator.
While I don't expect this implementation to be particularly fast, it is a useful generic implementation to which we can add optimized special cases as the need arises.
* Refactored host_callback to use the C++ runtime.
* The new runtime makes it unnecessary to start the outfeed_receiver
in the user's code
* We don't need msgpack anymore
* There is an interaction between host_callback and using lax.outfeed.
I am trying to solve this by (a) making host_callback_test stop the
outfeed receiver on finish and infeed_test on start, and (b)
telling pytest-xdist to run all the tests from one file into
a single worker.
* Use dynamic loading to locate CUDA libraries in jaxlib.
This should allow jaxlib CUDA wheels to be manylinux2010 compliant.
* Tag CUDA jaxlib wheels as manylinux2010.
Drop support for CUDA 9.2, add support for CUDA 11.0.
* Reorder CUDA imports.
We choose the same set as TensorFlow (minus 3.7, which TF is apparently considering dropping anyway).
This avoids a slow PTX -> SASS compilation on first time startup.
* raise minimum Bazel version to 2.0.0 to match TensorFlow.
* set --experimental_repo_remote_exec since it is required by the TF build.
* bump TF/XLA version.
* use the --config=short_logs trick from TF to suppress build warnings.
* work around https://github.com/bazelbuild/bazel/issues/9254 by setting BAZEL_LINKLIBS=-lstdc++
* drop CUDA 9.0 support, since we use a batched kernel only present in CUDA 9.2 or later.
* drop Python 2.7 support.
This PPA has been removed by the owner: https://launchpad.net/~jonathonf/+archive/ubuntu/python-3.6
This causes `apt-get update` to fail when generating the Docker image. We don't seem to need this repository, so just remove it before calling `apt-get update`.
In principle, JAX should not need a hand-written CUDA kernel for the ThreeFry2x32 algorithm. In practice XLA aggresively inlines, which causes compilation times on GPU blow up when compiling potentially many copies of the PRNG kernel in a program. As a workaround, we add a hand-written CUDA kernel mostly to reduce compilation time.
When XLA becomes smarter about compiling this particular hash function, we should be able to remove the hand-written kernel once again.
It turns out there are implicit CUDA dependencies inside the TF libraries used by JAX, so the attempt to disable GPU dependencies conditionally didn't work.
--action_env variables are passed to every build action. This means that if the variable changes, the entire build cache is invalidated. By contrast, --repo_env variables are only passed to repository rules and don't affect every action. In principle this means that we should be able to rebuild JAX for different Python versions without rebuilding 99% of the C++ code.
Update bazel release for build script to 0.29.1 (same as TensorFlow.)
Previously we lied claimed our wheels were manylinux1 compliant but they weren't.
Uses a cross-compilation toolchain from the TF folks that builds manylinux2010 compliant wheels from a Ubuntu 16.04 VM.
The CUDA wheels still aren't manylinux2010 compliant because they depend on CUDA libraries from the system.
