mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-14 10:56:06 +00:00
Code Issues Packages Projects Releases Wiki Activity
27,027 Commits 495 Branches 263 Tags
Commit Graph

218 Commits

Author SHA1 Message Date
Peter Hawkins
3736e5ba85 Bump the JAX version to v0.6.0, which will be the next release version. 
PiperOrigin-RevId: 746490665
 2025-04-11 09:34:42 -07:00
Nitin Srinivasan
64e4bf2632 Relax jax dependency constraints to be able to install RC wheels 
Also, add a job to the release test workflow that verifies that the release wheels can be installed.

TESTED:
1. Full release: https://github.com/jax-ml/jax/actions/runs/14315832784

2. jax only release: https://github.com/jax-ml/jax/actions/runs/14316157252

PiperOrigin-RevId: 744857804
 2025-04-07 14:49:55 -07:00
Peter Hawkins
9d534ad2cd Update version numbers after JAX 0.5.3 release. 2025-03-19 14:41:25 -04:00
Peter Hawkins
ed43119a86 JAX release v0.5.3 2025-03-18 21:38:14 -04:00
Skye Wanderman-Milne
cebedb9f1a Update version number after 0.5.2 release 2025-03-04 18:49:12 -08:00
Peter Hawkins
c8c4cfa04e Update version numbers after 0.5.1 release. 2025-02-24 16:18:25 -05:00
Yash Katariya
07440f4afa Prepare for JAX release 0.5.1 2025-02-24 10:59:04 -05:00
jax authors
d424f5b5b3 Refactor JAX wheel build rules to control the wheel filename and maintain reproducible wheel content and filename results. 
This change is a part of the initiative to test the JAX wheels in the presubmit properly.

The list of the changes:
1. JAX wheel build rule verifies that `--@local_config_cuda//cuda:include_cuda_libs=false` during the wheel build. There is a way to pass the restriction by providing `--@local_config_cuda//cuda:override_include_cuda_libs=true`.

