cannot import name 'masking' from 'jax.interpreters' error in jax >= 0.3.16

[cgarciae]

In jax>=0.3.16 I am getting this error when importing vit_jax:

ImportError                               Traceback (most recent call last)
/data/cristian/flax/docs/guides/transfer_learning.ipynb Cell 3 in <cell line: 1>()
----> 1 from vit_jax import models, checkpoint
      2 from pathlib import Path
      3 import tensorflow as tf

File /data/cristian/flax/.venv/lib/python3.8/site-packages/vit_jax/models.py:15, in <module>
      1 # Copyright 2022 Google LLC.
      2 #
      3 # Licensed under the Apache License, Version 2.0 (the "License");
   (...)
     12 # See the License for the specific language governing permissions and
     13 # limitations under the License.
---> 15 from vit_jax import models_lit
     16 from vit_jax import models_mixer
     17 from vit_jax import models_vit

File /data/cristian/flax/.venv/lib/python3.8/site-packages/vit_jax/models_lit.py:31, in <module>
     28 from vit_jax import models_vit
     29 from vit_jax import preprocess
---> 31 from flaxformer.architectures.bert import bert
     32 from flaxformer.architectures.bert import configs
     35 BASE_PATH = 'gs://vit_models/lit'
...
---> 27 from jax.interpreters import masking
     28 from jax.interpreters import mlir
     29 import numpy as onp

ImportError: cannot import name 'masking' from 'jax.interpreters' (/data/cristian/flax/.venv/lib/python3.8/site-packages/jax/interpreters/__init__.py)

Seems masking was removed from jax.interpreters in 0.3.16. flax>=0.6.0 now depends on jax>=0.3.16 so vit_jax is failing with the latest Flax version.

[andsteing]

Should be fixed with #231

Can you confirm?

[cgarciae]

#231 should fix it for regular users, however this use case is for a guide in the flax repo so vit_jax would have to work with the latest version.

[andsteing]

created new #232 for that