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",
 ]