mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-17 20:36:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
14,091 Commits 504 Branches 263 Tags
Commit Graph

135 Commits

Author SHA1 Message Date
Yash Katariya
0bdb7ec042 Finish jax and jaxlib release 0.4.0 
PiperOrigin-RevId: 494833878
 2022-12-12 14:43:35 -08:00
Yash Katariya
0118f8d568 Prepare for jax and jaxlib 0.4.0 release 
PiperOrigin-RevId: 493733609
 2022-12-07 16:02:24 -08:00
Jake VanderPlas
4389216d0c Remove typing_extensions dependency 2022-12-05 15:42:26 -08:00
Yash Katariya
25d1a0b4c6 Add cudnn 86 (for cuda 11.8) so that I can release cuda 11.8 nightlies. 
PiperOrigin-RevId: 493086060
 2022-12-05 12:50:09 -08:00
Jake VanderPlas
cb62a31653 Drop support for Python 3.7 2022-11-29 15:01:47 -08:00
Yash Katariya
8c42edfec1 Finish jax and jaxlib release 0.3.25. The next release will be 0.4.0 (since jax.Array will be enabled in that release) 
PiperOrigin-RevId: 488672395
 2022-11-15 09:02:53 -08:00
Yash Katariya
a683186570 Use the 11/09 libtpu build for jaxlib release since that passes all the tests. 
PiperOrigin-RevId: 488543322
 2022-11-14 20:37:41 -08:00
Yash Katariya
f36084acd3 Update the values for jaxlib release (again) 
PiperOrigin-RevId: 488522992
 2022-11-14 18:31:08 -08:00
Yash Katariya
0da02dd41c Update the values needed for a jaxlib release 
PiperOrigin-RevId: 488508360
 2022-11-14 17:08:59 -08:00
Yash Katariya
1d48c93b0e Finish the release of jax and jaxlib 0.3.24 
PiperOrigin-RevId: 486162090
 2022-11-04 09:43:12 -07:00
Yash Katariya
2e384ce58f Prepare for release of jax and jaxlib 0.3.24 
PiperOrigin-RevId: 485985460
 2022-11-03 15:13:23 -07:00
Jake VanderPlas
41b815bf33 Fix jax releases URL in setup.py comments 2022-10-25 10:39:39 -07:00
Peter Hawkins
fb72c38e19 Add Python 3.11 as a compatible Python version. 2022-10-14 14:56:07 +00:00
Nicholas Junge
efd61b73f6 Migrate JAX internals to builtin Python logging 
This commit changes the JAX codebase to use Python's builtin logging instead of ABSL logging. With the latter being used in JAX code as of now, the change to Python builtin logging is advised for the following reasons (among others):

- absl-py can be removed as an external dependency of JAX.
- Builtin logging brings the option of adding more log handlers, for example file handlers for log dumps or writers to different IO streams.

Logging in JAX is ported over to take place at the module level. While previously, some Python namespaces within JAX already used module-scoped logging via absl.vlog, the following idiom was adopted to provide the same functionality in Python builtin logging:

```py
import logging
logger = logging.getLogger(__name__)

logger.debug(...)
logger.info(...)
```

 The builtin root logger is left untouched, which is beneficial for downstream users planning to customize the Python root logger. All JAX internal code promises to log to descendants of the top-level "jax" logger by virtue of log propagation.

The package `absl-py` was removed from JAX's install requirements, and added into its test requirements.
 2022-10-13 21:32:44 +02:00
Peter Hawkins
91ea57d981 Make the etils[epath] dependency of JAX optional. 
If epath is installed then we will use it for file I/O. If it is not, we
will fall back to pathlib from the standard library.
 2022-10-12 21:16:48 +00:00
Skye Wanderman-Milne
a6e0e77624 Update version.py, setup.py and CHANGELOG post jax 0.3.22 release 2022-10-11 21:37:53 +00:00
Skye Wanderman-Milne
63be2201aa Update WORKSPACE and setup.py for jaxlib 0.3.22 release 2022-10-11 17:55:05 +00:00
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