From e57709bf08ce4d0a0f77b1e4f5afd72acf90fdad Mon Sep 17 00:00:00 2001
From: Chandra Devarakonda <chandrasekhard@google.com>
Date: Tue, 17 Oct 2023 00:15:45 +0000
Subject: [PATCH] Update Pax 1.2.0 requirements

---
 paxml/pip_package/requirements.txt | 4 ++--
 requirements.in                    | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/paxml/pip_package/requirements.txt b/paxml/pip_package/requirements.txt
index 91a6aba49..28a9857b2 100644
--- a/paxml/pip_package/requirements.txt
+++ b/paxml/pip_package/requirements.txt
@@ -124,7 +124,7 @@ executing==2.0.0
     # via stack-data
 fastjsonschema==2.18.1
     # via nbformat
-fiddle @ git+https://github.com/google/fiddle
+fiddle==0.2.11
     # via -r praxis-requirements.in
 filelock==3.12.4
     # via
@@ -206,7 +206,7 @@ ipywidgets==8.1.1
     # via jupyter
 isoduration==20.11.0
     # via jsonschema
-jax @ git+https://github.com/google/jax
+jax==0.4.18
     # via
     #   -r paxml-requirements.in
     #   -r praxis-requirements.in
diff --git a/requirements.in b/requirements.in
index 7c3da59fc..a75acc979 100644
--- a/requirements.in
+++ b/requirements.in
@@ -6,7 +6,7 @@ clu
 etils
 flax
 graphviz
-jax @ git+https://github.com/google/jax
+jax==0.4.18
 lingvo
 numpy
 orbax-checkpoint
@@ -20,4 +20,4 @@ tensorflow-text~=2.9.0
 tensorstore
 tensorflow-datasets==4.8.3
 tfds-nightly==4.8.3.dev202303280045
-tensorflow-metadata==1.12.0
\ No newline at end of file
+tensorflow-metadata==1.12.0