256 Commits

Author SHA1 Message Date
Peter Hawkins
81991d87c8 JAX release 0.4.35 2024-10-22 15:00:23 -04:00
Peter Hawkins
e4f3f8f064 Use libtpu releases rather than libtpu-nightly for jax[tpu].
PiperOrigin-RevId: 688632409
2024-10-22 11:47:07 -07:00
Jake VanderPlas
74fc7360f5 Require ml_dtypes >= 0.4.0
This is the minimum version that supports NumPy 2.0.
2024-10-22 09:04:28 -07:00
Peter Hawkins
b0b7a60e63 Merge branch 'release/0.4.34' 2024-10-04 10:56:18 -04:00
Peter Hawkins
affba367c5 JAX 0.4.34 release 2024-10-02 14:37:09 -04:00
Yu-Hang Tang
c88c3aecae add k8s cluster environment 2024-09-20 17:26:53 +00:00
Michael Hudgins
d4d1518c3d Update references to the GitHub url in JAX codebase to reflect move from google/jax to jax-ml/jax
PiperOrigin-RevId: 676843138
2024-09-20 07:52:33 -07:00
Peter Hawkins
ae0e403c60 Merge release/0.4.33 into main and update version numbers. 2024-09-16 18:46:24 +00:00
Peter Hawkins
80e1c94de6 Prepare for v0.4.33 release.
This release is branched off the v0.4.32 release, with two changes:
a) a fixed libtpu pin, and
b) a patch to revert an F64 tanh issue on CPU.
2024-09-16 13:30:35 +00:00
Peter Hawkins
3e81ae530d Update version numbers after v0.4.32 release. 2024-09-11 16:18:56 -04:00
Peter Hawkins
1594d2f30f Prepare for v0.4.32 release. 2024-09-11 14:29:55 -04:00
Yash Katariya
0b04dd022a Reverts 5e4250e64bb415be94ddc8a80dba6083a6a4123a
PiperOrigin-RevId: 673141373
2024-09-10 16:01:00 -07:00
Yash Katariya
5e4250e64b Prepare for jax 0.4.32 release
PiperOrigin-RevId: 673105544
2024-09-10 14:26:17 -07:00
Yash Katariya
2106a25977 Finish jax and jaxlib v0.4.31 release
PiperOrigin-RevId: 657388782
2024-07-29 17:57:37 -07:00
Yash Katariya
7fd9302785 Start JAX and jaxlib 0.4.31 release
PiperOrigin-RevId: 657295431
2024-07-29 12:49:58 -07:00
Peter Hawkins
fd23b8733d Bump minimum SciPy version to 1.10.
SciPy 1.9.0 was released July 29, 2022, which is 24 months ago

PiperOrigin-RevId: 657215038
2024-07-29 08:50:18 -07:00
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