2. The JAX version number (which is also used in the wheel filenames) is stored in `_version` variable in the file [version.py](https://github.com/jax-ml/jax/blob/main/jax/version.py). The custom repository rule `jax_python_wheel_version_repository` saves this value in `wheel_version.bzl`, so it becomes available in Bazel build phase.

3. The version suffix of the wheel in the build rule output depends on the environment variables.

   The version suffix chunks that are not reproducible shouldn’t be calculated as a part of the wheel binary: for example, the current date changes every day, thus the wheels built today and tomorrow on the same code version will be technically different. To maintain reproducible wheel content, we need to pass suffix chunks in a form of environment variables.

4. Environment variables combinations for creating wheels with different versions:
  * `0.5.1.dev0+selfbuilt` (local build, default build rule behavior): `--repo_env=ML_WHEEL_TYPE=snapshot`
  * `0.5.1` (release): `--repo_env=ML_WHEEL_TYPE=release`
  * `0.5.1rc1` (release candidate): `--repo_env=ML_WHEEL_TYPE=release --repo_env=ML_WHEEL_VERSION_SUFFIX=rc1`
  * `0.5.1.dev20250128+3e75e20c7` (nightly build): `--repo_env=ML_WHEEL_TYPE=custom --repo_env=ML_WHEEL_BUILD_DATE=20250128 --repo_env=ML_WHEEL_GIT_HASH=$(git rev-parse HEAD)`

PiperOrigin-RevId: 723552265
 2025-02-05 10:01:23 -08:00
Peter Hawkins
9fa2912254 Update version numbers after 0.5.0 release 2025-01-17 13:30:59 -05:00
Peter Hawkins
c25fb92c44 Release JAX 0.5.0 2025-01-17 10:28:03 -05:00
Peter Hawkins
3a8f31aa83 Update the JAX version to 0.5.0. 
This is because of the breaking change to PRNG key semantics, and the version follows JAX's new effver versioning scheme (https://jax.readthedocs.io/en/latest/jep/25516-effver.html).
 2025-01-15 14:08:15 -05:00
Peter Hawkins
ff52aedf67 Update version numbers after release. 2024-12-17 18:16:25 -05:00
Peter Hawkins
c4e3f5a356 Release v0.4.38 2024-12-17 14:35:08 -05:00
Peter Hawkins
820f51dc53 Merge branch 'release/0.4.37' into main. 2024-12-09 20:21:43 -05:00
Peter Hawkins
ffb07cdadb Update versions for v0.4.37 release. 2024-12-09 15:39:59 -05:00
Peter Hawkins
ba626fa650 Bump JAX version after release. 
PiperOrigin-RevId: 703472753
 2024-12-06 05:58:09 -08:00
Peter Hawkins
7e6620a577 JAX release 0.4.36. 2024-12-05 16:08:39 -05:00
Yunlong Liu
27bf80a506 Adds an env that can let users provide a custom version suffix for jax dev build. 
fix the local version

update to what Jake suggested
 2024-11-16 14:39:54 -08:00
Peter Hawkins
2aeda17829 Merge branch 'release/0.4.35' 2024-10-23 08:50:31 -04:00
Peter Hawkins
81991d87c8 JAX release 0.4.35 2024-10-22 15:00:23 -04: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
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
46a7be6d7a Lower the minimum jaxlib version to 0.4.31, since 0.4.32 was yanked from pypi. 
PiperOrigin-RevId: 673992844
 2024-09-12 13:46:23 -07: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
30037547d7 Bump minimum jaxlib version to 0.4.31. The corresponding xla_extension_version is 279 and mlir_api_version is 57 
PiperOrigin-RevId: 657400413
 2024-07-29 18:44:31 -07:00
Yash Katariya
2106a25977 Finish jax and jaxlib v0.4.31 release 
PiperOrigin-RevId: 657388782
 2024-07-29 17:57:37 -07:00
Peter Hawkins
07d24e7dcc Bump minimum jaxlib version to v0.4.30. 
This corresponds to xla_extension_version 271 and mlir_api_version 57.
 2024-06-18 12:35:08 -04: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
6fa31e59c4 Update version numbers after v0.4.29 release. 2024-06-10 14:37:53 -04:00
Jake VanderPlas
9ac1d38226 Finish jax and jaxlib 0.4.28 release 
PiperOrigin-RevId: 632653310
 2024-05-10 18:06:52 -07:00
Yash Katariya
395d3cb79e Bump minimum jaxlib version to 0.4.27 
xla_extension_version is 261 and mlir_api_version is 56

PiperOrigin-RevId: 631579739
 2024-05-07 16:07:59 -07:00
Yash Katariya
5031a1ddc4 Finish jax and jaxlib 0.4.27 release 
PiperOrigin-RevId: 631486157
 2024-05-07 11:14:09 -07:00
Sergei Lebedev
754fab91f7 Bumped the minimum jaxlib to 0.4.23 
jaxlib 0.4.23 has xla_extension_version 223 and mlir_api_version 54.
 2024-04-13 08:18:33 +01:00
Olli Lupton
b27b77446f _version_from_git_tree: avoid git describe 2024-04-10 12:47:47 +00:00
Yash Katariya
24517ca3e0 Finish jax and jaxlib 0.4.26 release 
PiperOrigin-RevId: 621658207
 2024-04-03 15:40:24 -07:00
Yash Katariya
e0fd29082d Finish jax and jaxlib 0.4.25 release 
PiperOrigin-RevId: 610413312
 2024-02-26 08:19:05 -08:00
Peter Hawkins
aad02dba7e Increase minimum jaxlib version to 0.4.20. 
jaxlib 0.4.20 has xla_extension_version 210 and mlir_api_version 54.

PiperOrigin-RevId: 609094229
 2024-02-21 12:58:57 -08:00
Skye Wanderman-Milne
b93772fc24 Update version numbers post-0.4.24 release 2024-02-06 16:28:23 -08:00
Parker Schuh
23b9c2a22f Add the githash that the jaxlib was built at to __init__.py. This is to allow identifying the githash of nightlies. 
PiperOrigin-RevId: 595529249
 2024-01-03 16:12:23 -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
25c16c0b78 Finish jax and jaxlib 0.4.22 release 
PiperOrigin-RevId: 590775311
 2023-12-13 18:26:47 -08:00
Yash Katariya
a9bfbd32e1 Finish jax and jaxlib 0.4.21 release 
PiperOrigin-RevId: 587866580
 2023-12-04 15:51:58 -08:00
Olli Lupton
e50c35d1e7 Fix repeatedly building JAX. 
Reproducer was essentially running `pip install .` twice in a row in the
same source directory. Closes google/jax#18252.
 2023-11-30 13:31:17 +01:00
Jake VanderPlas
cc9656d654 version: fix commit hash extraction 2023-11-22 09:35:03 -08:00