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

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?