google-deepmind / alphageometry

Apache License 2.0
4.18k stars 469 forks source link

cannot import name 'get_tensorstore_spec' from 'jax.experimental.array_serialization #98

Open XEROXyuyuyuyu opened 7 months ago

XEROXyuyuyuyu commented 7 months ago

Traceback (most recent call last): File "C:\users\lenovo\alphageometry\lm_inference_test.py", line 22, in import lm_inference as lm File "C:\users\lenovo\alphageometry\lm_inference.py", line 22, in from transformer import inference_utils File "C:\users\lenovo\alphageometry\$MELIAD_PATH\transformer\inference_utils.py", line 26, in import training_loop File "C:\users\lenovo\alphageometry\$MELIAD_PATH\training_loop.py", line 27, in from flax.training import checkpoints File "C:\Users\lenovo\alphageometry\venv\lib\site-packages\flax\training\checkpoints.py", line 33, in from jax.experimental.array_serialization import get_tensorstore_spec ImportError: cannot import name 'get_tensorstore_spec' from 'jax.experimental.array_serialization' (C:\Users\lenovo\alphageometry\venv\lib\site-packages\jax\experimental\array_serialization__init__.py)

It shows that "get_tensorstore_spec" doesn't exist. What should I do?