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
13,235 Commits 508 Branches 263 Tags
Commit Graph

118 Commits

Author SHA1 Message Date
Peter Hawkins
b49e31a012 Update version numbers after release. 2022-09-28 18:49:22 +00:00
jax authors
0282b4bfad Merge pull request #12538 from jakevdp:bundle-pyi 
PiperOrigin-RevId: 477453094
 2022-09-28 08:00:20 -07:00
Peter Hawkins
8d8643664c jax/jaxlib 0.3.20 release candidate. 2022-09-28 13:33:52 +00:00
Jake VanderPlas
6e6fb10ca3 setup: bundle *.pyi files with distribution 2022-09-27 12:55:42 -07:00
Skye Wanderman-Milne
3c0d280bc0 Update version and changelog for jax 0.3.18 release 2022-09-26 12:43:39 -07:00
Peter Hawkins
bcd36d8eb2 Jax and jaxlib 0.3.18 release candidate. 2022-09-26 14:10:57 +00:00
Peter Hawkins
ba557d5e1b Change JAX's copyright attribution from "Google LLC" to "The JAX Authors.". 
See https://opensource.google/documentation/reference/releasing/contributions#copyright for more details.

PiperOrigin-RevId: 476167538
 2022-09-22 12:27:19 -07:00
Sharad Vikram
b0309dc33a Bump libtpu nightly version 2022-08-11 15:51:41 -07:00
Parker Schuh
022fedde98 Move protobuf deps to be optional. 
PiperOrigin-RevId: 466123822
 2022-08-08 12:57:56 -07:00
Peter Hawkins
c735c6bf0e Increase minimum NumPy version to 1.20. 
Per NEP 29, support for 1.19 ended on Jun 21, 2022.
 2022-08-06 14:51:14 +00:00
Parker Schuh
6b1610ce9e Add dep on protobuf and build protobufs if protoc is available. 
PiperOrigin-RevId: 464645042
 2022-08-01 16:01:02 -07:00
Skye Wanderman-Milne
d840f54fe5 Bump version numbers after 0.3.15 release 
PiperOrigin-RevId: 463344160
 2022-07-26 08:39:17 -07:00
Skye Wanderman-Milne
26fbeb6e2a Update WORKSPACE and libtpu version for jaxlib 0.3.15, take 3 2022-07-22 11:41:39 -07:00
Skye Wanderman-Milne
186a4f83e3
Update libtpu version for 0.3.15 release 2022-07-19 17:11:43 -07:00
Skye Wanderman-Milne
9149c38e1e Update WORKSPACE and setup.py in preparation for 0.3.15 jax/jaxlib release 2022-07-14 10:12:59 -07:00
Jake VanderPlas
00d8ce6c4a Populate long_description for jax & jaxlib 2022-07-13 14:03:32 -07:00
Sharad Vikram
1daea700f2 Bump JAX/Jaxlib versions 2022-06-28 14:36:47 -07:00
Peter Hawkins
1e29b7b762 Update CHANGELOG.md and setup.py for 0.3.14 release. 2022-06-27 09:38:41 -04:00
Peter Hawkins
a560a29e12 Increase the minimum scipy version to 1.5. 
We don't have a formal support policy for scipy versions, but 1.5 dates from around the same date as the oldest supported NumPy release NEP-29 would have us support (1.20).
 2022-06-24 15:07:09 -04:00
Sharad Vikram
9bd1bd67e0 Update versions for jax/jaxlib release 2022-06-21 12:57:28 -07:00
Yash Katariya
e0ff842c2a Use epath from etils package. This CL also makes epath a required dep for JAX. 
This is being used in the following ways in this CL:

* To dump IR, you can now pass paths with `gs://` or `cns` and the HLO can be dumped to those paths.
* Removing the TF dep from gda serialization.

PiperOrigin-RevId: 452117007
 2022-05-31 12:47:35 -07:00
Jeppe Klitgaard
a11f15e3ec feat: officially support Python 3.10 2022-05-07 13:43:12 +01:00
Yash Katariya
d2e3d4278d Updates values after jax and jaxlib 0.3.10 release 
PiperOrigin-RevId: 446623299
 2022-05-04 21:17:37 -07:00
Yash Katariya
38ce6d027b Update TF commit for release 
PiperOrigin-RevId: 446555288
 2022-05-04 14:42:50 -07:00
Yash Katariya
ff1a3c40ba jax and jaxlib release 
PiperOrigin-RevId: 446295827
 2022-05-03 14:52:40 -07:00
Peter Hawkins
38ea5a6bc0 Copybara import of the project: 
--
391dea76bc8fe264cf26ec93d42147f87847894d by Peter Hawkins <phawkins@google.com>:

Update version numbers after jax/jaxlib 0.3.7 release.

COPYBARA_INTEGRATE_REVIEW=https://github.com/google/jax/pull/10324 from hawkinsp:jaxlib 391dea76bc8fe264cf26ec93d42147f87847894d
PiperOrigin-RevId: 442311051
 2022-04-16 22:37:09 -07:00
