From 19224093f4161273709a2a2687a538a6cdee3e68 Mon Sep 17 00:00:00 2001
From: Jongwook Choi <wookayin@gmail.com>
Date: Tue, 6 Dec 2022 21:14:49 -0500
Subject: [PATCH] Drop tensorstore dependency; require it for tests only

tensorstore is an optional dependency for flax. In #2520 it was added
back as the normal dependency, but on some Linux and macOS environments
tensorstore still fails to build (see #2341, google/tensorstore#63).
---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 7d2a9af16b..1fb732954b 100644
--- a/setup.py
+++ b/setup.py
@@ -31,7 +31,6 @@
     "msgpack",
     "optax",
     "orbax",
-    "tensorstore",
     "rich>=11.1",
     "typing_extensions>=4.1.1",
     "PyYAML>=5.4.1",
@@ -54,6 +53,7 @@
     "tensorflow_text>=2.4.0",  # WMT example.
     "tensorflow_datasets",
     "tensorflow",
+    "tensorstore",
     "torch",
 ]