mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
21,727 Commits 506 Branches 263 Tags
Commit Graph

240 Commits

Author SHA1 Message Date
Peter Hawkins
945fde41e4 Update minimum Python version to 3.10. 2024-06-26 13:47:14 -04:00
Peter Hawkins
7f24837eef Update minimum NumPy version to v1.24. 2024-06-21 15:17:17 -07:00
Yash Katariya
045ea944c8 Finish jax and jaxlib 0.4.30 release 
PiperOrigin-RevId: 644402635
 2024-06-18 08:56:39 -07:00
Peter Hawkins
3f4f79c83e Prepare for 0.4.30 release. 2024-06-18 07:34:45 -04:00
Peter Hawkins
dcb7b3c3f1 Readd cuda12_pip extra to keep CI users happy. 2024-06-13 12:37:06 -04:00
Peter Hawkins
b13733c13f Update JAX dependencies, extras, and documentation for plugins. 
* Make jaxlib a direct dependency of jax.
* Remove mentions of monolithic CUDA installations from the JAX documentation.
* Drop the cuda12_pip extra and the cudnn version specific extras.
* Add a with_cuda extra to the jax-cuda12-plugin package, use it in jax's setup.py. This allows us to specify cuda extras in one place.
* Make a few small doc improvements.
 2024-06-13 11:36:23 -04:00
Yash Katariya
b1f7627c71 [Rollback] Bumped the minimum ml_dtypes version to 0.4.0 
Reverts e86c436e7f8e4e0546eff8bc2d3756a7c49dc83b

PiperOrigin-RevId: 642741832
 2024-06-12 14:40:40 -07:00
Peter Hawkins
6fa31e59c4 Update version numbers after v0.4.29 release. 2024-06-10 14:37:53 -04:00
Peter Hawkins
e071053d97 Prepare for 0.4.29 release. 2024-06-10 08:54:28 -04:00
Henning Becker
437263659c Fix pip dependency after cuDNN 9 upgrade 
PiperOrigin-RevId: 637959308
 2024-05-28 11:06:48 -07:00
Henning Becker
15a1985445 Update cuDNN to version 9.1.1 in JAX 
PiperOrigin-RevId: 636956696
 2024-05-24 10:10:21 -07:00
jax authors
e86c436e7f Merge pull request #21404 from superbobry:maint 
PiperOrigin-RevId: 636862158
 2024-05-24 03:53:58 -07:00
Parker Schuh
c1f5a32875 Drop refs to australis in build rules and exporter. 
PiperOrigin-RevId: 636720089
 2024-05-23 16:52:58 -07:00
Sergei Lebedev
0a694a1b42 Bumped the minimum ml_dtypes version to 0.4.0 2024-05-23 21:51:00 +01:00
Jake VanderPlas
9ac1d38226 Finish jax and jaxlib 0.4.28 release 
PiperOrigin-RevId: 632653310
 2024-05-10 18:06:52 -07:00
Peter Hawkins
038dfeec15 Prepare 0.4.28 release. 2024-05-09 19:25:33 +00:00
Yash Katariya
5031a1ddc4 Finish jax and jaxlib 0.4.27 release 
PiperOrigin-RevId: 631486157
 2024-05-07 11:14:09 -07:00
Yash Katariya
70b4477296 Start jax and jaxlib 0.4.27 release 
PiperOrigin-RevId: 631409685
 2024-05-07 07:01:24 -07:00
Yash Katariya
24517ca3e0 Finish jax and jaxlib 0.4.26 release 
PiperOrigin-RevId: 621658207
 2024-04-03 15:40:24 -07:00
Peter Hawkins
61493263a9 Prepare for 0.4.26 release. 2024-04-03 14:38:58 -04:00
Peter Hawkins
478cfa9944 Add an upper bound on JAX's CUDNN version constraint. 
Major releases of CUDNN break ABI compatibility, so we cannot allow new major versions.

PiperOrigin-RevId: 620030416
 2024-03-28 13:00:36 -07:00
jax authors
0be07e6aec Remove support for CUDA 11. 
Pin minimal required versions for CUDA to 12.1.

Reverts 910a31d7b7510e3375718ab1ea0d38df7bd2c0d5

PiperOrigin-RevId: 618911489
 2024-03-25 11:46:39 -07:00
jax authors
910a31d7b7 Reverts bed4f65438a62777ed100ecec2b0eb3f7cf87a0e 
PiperOrigin-RevId: 618249855
 2024-03-22 12:10:53 -07:00
jax authors
bed4f65438 Remove support for CUDA 11. 
Pin minimal required versions for CUDA to 12.1.

PiperOrigin-RevId: 618195554
 2024-03-22 09:05:39 -07:00