Peter Hawkins
52a97f2e06 Jax 0.3.7 and jaxlib 0.3.7 release. 2022-04-15 12:02:05 -04:00
Yash Katariya
5fd78eaf02
Bump the libtpu version to prepare for JAX release 2022-04-12 11:41:07 -07:00
Peter Hawkins
4dc69034b0 Update version numbers after jax/jaxlib release. 2022-04-07 16:40:19 -04:00
Peter Hawkins
7f751c5523 Update libtpu version for jax 0.3.5 release. 2022-04-07 16:14:13 -04:00
Peter Hawkins
96ba290faf Jax 0.3.5 and jaxlib 0.3.5 release. 2022-04-06 23:56:41 +00:00
Skye Wanderman-Milne
d7087abce6 Bump jax and jaxlib versions for 0.3.2 release 
Also add CPU pjit to changelog
 2022-03-16 14:31:00 -07:00
Skye Wanderman-Milne
f9775a2ced Update CHANGELOG and setup.py for jax + jaxlib 0.3.2 releases 2022-03-16 10:17:42 -07:00
Yash Katariya
2162868ed9 Update values after release 
PiperOrigin-RevId: 427910510
 2022-02-10 20:32:53 -08:00
Yash Katariya
1ad3551ec9 Release jax and jaxlib 0.3.0 as per the new release process. 
PiperOrigin-RevId: 427809845
 2022-02-10 11:59:13 -08:00
Yash Katariya
d82bcc2a0c Add jax[_ci] option to account for the new release process. 
PiperOrigin-RevId: 427802081
 2022-02-10 11:29:51 -08:00
Skye Wanderman-Milne
d096d9a758 Update pinned libtpu-nightly version for jaxlib 0.1.76 2022-01-28 17:29:04 -08:00
Peter Hawkins
6791446bb1 Update development jaxlib version to 0.1.77, update jaxlib version in setup.py to 0.1.76. 
Changelog entry for jaxlib 0.1.77 was already added in a previous PR.

PiperOrigin-RevId: 424872047
 2022-01-28 08:10:58 -08:00
Peter Hawkins
04369a3588 Drop support for NumPy 1.18. 
Per NEP-29, we can drop NumPy 1.18 support on Dec 22, 2021.

The next NumPy deprecation will be 1.19 on Jun 21, 2022.

PiperOrigin-RevId: 419651428
 2022-01-04 12:11:38 -08:00
Yash Katariya
1b5630eed6 Update jaxlib version number to 0.1.76 
PiperOrigin-RevId: 415050863
 2021-12-08 11:14:12 -08:00
Peter Hawkins
7902ddaca2 Update jaxlib versions. 2021-11-17 11:46:41 -05:00
Yash Katariya
ee752b32f7 Use cuda11_cudnn82 instead of cuda=11,cudnn=82 because the latter one is a syntax error 
PiperOrigin-RevId: 404240654
 2021-10-19 06:24:53 -07:00
Yash Katariya
4d8bce1b85 Add a default cuda installation path and more explicit installation paths for CUDA jaxlib. 
```
# Installs Cuda 11 with Cudnn 8.2
$ pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=82] -f https://storage.googleapis.com/jax-releases/jax_releases.html

$ pip install jax[cuda=11,cudnn=805] -f https://storage.googleapis.com/jax-releases/jax_releases.html
```

PiperOrigin-RevId: 404134291
 2021-10-18 19:56:22 -07:00
Peter Hawkins
2bd010ae88 Cleanup internal representation of XLA translation rules. 
Over time JAX has sprouted many variants of XLA translation rules, each with slightly different but overlapping arguments. This change consolidates them into a new xla.TranslationRule signature:
rule(ctx, avals_in, avals_out, *args, **params)
where ctx contains the parts of the other signatures that were typically not specific to a particular equation.

Since there are many JAX rules to migrate, and even a number of translation rules belonging to projects downstream of JAX, we leave backwards compatibility shims around `xla.translations`, `xla.backend_specific_translations`, and `xla.call_translations` which seem to be the only ones used outside JAX itself.

In passing, this change alters the semantics of `backend` arguments to nested `jit` blocks. We now always canonicalize the backend to a specific backend at the outermost `jit`, and do not complain if an inner `jit` has an explicit `backend` that matches the current default choice.

PiperOrigin-RevId: 403607667
 2021-10-16 07:53:24 -07:00
Yash Katariya
66a4a9ff3f Remove 10.2 cuda support 
PiperOrigin-RevId: 402707900
 2021-10-12 18:44:07 -07:00
Skye Wanderman-Milne
0072c32546 Update CHANGELOG and verson numbers for jaxlib 0.1.72 release 2021-10-12 17:37:29 -07:00
yashkatariya
be824a792e Update files after new jaxlib release 0.1.71 2021-09-01 10:43:20 -07:00
Jake VanderPlas
062f8d2261 Specify scipy in setup.py install_requires 2021-08-20 10:16:17 -07:00
Yash Katariya
bf967d88d8 Upgrade versions after jaxlib release 
PiperOrigin-RevId: 389753047
 2021-08-09 16:37:44 -07:00
Peter Hawkins
6e9169d100 Drop support for NumPy 1.17. 2021-07-29 09:18:01 -04:00