From ff1a3c40bad4da4429793ea2773b1d318507f28b Mon Sep 17 00:00:00 2001
From: Yash Katariya <yashkatariya@google.com>
Date: Tue, 3 May 2022 14:52:12 -0700
Subject: [PATCH] jax and jaxlib release

PiperOrigin-RevId: 446295827
---
 CHANGELOG.md | 18 ++++++++++++++----
 WORKSPACE    |  6 +++---
 setup.py     |  4 ++--
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7faa6c9a1..6f42b033a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,11 +8,21 @@ 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.10 (Unreleased)
-* [GitHub commits](https://github.com/google/jax/compare/jax-v0.3.9...main).
+## jax 0.3.11 (Unreleased)
+* [GitHub commits](https://github.com/google/jax/compare/jax-v0.3.10...main).
 
-## jaxlib 0.3.8 (Unreleased)
-* [GitHub commits](https://github.com/google/jax/compare/jaxlib-v0.3.7...main).
+## jaxlib 0.3.11 (Unreleased)
+* [GitHub commits](https://github.com/google/jax/compare/jaxlib-v0.3.10...main).
+
+## jax 0.3.10 (May 3, 2022)
+* [GitHub commits](https://github.com/google/jax/compare/jax-v0.3.9...jax-v0.3.10).
+
+## jaxlib 0.3.10 (May 3, 2022)
+* [GitHub commits](https://github.com/google/jax/compare/jaxlib-v0.3.7...jaxlib-v0.3.10).
+* Changes
+  * [TF commit](https://github.com/tensorflow/tensorflow/commit/207d50d253e11c3a3430a700af478a1d524a779a)
+    fixes an issue in the MHLO canonicalizer that caused constant folding to
+    take a long time or crash for certain programs.
 
 ## jax 0.3.9 (May 2, 2022)
 * [GitHub commits](https://github.com/google/jax/compare/jax-v0.3.8...jax-v0.3.9).
diff --git a/WORKSPACE b/WORKSPACE
index 22ee094e7..fd66ebadc 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -7,10 +7,10 @@ load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 #    and update the sha256 with the result.
 http_archive(
     name = "org_tensorflow",
-    sha256 = "ef5a001226c37f59eca9c9bf0506b962cf220906bcc8f24df4d1ed6011a593e9",
-    strip_prefix = "tensorflow-e011187e26dbeed80dede66eb48729d184e3bb1d",
+    sha256 = "bee3000a7892008b32c65a5d6ba5d0a1f490e9e749d6789c5100d47068b1bbd8",
+    strip_prefix = "tensorflow-46cb56d405f0800e1741f68bba098d8eea5d0025",
     urls = [
-        "https://github.com/tensorflow/tensorflow/archive/e011187e26dbeed80dede66eb48729d184e3bb1d.tar.gz",
+        "https://github.com/tensorflow/tensorflow/archive/46cb56d405f0800e1741f68bba098d8eea5d0025.tar.gz",
     ],
 )
 
diff --git a/setup.py b/setup.py
index 3b70e1571..0be99846b 100644
--- a/setup.py
+++ b/setup.py
@@ -14,14 +14,14 @@
 
 from setuptools import setup, find_packages
 
-_current_jaxlib_version = '0.3.7'
+_current_jaxlib_version = '0.3.10'
 # The following should be updated with each new jaxlib release.
 _latest_jaxlib_version_on_pypi = '0.3.7'
 _available_cuda_versions = ['11']
 _default_cuda_version = '11'
 _available_cudnn_versions = ['82', '805']
 _default_cudnn_version = '82'
-_libtpu_version = '0.1.dev20220415'
+_libtpu_version = '0.1.dev20220503'
 
 _dct = {}
 with open('jax/version.py') as f: