From 10639c9b1461204ec1d529e086f571850cd52da3 Mon Sep 17 00:00:00 2001
From: Iurii Kemaev <iukemaev@google.com>
Date: Thu, 15 Apr 2021 00:53:56 -0700
Subject: [PATCH] Update jax version in requirements.

PiperOrigin-RevId: 368589236
---
 requirements/requirements.txt | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/requirements/requirements.txt b/requirements/requirements.txt
index 78a865e3..fd452627 100644
--- a/requirements/requirements.txt
+++ b/requirements/requirements.txt
@@ -1,7 +1,6 @@
 absl-py>=0.9.0
 chex>=0.0.5
-# Change the sign once TFP is fixed
-jax<=0.2.11
+jax>=0.2.11
 jaxlib>=0.1.37
 numpy>=1.18.0
-tensorflow-probability>=0.12.1
+tensorflow-probability>=0.12.2