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
5,166 Commits 508 Branches 263 Tags
Commit Graph

142 Commits

Author SHA1 Message Date
Qiao Zhang
82af356b4c
Bump TF hash to get an upstream LLVM GCC fix. (#4251) 2020-09-10 10:10:33 -07:00
Qiao Zhang
a14133aa33
Update TF dep to a passing commit hash. (#4239) 2020-09-09 16:36:30 -07:00
Qiao Zhang
4600dd7957
Update jaxlib version for dlpack fix. (#4231) 2020-09-08 17:20:48 -07:00
Jean-Baptiste Lespiau
bdd65453b4
Add more features to the C++ jax.jit. (#4169) 
This mainly follows https://github.com/google/jax/pull/4089 by adding:

- support for disable_jit from C++
- support for jax._cpp_jit on methods.
- supporting applying @jax.jit on top-level functions, by delaying the retrieval of the device and backend.
- concurrency support.

I am not aware of any feature missing (but I suspect there are still some differences due to the differences between xla_computation and _xla_callable.)

See:

- https://i.ibb.co/ZMvZ4nK/benchmark.png for the benchmarking comparison (see
 cr/328899906 + benchmarks for how numbers were generated)
- The results of the Jax tests when enabling this:
http://sponge2/4a67d132-209f-45c5-ab7b-83716d329ec2 (110 fails, 92 passes, but many common cause of failure).
 2020-09-01 10:34:47 +03:00
Jean-Baptiste Lespiau
2ab6b42a45
Use pytree defined in tensorflow. (#4087) 
It also adds some tests on the scalar C++ conversion.
 2020-08-18 08:58:43 +03:00
George Necula
c7aff1da06
Revert "Use pytree from xla_client. (#4063)" (#4081) 
This reverts commit d8de6b61411179dcd2f63d7639bbcd69b30ac15f.

Tryting to revert because it seems that this produces test
failures in Google.
 2020-08-17 12:53:18 +03:00
Jean-Baptiste Lespiau
d8de6b6141
Use pytree from xla_client. (#4063) 2020-08-14 11:44:03 -04:00
Peter Hawkins
90b353260a
Update XLA. (#3834) 2020-07-23 12:10:39 -04:00
Jake Vanderplas
17c21a79c3
update jaxlib version and changelog for pypi (#3824) 2020-07-22 15:21:30 -07:00
Peter Hawkins
c638264479
Update XLA. (#3671) 2020-07-06 14:35:10 -04:00
George Necula
bfe8f8732b
Change the version of tensorflow, the one in jaxlib 0.1.51 does not build (#3637) 2020-07-02 12:47:22 +03:00
George Necula
a5ed161550
Release jaxlib 0.1.51 (#3636) 2020-07-02 10:33:00 +03:00
Peter Hawkins
47ac612214
Update XLA. (#3623) 2020-07-01 11:14:34 -04:00
Peter Hawkins
2a6fc316c3
Update XLA in preparation for a new jaxlib release (0.1.50). (#3560) 2020-06-25 15:12:47 -04:00
Skye Wanderman-Milne
a088c023ff
Bump jaxlib version to 0.1.49 and update WORKSPACE (#3495) 2020-06-19 13:27:12 -07:00
Peter Hawkins
dd040de18d
Bump XLA version. (#3424) 
Update jaxlib release notes.
 2020-06-12 14:46:45 -04:00
Peter Hawkins
a06b122e4a
Add support for 64-bit FFTs. (#3290) 2020-06-02 09:41:44 -04:00
George Necula
e1cb032456
Prepare version 0.1.47 for jaxlib (#3008) 2020-05-08 10:01:28 +03:00
Peter Hawkins
0534b6595c
Update XLA. (#2977) 2020-05-05 20:47:40 -04:00
Peter Hawkins
4d236b5c47
Update XLA to fix build failures. (#2950) 2020-05-04 09:17:07 -04:00
Peter Hawkins
a18257860a
Update XLA. (#2932) 
Mention illegal instruction fix in changelog.
 2020-05-02 12:47:07 -04:00
Peter Hawkins
ee38e1b3b6
Update XLA. (#2929) 
Includes a fix that may help with issue #2906.
 2020-05-02 11:09:21 -04:00
Peter Hawkins
9802d7321c
Update XLA. (#2927) 2020-05-01 21:08:56 -04:00
Skye Wanderman-Milne
4703b0fa3c
Bump jaxlib version to 0.1.45 and update WORKSPACE and CHANGELOG. (#2785) 2020-04-21 12:01:02 -07:00
Peter Hawkins
14618d1af4
Update XLA. (#2733) 2020-04-15 20:45:20 -04:00
Peter Hawkins
5ba9b6b550
Update XLA. (#2703) 2020-04-13 20:05:33 -04:00
Peter Hawkins
04de8339a8
Update XLA. (#2699) 2020-04-13 15:19:50 -04:00
Peter Hawkins
e904780c71
Update XLA. (#2668) 2020-04-10 08:37:32 -04:00
Peter Hawkins
acb241468d
Update XLA to fix CUDA 9.2 build problem. (#2547) 2020-03-30 19:43:23 -04:00
Peter Hawkins
0c7e134cf4
Update XLA and increment jaxlib version to 0.1.43. (#2546) 2020-03-30 16:56:54 -04:00
Peter Hawkins
cecfb37e6c
Increment jaxlib version to 0.1.42. (#2457) 
Update XLA.
 2020-03-19 09:57:11 -04:00
Peter Hawkins
388e78f8e5
Increment jaxlib version to 0.1.41. (#2428) 
Update XLA.
 2020-03-16 11:24:02 -04:00
Peter Hawkins
7243ddf05c
Bump XLA version. (#2353) 2020-03-04 16:23:48 -05:00
Peter Hawkins
9d61d4b616
Bump jaxlib version to 0.1.40. (#2352) 
Update XLA version.
 2020-03-04 14:49:11 -05: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
Skye Wanderman-Milne
0a78c8c301 Update WORKSPACE to pick up TF grpc upgrade and cleanup. 2020-02-20 09:09:32 -08:00
Skye Wanderman-Milne
cfb5666ac1
Update WORKSPACE to setup upstream dependencies. (#2267) 
This is necessary as of f396035891.

Many thanks to @lezh for this fix!
 2020-02-19 14:58:26 -08:00
Peter Hawkins
9b36238003
Fix sha256 sum for XLA release. (#2230) 2020-02-14 15:33:20 -05:00
Peter Hawkins
71adf8cf7c
Update XLA. (#2229) 
.
 2020-02-14 15:20:06 -05:00
Skye Wanderman-Milne
2ff7019923 Update WORKSPACE 2020-02-13 18:09:53 -08:00
Skye Wanderman-Milne
89edd5eaf0 Update WORKSPACE 2020-02-13 16:05:13 -08:00
Skye Wanderman-Milne
c6b65c0a11 Bump jaxlib version to 0.1.39 and update WORKSPACE. 2020-02-11 18:58:28 -08:00
Peter Hawkins
843e22dd17
Support __cuda_array_interface__ on JAX DeviceArrays. (#2133) 
Allows exporting GPU device-resident arrays to other libraries, e.g., CuPy.
 2020-01-31 10:09:40 -05:00
Peter Hawkins
511d33983a
Initial implementation of DLPack support. (#2123) 
* Initial implementation of DLPack support.

Unfortunately there are still a few bugs in the jaxlib DLPack support, so this code won't be ready to use until jaxlib 0.1.39.

* Fix test failures.

* Update XLA.

Fix failing torch test.
 2020-01-29 23:19:14 -05:00
Skye Wanderman-Milne
6aaf257d8a Update WORKSPACE 2020-01-28 18:04:52 -08:00
Skye Wanderman-Milne
393938f387 Update WORKSPACE to TF commit that builds with build_jaxlib_wheels.sh 2020-01-23 16:35:35 -08:00
Skye Wanderman-Milne
f04348ed53 Bump jaxlib version to 0.1.38 and update WORKSPACE. 2020-01-21 16:59:27 -08:00
Peter Hawkins
653001aa64
Update references to bazel repositories in WORKSPACE to match TF head. (#2005) 2020-01-15 10:51:46 -05:00
Peter Hawkins
64bf55dc6f
Update XLA. (#1997) 
Drop six dependency from jaxlib, since xla_client.py no longer uses six.
 2020-01-14 11:05:54 -05:00
Peter Hawkins
94203bf022
Update XLA. (#1837) 
Update jaxlib BUILD for ead06270dc
 2019-12-10 11:25:09 -05:00