ModuleNotFoundError: No module named 'flax.deprecated'

[DaShenZi721]

Hello! Thanks for offering this benchmark! However, when I use the following commands

PYTHONPATH="$(pwd)":"$PYTHON_PATH" python lra_benchmarks/listops/train.py \
      --config=lra_benchmarks/listops/configs/transformer_base.py \
      --model_dir=/tmp/listops \
      --task_name=basic \
      --data_dir=$HOME/lra_data/listops/

I got this message

Traceback (most recent call last):
  File "lra_benchmarks/listops/train.py", line 28, in <module>
    from flax.deprecated import nn
ModuleNotFoundError: No module named 'flax.deprecated'

Is there any solutions?

[DaShenZi721]

This is my setting:

• cuda: 11.2
• jax: 0.2.13
• jaxlib: 0.1.65+cuda112
• flax: 0.2.2
• tensorflow: 2.7.0

[DaShenZi721]

I try to replace from flax.deprecated import nn with from flax import nn, and it works!

[AlexKay28]

@DaShenZi721 which python version did you use?