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

153 Commits

Author SHA1 Message Date
Yash Katariya
e9cac5eb47 Prepare for jax and jaxlib 0.4.7 release 
PiperOrigin-RevId: 519785176
 2023-03-27 11:45:22 -07:00
Peter Hawkins
40fb646e35 Fix duplicate definition of 'cuda' extra in setup.py. 
PiperOrigin-RevId: 519750659
 2023-03-27 09:52:37 -07:00
Peter Hawkins
b62f114524 Add support for using pip-installed CUDA wheels. 
Add a currently undocumented jax[cuda11_pip] and jax[cuda12_pip] that depend on the pip CUDA wheels.
Add a currently undocumented jax[cuda11_local] and jax[cuda12_local] that avoid the CUDA wheel dependency.
 2023-03-26 12:35:00 +00:00
Peter Hawkins
b7375b316b Increase minimum NumPy version to 1.21. 
Also increase minimum SciPy version to 1.7, which was released just before NumPy 1.21.
 2023-03-23 21:15:10 -04:00
Peter Hawkins
8bb90b5fbe [XLA:Python] Change JAX and the XLA Python extension to get NumPy bfloat16/float8 types from ml_dtypes. 
PiperOrigin-RevId: 518830467
 2023-03-23 05:13:39 -07:00
Skye Wanderman-Milne
6560bf8c36 Update versions and changelog for jax + jaxlib 0.4.6 release 2023-03-09 14:50:42 -08:00
Skye Wanderman-Milne
1aa08fd4a0 Update WORKSPACE and setup.py for jax/jaxlib 0.4.6 release 2023-03-09 10:44:58 -08:00
pizzud
631e4ed7e0 lax_test: Create a separate module for lax-specific test utils in a new package. 
These utils are currently shared with lax_vmap_test by importing lax_test as a
library, which is an odd thing to do.

The new package and the module within it are not built into the wheel, as these
are internal utilities for JAX's tests, not utilities for JAX users writing
their own tests.

Followup changes will add additional existing internal test utilities to this
package. This will allow removing sys.path manipulation from
deprecation_module_test and hopefully lazy_loader_test, as well as removing
the non-public test_util.py from _src to make it clearer that it should not be
used from outside JAX.

PiperOrigin-RevId: 510260230
 2023-02-16 15:29:41 -08:00
Yash Katariya
941722f7db Finish jax and jaxlib 0.4.4 release 
PiperOrigin-RevId: 510234171
 2023-02-16 13:54:56 -08:00
Yash Katariya
58e46b48e6 Prepare for jax and jaxlib 0.4.4 release 
PiperOrigin-RevId: 510152471
 2023-02-16 08:37:15 -08:00
Skye Wanderman-Milne
21f12183bf Post 0.4.3 release updates 2023-02-08 10:08:59 -08:00
Skye Wanderman-Milne
8ab158574d Update WORKSPACE and setup.py for jax/jaxlib 0.4.3 release 2023-02-07 15:45:28 -08:00
Skye Wanderman-Milne
b6a8aa6394 Update versions for jaxlib 0.4.2 release. 
I also screwed up the CHANGELOG before (I shouldn't have added a
date), so I'm fixing the dates now.
 2023-01-26 01:12:06 +00:00
Skye Wanderman-Milne
c4ad27c363 Update libtpu version for jaxlib 0.4.2 release (again) 2023-01-25 01:34:16 +00:00
Skye Wanderman-Milne
3f4bd5f449 Updates for jax + jaxlib 0.4.2 release 2023-01-20 19:04:46 +00:00
Jake VanderPlas
9d100ae9f4 Explicitly set utf-8 encoding in setup.py 2023-01-05 09:41:18 -08:00
Yash Katariya
835d0c979a Finish jax and jaxlib 0.4.2 release 
PiperOrigin-RevId: 495068000
 2022-12-13 10:51:13 -08:00
Yash Katariya
c4d590b1b6 Update values for release 0.4.1 
PiperOrigin-RevId: 494889744
 2022-12-12 19:04:38 -08:00
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