Prepare for jax and jaxlib 0.4.4 release

PiperOrigin-RevId: 510152471
abaelhe · Feb 16, 2023 · 58e46b4 · 58e46b4
1 parent c6a99b6
commit 58e46b4
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git 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 = "8d85c2a03b60fdb7152e810ac9e7380f9ebd677f382f7ea93f2e88af6f23ed26",
-    strip_prefix = "tensorflow-5e52982c02c0f5d88eaa449408abe3bdcac9ee05",
+    sha256 = "23aae276b2705bfbdaea3c472da24130598f13ac0439cfb7149befb781d97a8f",
+    strip_prefix = "tensorflow-43e9d313548ded301fa54f25a4192d3bcb123330",
     urls = [
-        "https://github.com/tensorflow/tensorflow/archive/5e52982c02c0f5d88eaa449408abe3bdcac9ee05.tar.gz",
+        "https://github.com/tensorflow/tensorflow/archive/43e9d313548ded301fa54f25a4192d3bcb123330.tar.gz",
     ],
 )
 

diff --git a/jax/tools/colab_tpu.py b/jax/tools/colab_tpu.py
@@ -22,7 +22,7 @@
 TPU_DRIVER_MODE = 0
 
 
-def setup_tpu(tpu_driver_version='tpu_driver_20230207'):
+def setup_tpu(tpu_driver_version='tpu_driver_20230216'):
   """Sets up Colab to run on TPU.
 
   Note: make sure the Colab Runtime is set to Accelerator: TPU.

diff --git a/setup.py b/setup.py
@@ -19,14 +19,14 @@
 
 from setuptools import setup, find_packages
 
-_current_jaxlib_version = '0.4.3'
+_current_jaxlib_version = '0.4.4'
 # The following should be updated with each new jaxlib release.
 _latest_jaxlib_version_on_pypi = '0.4.3'
 _available_cuda_versions = ['11']
 _default_cuda_version = '11'
 _available_cudnn_versions = ['82', '86']
 _default_cudnn_version = '86'
-_libtpu_version = '0.1.dev20230207'
+_libtpu_version = '0.1.dev20230216'
 
 _dct = {}
 with open('jax/version.py', encoding='utf-8') as f: