Drop tensorstore dependency; require it for tests only

tensorstore is an optional dependency for flax. In google#2520 it was added
back as the normal dependency, but on some Linux and macOS environments
tensorstore still fails to build (see google#2341, google/tensorstore#63).

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