ImportError: cannot import name 'masking' from 'jax.interpreters'

[csxrzhang]

Traceback (most recent call last):
File "/mnt/data1/zhangxianrong/d3pm/text/main.py", line 35, in
from text import diffusion # pylint: disable=unused-import
File "/mnt/data1/zhangxianrong/d3pm/text/diffusion.py", line 38, in
from text import models
File "/mnt/data1/zhangxianrong/d3pm/text/models.py", line 33, in
from flaxformer.architectures.t5 import t5_architecture
File "/opt/miniconda3/envs/python3.9-zxr/lib/python3.9/site-packages/flaxformer/architectures/t5/t5_architecture.py", line 35, in
from flaxformer.components import rich_attention_position_scores
File "/opt/miniconda3/envs/python3.9-zxr/lib/python3.9/site-packages/flaxformer/components/rich_attention_position_scores.py", line 30, in
from flaxformer.components import dense
File "/opt/miniconda3/envs/python3.9-zxr/lib/python3.9/site-packages/flaxformer/components/dense.py", line 23, in
from aqt.jax_legacy.jax import flax_layers as aqt_flax_layers
File "/opt/miniconda3/envs/python3.9-zxr/lib/python3.9/site-packages/aqt/jax_legacy/jax/flax_layers.py", line 26, in
from aqt.jax_legacy.jax import compute_cost_utils
File "/opt/miniconda3/envs/python3.9-zxr/lib/python3.9/site-packages/aqt/jax_legacy/jax/compute_cost_utils.py", line 27, in
from jax.interpreters import masking
ImportError: cannot import name 'masking' from 'jax.interpreters' (/opt/miniconda3/envs/python3.9-zxr/lib/python3.9/site-packages/jax/interpreters/init.p