From d4a1bb9292f28d2ad1424ebdbf9ee32384169d36 Mon Sep 17 00:00:00 2001
From: Skye Wanderman-Milne <skyewm@google.com>
Date: Fri, 6 Oct 2023 13:10:59 -0700
Subject: [PATCH] Update setup.py and CHANGELOG for jax 0.4.18 release

---
 CHANGELOG.md | 3 +++
 setup.py     | 4 ++--
 2 files changed, 5 insertions(+), 2 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1ba2211d8b6c..5b39f112e69c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -26,6 +26,9 @@ Remember to align the itemized text with the first line of an item within a list
       `standard_primitive`, `standard_unop`, `unop`, and `unop_dtype_rule` are
       internal utilities, now deprecated without replacement.
 
+* Bug fixes
+  * Fixed Cloud TPU regression where compilation would OOM due to smem.
+
 # jax 0.4.17 (Oct 3, 2023)
 
 * New features
diff --git a/setup.py b/setup.py
index 987f79ed7a77..78b0a9b69322 100644
--- a/setup.py
+++ b/setup.py
@@ -22,13 +22,13 @@
 
 project_name = 'jax'
 
-_current_jaxlib_version = '0.4.17'
+_current_jaxlib_version = '0.4.18'
 # The following should be updated with each new jaxlib release.
 _latest_jaxlib_version_on_pypi = '0.4.17'
 _available_cuda11_cudnn_versions = ['86']
 _default_cuda11_cudnn_version = '86'
 _default_cuda12_cudnn_version = '89'
-_libtpu_version = '0.1.dev20231003'
+_libtpu_version = '0.1.dev20231006'
 
 def load_version_module(pkg_path):
   spec = importlib.util.spec_from_file_location(