mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-22 08:06:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
7,970 Commits 508 Branches 263 Tags
Commit Graph

191 Commits

Author SHA1 Message Date
Skye Wanderman-Milne
c5e199a42d Update WORKSPACE for jaxlib 0.1.67 2021-05-17 12:15:54 -07:00
Skye Wanderman-Milne
45cbfd6154 Update WORKSPACE to include tpu_driver fix 2021-05-10 16:17:01 -07:00
Qiao Zhang
ce5840969a Update XLA. 2021-05-10 13:15:26 -07:00
Qiao Zhang
8e45a9842f Update XLA. 2021-05-07 10:29:11 -07:00
Qiao Zhang
a0f0586383 Update XLA. 2021-05-06 15:35:10 -07:00
Qiao Zhang
c9d40a45b0 Update XLA. 2021-05-05 19:22:14 -07:00
Qiao Zhang
c2e2953a9c Update XLA. 2021-05-05 15:02:49 -07:00
Qiao Zhang
984f2c634d Fix XLA commit sha256. 2021-05-03 17:55:15 -07:00
Qiao Zhang
24d126b6a9 Update XLA. 2021-05-03 17:20:15 -07:00
Peter Hawkins
efe4a1f920 Update XLA. 2021-05-03 13:38:12 -04:00
Peter Hawkins
cacc5afc4a Update XLA. 2021-05-01 17:05:14 -04:00
Peter Hawkins
a0c96b5ca5 Fix cusparse kernel build. 
The build_wheel.py script was copying the wrong module.

In addition the CUDA stubs from the TF repo were missing a number of cusparse symbols. The updated TF includes the correct stubs.
 2021-04-27 20:28:30 -04:00
Peter Hawkins
79a7f7bca8 Don't build CUDA 11.2 wheels. 
Update XLA.

CUDA 11.1 wheels are compatible with CUDA versions 11.1+, since NVidia now promises enhanced version compatibility between CUDA minor releases starting with CUDA 11.1
 2021-04-26 09:43:29 -04:00
Peter Hawkins
8654a9816f Update XLA. 2021-04-21 09:51:29 -04:00
Peter Hawkins
b37158f9ce Update XLA version. 2021-04-13 16:50:34 -04:00
Skye Wanderman-Milne
55f8c398c9 Update WORKSPACE for jaxlib 0.1.65 release 2021-04-05 14:39:18 -07:00
Skye Wanderman-Milne
6965ab29e9 Update WORKSPACE for profiler fix 2021-03-18 08:41:26 -07:00
Skye Wanderman-Milne
7ad51264fd Update WORKSPACE for jaxlib 0.1.63 release, take 2 
PiperOrigin-RevId: 363315929
 2021-03-16 18:26:52 -07:00
Skye Wanderman-Milne
b319d23431 Update WORKSPACE for jaxlib 0.1.63 release 2021-03-16 09:43:22 -07:00
Skye Wanderman-Milne
482c933db7 Update WORKSPACE for jaxlib 0.1.62 release 2021-03-08 10:47:57 -08:00
Skye Wanderman-Milne
860951b5db Update WORKSPACE to move past build breakage. 
Something was wrong earlier that appears to be fixed now. Fixes #5901.
 2021-03-02 17:49:44 -08:00
Peter Hawkins
b131038094 Update XLA. 2021-02-19 09:55:44 -05:00
Skye Wanderman-Milne
95444d9fae Change TF version for jaxlib-v0.1.61 
This is to line up the new jaxlib with a libtpu nightly build.
 2021-02-11 09:01:35 -08:00
Skye Wanderman-Milne
3a50788ec5 Update WORKSPACE to latest TF commit for jaxlib 0.1.61 release. 2021-02-10 10:05:56 -08:00
Peter Hawkins
83f4069d00 Update XLA. 
Fixes #4687.
 2021-02-08 09:28:14 -05:00
Peter Hawkins
3829bb9cbb Release jaxlib v0.1.60. 2021-02-03 12:00:33 -05:00
Peter Hawkins
fef73075b9 Update XLA. 
Includes 53ab21aa86, which fixes a number of test failures under NumPy 1.20.0.
 2021-02-01 16:42:29 -05:00
