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

230 Commits

Author SHA1 Message Date
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
Skye Wanderman-Milne
b9c602b8d7 Update XLA commit and versions for jax 0.4.17 release 2023-10-03 13:20:25 -07:00
Peter Hawkins
a2e1f1f24e Update changelog. 
Bump the minimum CUDA 12 pip package versions to the current releases.
 2023-09-26 18:21:51 -04:00
Yash Katariya
dcc465b4de Finish jax and jaxlib 0.4.16 release 
PiperOrigin-RevId: 566477931
 2023-09-18 19:09:19 -07:00
Yash Katariya
d477b92121 Start jax and jaxlib 0.4.16 release 
PiperOrigin-RevId: 566354738
 2023-09-18 11:05:11 -07:00
jax authors
f64235acc8 Merge pull request #17453 from jakevdp:fix-version-string 
PiperOrigin-RevId: 563466394
 2023-09-07 10:06:32 -07:00
Jake VanderPlas
6f3f0d5e57 build: write appropriate version strings to build artifacts 2023-09-07 08:45:48 -07:00
Skye Wanderman-Milne
f71ba0a2e7 Update versions and changelog post 0.4.15 release 2023-08-30 16:20:25 -04:00
Skye Wanderman-Milne
2c5e8577e3 Update versions for jax/jaxlib 0.4.15 release 2023-08-30 09:10:07 -04:00
Jake VanderPlas
f2e41b3c27 Add requests dependency to jax[tpu] 
Fixes #17075
 2023-08-11 11:46:58 -07:00
Skye Wanderman-Milne
011fc88c03 Update versions and changelog for jax 0.4.14 release 2023-07-27 16:22:53 -07:00