Update version numbers after jax/jaxlib release.

CHANGELOG.md

@@ -8,9 +8,16 @@ 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.5 (April 6, 2022)
+## jax 0.3.6 (unreleased)
 * [GitHub
-  commits](https://github.com/google/jax/compare/jax-v0.3.4...main).
+  commits](https://github.com/google/jax/compare/jax-v0.3.5...main).
+
+## jaxlib 0.3.6 (unreleased)
+
+
+## jax 0.3.5 (April 7, 2022)
+* [GitHub
+  commits](https://github.com/google/jax/compare/jax-v0.3.4...jax-v0.3.5).
 * Changes:
   * added {func}`jax.random.loggamma` & improved behavior of {func}`jax.random.beta`
     and {func}`jax.random.dirichlet` for small parameter values ({jax-issue}`#9906`).
@@ -24,7 +31,7 @@ PLEASE REMEMBER TO CHANGE THE '..main' WITH AN ACTUAL TAG in GITHUB LINK.
   * {func}`jax.tree_util.tree_multimap` is deprecated. Use {func}`jax.tree_util.tree_map` instead ({jax-issue}`#5746`).
   * `jax.experimental.sharded_jit` is deprecated. Use `pjit` instead.
 
-## jaxlib 0.3.5 (April 6, 2022)
+## jaxlib 0.3.5 (April 7, 2022)
 * Bug fixes
   * Fixed a bug where double-precision complex-to-real IRFFTs would mutate their
     input buffers on GPU ({jax-issue}`#9946`).

jax/version.py

@@ -15,7 +15,7 @@
 def _version_as_tuple(version_str):
   return tuple(int(i) for i in version_str.split(".") if i.isdigit())
 
-__version__ = "0.3.5"
+__version__ = "0.3.6"
 __version_info__ = _version_as_tuple(__version__)
 
 _minimum_jaxlib_version = "0.3.0"

jaxlib/version.py

@@ -17,4 +17,4 @@
 # reflect the most recent available binaries.
 # __version__ should be increased after releasing the current version
 # (i.e. on main, this is always the next version to be released).
-__version__ = "0.3.5"
+__version__ = "0.3.6"

setup.py

@@ -16,7 +16,7 @@ from setuptools import setup, find_packages
 
 _current_jaxlib_version = '0.3.5'
 # The following should be updated with each new jaxlib release.
-_latest_jaxlib_version_on_pypi = '0.3.2'
+_latest_jaxlib_version_on_pypi = '0.3.5'
 _available_cuda_versions = ['11']
 _default_cuda_version = '11'
 _available_cudnn_versions = ['82', '805']