Merge pull request #12565 from hawkinsp:release

PiperOrigin-RevId: 477549228

CHANGELOG.md

@@ -8,13 +8,17 @@ Remember to align the itemized text with the first line of an item within a list
 PLEASE REMEMBER TO CHANGE THE '..main' WITH AN ACTUAL TAG in GITHUB LINK.
 -->
 
-## jax 0.3.20
+## jax 0.3.21
+
+## jaxlib 0.3.21
+
+## jax 0.3.20 (Sep 28, 2022)
 * Bug fixes:
   * Adds missing `.pyi` files that were missing from the previous release ({jax-issue}`#12536`).
   * Fixes an incompatibility between `jax` 0.3.19 and the libtpu version it pinned ({jax-issue}`#12550`). Requires jaxlib 0.3.20.
   * Fix incorrect `pip` url in `setup.py` comment ({jax-issue}`#12528`).
 
-## jaxlib 0.3.20
+## jaxlib 0.3.20 (Sep 28, 2022)
 * Bug fixes
   * Fixes support for limiting the visible CUDA devices via
    `jax_cuda_visible_devices` in distributed jobs. This functionality is needed for

jax/version.py

@@ -15,7 +15,7 @@
 # This file is included as part of both jax and jaxlib. It is also
 # eval()-ed by setup.py, so it should not have any dependencies.
 
-__version__ = "0.3.20"
+__version__ = "0.3.21"
 _minimum_jaxlib_version = "0.3.15"
 
 def _version_as_tuple(version_str):

setup.py

@@ -21,7 +21,7 @@
 
 _current_jaxlib_version = '0.3.20'
 # The following should be updated with each new jaxlib release.
-_latest_jaxlib_version_on_pypi = '0.3.15'
+_latest_jaxlib_version_on_pypi = '0.3.20'
 _available_cuda_versions = ['11']
 _default_cuda_version = '11'
 _available_cudnn_versions = ['82', '805']