59 Commits

Author SHA1 Message Date
Peter Hawkins
dacf31f202 Check for NumPy and SciPy versions during jaxlib builds. 2021-05-24 12:39:37 -04:00
Peter Hawkins
86887a21e2 Add Mac M1 support to build.py. 2021-05-18 21:53:31 -04:00
Peter Hawkins
7d4551dd45 Add support for arm64 builds to build.py and build_wheel scripts. 2021-05-10 11:08:09 -04:00
Skye Wanderman-Milne
7b7a2a1ba8 Don't build with Cloud TPU support for GPU or Mac wheels.
This should make our builds simpler and less failure-prone.
2021-03-08 16:19:19 -08:00
Christian Sigg
c0e42e59eb
Switch to new flag to enable CUDA support in TF 2021-03-03 08:55:08 +01:00
Situ Yi
77905ef5fd
Update build.py with Cloud TPU support
Update build.py to have Cloud TPU support by default.
b/180656002
2021-02-19 12:57:57 -08:00
jax authors
15359e018d Merge pull request #5738 from hawkinsp:avx
PiperOrigin-RevId: 357785312
2021-02-16 12:53:41 -08:00
Skye Wanderman-Milne
3caf52d54a Fix build.py version comparison 2021-02-16 11:32:50 -08:00
Peter Hawkins
a27545497b Enable AVX by default in release builds.
TensorFlow also enables AVX by default, so this seems like a reasonable baseline CPU feature to expect.
2021-02-16 14:12:01 -05:00
jax authors
5fd3deb0be Merge pull request #5580 from daskol:search-bazel-in-path
PiperOrigin-RevId: 357755577
2021-02-16 10:49:33 -08:00
Skye Wanderman-Milne
dd513c1585
Small simplification 2021-02-16 08:39:32 -08:00
Daniel Bershatsky
96193cef12 Simplify bazel version check 2021-02-13 16:32:21 +03:00
Skye Wanderman-Milne
74db363830 Bump required bazel version to 3.7.2
Untested, but I'm about to build a new jaxlib and will make any fixes then.
2021-02-11 09:07:04 -08:00
Daniel Bershatsky
c71a4f5871 Try to download bazel as the last resort
Resolution order of paths to bazel binary is as follows.

1. Use --bazel_path command line option.
2. Search bazel binary in PATH environment variable.
3. Download required bazel release.
2021-02-01 15:15:38 +03:00
Clemens Giuliani
330e69dd30 Add support for ROCm 2020-12-06 22:55:53 +01:00
Cloud Han
ea340eed93 use strict action env on windows to avoid constant full rebuilding 2020-11-24 23:44:26 +08:00
Peter Hawkins
c06ead6b04 Change jaxlib build rules to build a wheel, rather than writing output to the source directory. 2020-11-20 11:47:00 -05:00
Peter Hawkins
1fcbd2f083 Add -Wno-stringop-truncation to build flags on Linux.
Works around https://github.com/tensorflow/tensorflow/issues/39467 for gcc 10+ builds.
2020-11-19 22:50:08 -05:00
Cloud Han
a6acce58e0 Build on Windows
1. Build on Windows

2. Fix OverflowError

    When calling `key = random.PRNGKey(0)` OverflowError: Python int too
    large to convert to C long for casting value 4294967295 (0xFFFFFFFF)
    from python int to int32.

3. fix file path in regex of errors_test

