Finish jax and jaxlib release 0.4.0

PiperOrigin-RevId: 494833878

CHANGELOG.md

@@ -6,7 +6,11 @@ Best viewed [here](https://jax.readthedocs.io/en/latest/changelog.html).
 Remember to align the itemized text with the first line of an item within a list.
 -->
 
-## jax 0.4.0
+## jax 0.4.1
+
+## jaxlib 0.4.1
+
+## jax 0.4.0 (Dec 12, 2022)
 * Changes
   * Support for Python 3.7 has been dropped, in accordance with JAX's
     {ref}`version-support-policy`.
@@ -39,7 +43,7 @@ Remember to align the itemized text with the first line of an item within a list
     It will be removed in a future release, in accordance with the {ref}`api-compatibility`
     policy. It can be replaced with `jnp.sort(a, axis=0)`.
 
-## jaxlib 0.4.0
+## jaxlib 0.4.0 (Dec 12, 2022)
 * Changes
   * Support for Python 3.7 has been dropped, in accordance with JAX's
     {ref}`version-support-policy`.

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.4.0"
+__version__ = "0.4.1"
 _minimum_jaxlib_version = "0.3.22"
 
 def _version_as_tuple(version_str):

setup.py

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