Yash Katariya
e0fd29082d Finish jax and jaxlib 0.4.25 release 
PiperOrigin-RevId: 610413312
 2024-02-26 08:19:05 -08:00
Yash Katariya
f9e20d5875 Start the release of jax and jaxlib 0.4.25 
PiperOrigin-RevId: 610020972
 2024-02-24 10:25:17 -08:00
Yash Katariya
f12550964d Update the cuda 12 dependencies since we upgraded to cuda 12.3 
PiperOrigin-RevId: 607453817
 2024-02-15 14:29:45 -08:00
Skye Wanderman-Milne
b93772fc24 Update version numbers post-0.4.24 release 2024-02-06 16:28:23 -08:00
Skye Wanderman-Milne
4cca5fa1da jax 0.4.24 release take 2: update libtpu version 2024-02-05 10:15:11 -08:00
Skye Wanderman-Milne
368f594cc4 Update version numbers for 0.4.24 release 2024-01-23 16:36:10 -08:00
jax authors
f31d1e7599 Bump NCCL version on JAX OSS 
PiperOrigin-RevId: 597257455
 2024-01-10 08:28:15 -08:00
Yash Katariya
6e1ab7ca3f Finish release of jax and jaxlib 0.4.23 
PiperOrigin-RevId: 590833947
 2023-12-13 23:39:08 -08:00
Yash Katariya
1c68577dcd Prepare for emergency jax and jaxlib 0.4.23 release 
PiperOrigin-RevId: 590780824
 2023-12-13 19:02:24 -08:00
Yash Katariya
25c16c0b78 Finish jax and jaxlib 0.4.22 release 
PiperOrigin-RevId: 590775311
 2023-12-13 18:26:47 -08:00
Yash Katariya
2e633522a0 Start jax and jaxlib 0.4.22 release 
PiperOrigin-RevId: 590686003
 2023-12-13 12:57:23 -08:00
Jieying Luo
4fe9e59644 Add jax[cuda12] install variation for using cuda plugin. 
PiperOrigin-RevId: 590342149
 2023-12-12 14:40:54 -08:00
Yash Katariya
a9bfbd32e1 Finish jax and jaxlib 0.4.21 release 
PiperOrigin-RevId: 587866580
 2023-12-04 15:51:58 -08:00
Yash Katariya
5942e1580e Prepare for 0.4.21 release 
PiperOrigin-RevId: 587767502
 2023-12-04 10:26:48 -08:00
Skye Wanderman-Milne
55e3072d2e Update versions and CHANGELOG after jax 0.4.20 release 2023-11-02 16:30:56 -07:00
Skye Wanderman-Milne
6813819187 Update versions for jax + jaxlib 0.4.20 release 2023-11-02 09:34:16 -07:00
Yash Katariya
613369fc22 Finish 0.4.19 jax and jaxlib release 
PiperOrigin-RevId: 574983871
 2023-10-19 13:27:52 -07:00
Yash Katariya
f10e3338d6 Start release for jax 0.4.19 
PiperOrigin-RevId: 574575158
 2023-10-18 13:28:52 -07:00
Peter Hawkins
927a182d03 Fix typo in setup.py 2023-10-11 15:15:18 -07:00
Peter Hawkins
b8c0ca1185 Add a version constraint on nvidia-nvjitlink-cu12. 
This works around a missing version constraint on NVIDIA's CUDA packages, for example nvidia-cusolver-cu12 should have a versioned dependency on nvidia-nvjitlink-cu12.

Issue https://github.com/google/jax/issues/18027
 2023-10-11 09:52:20 -07:00
Peter Hawkins
84b58ec7f3 Increase minimum scipy version to 1.9. 
Scipy 1.9 appears to fix some crashes on Mac ARM.

PiperOrigin-RevId: 571977068
 2023-10-09 10:37:35 -07:00
Skye Wanderman-Milne
a06beaa1a2 Update versions post jax 0.4.18 release 2023-10-06 17:20:34 -07:00
Skye Wanderman-Milne
d4a1bb9292 Update setup.py and CHANGELOG for jax 0.4.18 release 2023-10-06 13:13:33 -07:00
Cristian Garcia
7498ffe843 condition numpy version based on python version 2023-10-04 21:06:01 +00:00
Peter Hawkins
efc18e4147 [JAX] Obtain NCCL via a stub, rather than linking it statically or dynamically. 
This shrinks the CUDA jaxlib wheel size by around 80MB.

PiperOrigin-RevId: 570554454
 2023-10-03 18:33:58 -07:00
Skye Wanderman-Milne
82b58386b7 Update versions and CHANGELOG after jax 0.4.17 release 2023-10-03 17:54:35 -07:00