4. handle ValueError of os.path.commonpath
2020-11-19 23:33:06 +08:00
Peter Hawkins
b222e5e05f Update bazel version to 3.1.0. 2020-11-17 15:50:30 -05:00
Peter Hawkins
2c6f932e0e Add Python 3.9 support to jaxlib build. 2020-11-09 16:25:58 -05:00
Srijan Saurav
40e20242db
Fix code quality issues (#4302)
Changes:
- Fix unnecessary generator
- Iterate dictionary directly instead of calling .keys()
- Remove global statement at the module level
- Use list() instead of a list comprehension
- Use with statement to open the file
- Merge isinstance calls
2020-09-17 09:21:18 -07:00
Peter Hawkins
a141cc6e8d
Make CUDA wheels manylinux2010 compliant, add CUDA 11, drop CUDA 9.2 (#3555)
* 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.
2020-06-25 14:37:14 -04:00
Peter Hawkins
821193b2b7
Explicitly build specific CUDA capabilities. (#2722)
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.
2020-04-15 10:57:53 -04:00
Peter Hawkins
cbdf9a5a43
Drop support for Python 3.5. (#2445) 2020-03-18 10:54:28 -04:00
Peter Hawkins
219d503e71
Don't show progress bar in build script if output is not a terminal. (#2429) 2020-03-16 11:01:08 -04:00
Peter Hawkins
80abdf0c53
Unbreak build and update XLA. (#2289)
* 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.
2020-02-22 09:45:24 -08:00
Peter Hawkins
b6e8341176
Improve developer documentation. (#2247)
Add Python version test to build.py.
2020-02-17 11:24:03 -08:00
Skye Wanderman-Milne
bf91ebf67a
Return error number in build.py on bad bazel version. (#2218)
This prevents our build scripts from continuing on error.
2020-02-12 09:57:54 -08:00
Skye Wanderman-Milne
96a65de6c8
Try downloading bazel before using pre-installed bazel. (#2217)
This ensures we're using the right bazel version.
2020-02-12 09:49:33 -08:00
Peter Hawkins
fe041c7590 Set minimum Bazel version to 1.2.1. 2020-02-03 10:13:51 -05:00
Ruizhe Zhao
8c7fc3919d
Upgrade bazel from 0.29.1 to 1.2.1 (#2137) 2020-02-03 10:12:40 -05:00
Peter Hawkins
e60d5dd54c
Remove "from __future__" uses from JAX. (#2117)
The future (Python 3) has arrived; no need to request it explicitly.
2020-01-29 12:29:03 -05:00
Skye Wanderman-Milne
b6dfb8bf18
Bump minimum bazel version to 0.26.0. (#2060)
Fixes #2044
2020-01-23 16:46:45 -08:00
Peter Hawkins
ea91c96a9d
Specify a minimum Mac OS version in builds to avoid backward compatibility problems. (#1807) 2019-12-03 11:59:31 -05:00
Peter Hawkins
534d812b57
Add a handwritten ThreeFry2x32 CUDA kernel. (#1756)
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.
2019-11-24 13:06:23 -05:00
Peter Hawkins
67038321f8
Revert support for building a non-GPU build with --config=cuda enabled. (#1757)
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.
2019-11-24 13:06:10 -05:00
Peter Hawkins
9b853a4255
Update XLA. (#1702)
Add support for building a CPU-only jaxlib with a CUDA-enabled toolchain.
2019-11-16 11:01:36 -05:00
Skye Wanderman-Milne
44ccca05df Revert "Update bazel min version to 0.26.0."
It turns out we can't build TF with this bazel version yet.

This reverts commit 807a1958bb2a276f012901ba6cd9226371099005.
2019-11-13 10:05:58 -08:00
Skye Wanderman-Milne
807a1958bb Update bazel min version to 0.26.0.
I think this was the first release including the --repo_env arg:
d7702b16eb
2019-11-13 09:59:38 -08:00
Peter Hawkins
3e9ce2f69f
Use --repo_env instead of --action_env to configure Python and CUDA. (#1619)
--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.)
2019-11-01 10:41:51 -04:00
Peter Hawkins
1428c11a2c Update Jaxlib version to 0.1.29.
Bump XLA version. Enable C++14 mode since it is required by the new XLA version.
2019-09-28 15:11:09 -04:00
Peter Hawkins
fef315b6e6 Add ability to pass extra bazel options to build script.
Remove cublas/cusolver dependencies from Jaxlib python code.
2019-08-08 16:14:45 -04:00
Peter Hawkins
40c517e9ef Add --spawn_strategy from TensorFlow configuration. 2019-08-01 21:43:03 -04:00
Peter Hawkins
fafe69bd8f Relax Bazel maximum version check.
Remove debug print from lax_numpy_test.py.
2019-05-29 19:38:55 -04:00
Peter Hawkins
8cbc475a8b Add "bazel shutdown" to the build script.
Fixes #739.
2019-05-20 20:53:44 -04:00
Peter Hawkins
5599f473f0 Revert to requiring Bazel 0.24 for JAX builds.
It turns out TF does not fully support 0.25 yet.

Add a check for a maximum bazel version to the build script.

Fixes #702.
2019-05-12 09:25:10 -04:00
Peter Hawkins
7f9e1809bb Fix build with Bazel 0.25.
Update minimum Bazel version to 0.25.0.
2019-05-10 11:53:15 -04:00
Peter Hawkins
f60d927df8 Update XLA version to include CUDA version check fix and newer XLA.
Make specification of CUDA environment variables optional.

Fixes #627.
Fixes #276, although the fix requires a new Jaxlib release.
2019-04-25 16:19:53 -07:00
Peter Hawkins
e8638f73cc Update XLA.
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.
2019-03-31 10:56:47 -07:00