From 0bdb7ec04299fcee0c77b3b1cb306e78620b2279 Mon Sep 17 00:00:00 2001
From: Yash Katariya <yashkatariya@google.com>
Date: Mon, 12 Dec 2022 14:42:48 -0800
Subject: [PATCH] Finish jax and jaxlib release 0.4.0

PiperOrigin-RevId: 494833878
---
 CHANGELOG.md   | 8 ++++++--
 jax/version.py | 2 +-
 setup.py       | 2 +-
 3 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index ea901c9b4..ed2fb737d 100644
--- a/CHANGELOG.md
+++ b/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`.
diff --git a/jax/version.py b/jax/version.py
index af14f7598..924fff6e1 100644
--- a/jax/version.py
+++ b/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):
diff --git a/setup.py b/setup.py
index 72aa11da1..ac1cedc72 100644
--- a/setup.py
+++ b/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']