Peter Hawkins
f78cefb01b Update WORKSPACE file to comply with TensorFlow's downstream WORKSPACE requirements. 
See:
https://groups.google.com/a/tensorflow.org/g/developers/c/F8-6HPanpnw
 2021-01-15 16:22:48 -05:00
Skye Wanderman-Milne
62b864a654 Update WORKSPACE and jaxlib version to 0.1.59 2021-01-14 13:38:17 -08:00
Peter Hawkins
3416d2082c Update XLA. 2021-01-14 14:17:26 -05:00
Skye Wanderman-Milne
345267bef1 Update TF version for jaxlib 0.1.58 release 2021-01-11 09:39:36 -08:00
Peter Hawkins
7d78d1e06c Update XLA. 2021-01-11 10:13:34 -05:00
Peter Hawkins
1eaf533a28 Update XLA version. 2021-01-05 15:20:35 -05:00
Adam Paszke
f3bfdf8968 Expose is_leaf predicate for pytree.flatten 
and add tests for it. The change has already been landed in the TF code,
where the C++ pytree components live. This is why I needed to bump the
commit.
 2020-12-11 11:26:18 +00:00
Jean-Baptiste Lespiau
e7d6d97b92 Upgrade the TF version. 
Fixes #5135.
 2020-12-08 16:31:54 +01:00
Peter Hawkins
073f13296a Update XLA and fix CI breakage at TF head. 
It seems we must now explicitly load tf_toolchains in our WORKSPACE file.
 2020-12-01 21:44:08 -05:00
Peter Hawkins
409ee76208 Update XLA. 2020-11-30 10:42:07 -05:00
Peter Hawkins
ba84168c23 Update XLA. 2020-11-21 11:42:32 -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
1e7bf5547d Update XLA. 2020-11-16 12:15:34 -05:00
Peter Hawkins
2c6f932e0e Add Python 3.9 support to jaxlib build. 2020-11-09 16:25:58 -05:00
Peter Hawkins
a4669370b9 Update XLA. 2020-11-09 12:07:56 -05:00
Peter Hawkins
608cfde949 Update XLA. 
Update changelog.
 2020-11-04 14:15:33 -05:00
Jean-Baptiste Lespiau
dad8377b78 Enable C++ jit again. 
PiperOrigin-RevId: 339419691
 2020-10-28 02:21:43 -07:00
Peter Hawkins
f58f1ee456 [JAX] Use PocketFFT for FFTs on CPU instead of Eigen. 
PocketFFT is the same FFT library used by NumPy (although we are using the C++ variant rather than the C variant.)

For the benchmark in #2952 on my workstation:

Before:
```
907.3490574884647
max:     4.362646594533903e-08
mean:    6.237288307614869e-09
min:     0.0
numpy fft execution time [ms]:   37.088446617126465
jax fft execution time [ms]:     74.93342399597168
```

After:
```
907.3490574884647
max:     1.9057386696477137e-12
mean:    3.9326737908882566e-13
min:     0.0
numpy fft execution time [ms]:   37.756404876708984
jax fft execution time [ms]:     28.128278255462646
```

Fixes https://github.com/google/jax/issues/2952

PiperOrigin-RevId: 338743753
 2020-10-23 14:20:32 -07:00
Peter Hawkins
29a83ddd70 Update XLA. 
Incorporates fix for #4651.

Fixes #4651.
 2020-10-20 14:11:47 -04:00
Skye Wanderman-Milne
b8cac03ca0 Update XLA in WORKSPACE 2020-10-12 11:57:59 -07:00
Peter Hawkins
72be97c861 Update XLA. 2020-10-08 14:59:06 -04:00
Peter Hawkins
49d851d2e1 Update XLA. 2020-10-05 16:14:32 -04:00
Qiao Zhang
82af356b4c
Bump TF hash to get an upstream LLVM GCC fix. (#4251) 2020-09-10 10:10:33 -07:00