mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #16527 from skye:version

Browse Source 
PiperOrigin-RevId: 542678766
This commit is contained in:
jax authors 2023-06-22 15:08:21 -07:00
commit c3e2427008
5 changed files with 16 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
.github/workflows/wheel_win_x64.yml vendored
View File

@ -1,8 +1,6 @@
name: Wheel build - Windows CPU x86_64
on:
workflow_dispatch: # allows triggering the workflow run manually
release:
types: [published]
env:
DISTUTILS_USE_SDK: 1

8
CHANGELOG.md
View File

@ -25,7 +25,9 @@ Remember to align the itemized text with the first line of an item within a list
determine the output shardings.
* If the mesh context manager is provided, None will imply that the value
will be replicated on all devices of the mesh.
* Executable.cost_analysis() works on Cloud TPU
* Executable.cost_analysis() works on Cloud TPU
* Added a warning if a non-allowlisted `jaxlib` plugin is in use.
* Added `jax.tree_util.tree_leaves_with_path`.
* Bug fixes
* Fixed incorrect wheel name in CUDA 12 releases (#16362); the correct wheel
@ -38,9 +40,13 @@ Remember to align the itemized text with the first line of an item within a list
## jaxlib 0.4.13
* Changes
* Added Windows CPU-only wheels to the `jaxlib` Pypi release.
* Bug fixes
* `__cuda_array_interface__` was broken in previous jaxlib versions and is now
fixed ({jax-issue}`16440`).
* Concurrent CUDA kernel tracing is now enabled by default on NVIDIA GPUs.
## jax 0.4.12 (June 8, 2023)

8
README.md
View File

@ -526,10 +526,10 @@ the following in your cloud TPU VM:
pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
```
For interactive notebook users: Colab TPUs are no longer supported by JAX as of
version 0.4. However, for an interactive TPU notebook in the cloud, you can use
[Kaggle TPU notebooks](https://www.kaggle.com/docs/tpu), which are fully
supported by JAX.
For interactive notebook users: Colab TPUs no longer support JAX as of
JAX version 0.4. However, for an interactive TPU notebook in the cloud, you can
use [Kaggle TPU notebooks](https://www.kaggle.com/docs/tpu), which fully
support JAX.
### pip installation: Apple GPUs

6
WORKSPACE
View File

@ -7,10 +7,10 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
# and update the sha256 with the result.
http_archive(
name = "xla",
sha256 = "f8f6efd93237c94ee3c8be65dd0702a065c3e3e7b0f732181e62d33813f1bd35",
strip_prefix = "xla-34521d16cd91423a041be86bcb35319dfee3bedb",
sha256 = "4ec16aff3862c5a243db956ce558d7a62eb79f5e20747b0e80802a3b0d12e419",
strip_prefix = "xla-12de6ec958419b57be248d0acd2d9f757e71748c",
urls = [
"https://github.com/openxla/xla/archive/34521d16cd91423a041be86bcb35319dfee3bedb.tar.gz",
"https://github.com/openxla/xla/archive/12de6ec958419b57be248d0acd2d9f757e71748c.tar.gz",
],
)

4
setup.py
View File

@ -19,13 +19,13 @@ import sys
from setuptools import setup, find_packages
_current_jaxlib_version = '0.4.12'
_current_jaxlib_version = '0.4.13'
# The following should be updated with each new jaxlib release.
_latest_jaxlib_version_on_pypi = '0.4.12'
_available_cuda11_cudnn_versions = ['86']
_default_cuda11_cudnn_version = '86'
_default_cuda12_cudnn_version = '89'
_libtpu_version = '0.1.dev20230608'
_libtpu_version = '0.1.dev20230622'
_dct = {}
with open('jax/version.py', encoding='utf-8') as f: