Finish jax and jaxlib 0.4.4 release

PiperOrigin-RevId: 510234171

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.4
+## jax 0.4.5
+
+## jaxlib 0.4.5
+
+## jax 0.4.4 (Feb 16, 2023)
 
 * Changes
   * The implementation of `jit` and `pjit` has been merged. Merging jit and pjit
@@ -44,7 +48,7 @@ Remember to align the itemized text with the first line of an item within a list
     The previous behavior of broadcating the output against non-scalar `initial`
     values was an unintentional implementation detail ({jax-issue}`#14446`).
 
-## jaxlib 0.4.4
+## jaxlib 0.4.4 (Feb 16, 2023)
   * Breaking changes
     * Support for NVIDIA Kepler series GPUs has been removed from the default
       `jaxlib` builds. If Kepler support is needed, it is still possible to

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.4"
+__version__ = "0.4.5"
 _minimum_jaxlib_version = "0.4.2"
 
 def _version_as_tuple(version_str):

setup.py

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