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

206 Commits

Author SHA1 Message Date
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
Skye Wanderman-Milne
0b24b2ba6a Update WORKSPACE and setup.py in preparation for jax/jaxlib 0.4.14 release 2023-07-27 13:35:04 -07:00
Jake VanderPlas
9962065deb Require ml_dtypes>=0.2 2023-07-07 12:07:44 -07:00
Jake VanderPlas
ad35702934 Drop support for numpy 1.21 
This is in accordance with NEP 29 and https://jax.readthedocs.io/en/latest/deprecation.html
 2023-06-23 10:28:26 -07:00
Peter Hawkins
bfa113ba60 Remove references to Python 3.8. 
Remove the old build scripts/Dockerfile, since they are unused and broken.

PiperOrigin-RevId: 542870354
 2023-06-23 08:48:57 -07:00
Yash Katariya
fc0dcd15a2 Copybara import of the project: 
--
57af5360a1ca1356dbf7760c76e241f7134ef6dd by Jake VanderPlas <jakevdp@google.com>:

[Roll forward] Update required Python version to 3.9

PiperOrigin-RevId: 542728213
 2023-06-22 18:58:30 -07:00
Skye Wanderman-Milne
20095ab9da Update version numbers and changelog post 0.4.13 release 2023-06-22 17:54:57 -07:00
Peter Hawkins
487b640acf Jax 0.4.13 release. 2023-06-22 14:59:36 -07:00
Peter Hawkins
0ded163027 Fix cuda12 pip install. 
The wheel is now called cudnn89.

Fixes #16362
 2023-06-12 15:38:16 -04:00
Skye Wanderman-Milne
4b80103077 Update versions and changelog post 0.4.12 release 2023-06-08 16:53:22 -07:00
Skye Wanderman-Milne
8223e452f2 Update WORKSPACE and setup.py for jax + jaxlib 0.4.12 release 2023-06-08 13:58:09 -07:00
Yash Katariya
4c48611fba Finish jax and jaxlib 0.4.11 release 
PiperOrigin-RevId: 536931532
 2023-05-31 23:49:32 -07:00
Yash Katariya
48ad9a6f3e Start jax and jaxlib 0.4.11 release 
PiperOrigin-RevId: 536860076
 2023-05-31 16:48:52 -07:00
Skye Wanderman-Milne
968237080f Add importlib_metadata to project requirements. 
This is necessary to ensure we can correctly detect PJRT plugins via
entry_points without compatibility errors.

Prior to this change, there was conditional logic to handle if
importlib_metadata wasn't installed at all. However, it doesn't handle
the case where importlib_metadata is installed by not high enough
version to support Python 3.10 compat. This change gets rid of that
logic and just ensures the right version is installed.

All of this logic can be removed if/when jax requires Python version
>= 3.10

This also removes an unnecessary `requests` dep for the [tpu] install.
 2023-05-31 21:03:12 +00:00
Peter Hawkins
69cf67f252 Bump the minimum CUDNN version for CUDA 12 wheels to 8.9. 2023-05-26 10:04:34 -04:00
Peter Hawkins
2b7790290b Bump minimum CUDNN version in pip installation to 8.8. 
There are known wrong output bugs observed in JAX for earlier versions, in particular related to RNNs.
 2023-05-25 14:46:39 -04:00
Skye Wanderman-Milne
533a7c05f1 Update versions and changelog post 0.4.10 release 2023-05-11 18:16:02 -07:00
Skye Wanderman-Milne
82bbeef519 Update setup.py, WORKSPACE, and CHANGELOG for jax/jaxlib 0.4.10 release 2023-05-11 14:46:06 -07:00
Skye Wanderman-Milne
b02b043e7f Update versions and changelog for 0.4.9 release 2023-05-09 17:06:59 -07:00
Skye Wanderman-Milne
5bcd9dcc46 Update WORKSPACE and setup.py in preparation for jax/jaxlib 0.4.9 release, take 2 2023-05-09 14:49:54 -07:00
Skye Wanderman-Milne
5e9364abc6 Revert setup.py changes. 
This reverts the setup.py changes from
f28b20175f307d5a56502446a9706480126a5bd4. We actually need to fix some
more issues before releasing 0.4.9, so fix the install at HEAD in the
meantime.
 2023-05-08 09:58:51 -07:00
Skye Wanderman-Milne
f28b20175f Update WORKSPACE and setup.py in preparation for jax/jaxlib 0.4.9 release 2023-05-04 14:38:46 -07:00
Jake VanderPlas
59e6ed213e Use ml_dtypes definition for jnp.finfo 2023-05-04 10:40:44 -07:00
Yash Katariya
6506ee2a40 Copybara import of the project: 
--
57af5360a1ca1356dbf7760c76e241f7134ef6dd by Jake VanderPlas <jakevdp@google.com>:

[Rollback] Update required Python version to 3.9

PiperOrigin-RevId: 528905991
 2023-05-02 15:33:29 -07:00
Jake VanderPlas
57af5360a1 Update required Python version to 3.9 2023-05-01 10:00:57 -07:00