The script inference_pipeline.ipynb not work! ImportError: cannot import name 'partitioning'

[rmpku]

The output errors

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.) Traceback (most recent call last): File "/workspace/sd_a100/unlearning_datasets/unlearning_dataset/dalle-mini/tools/inference/inference.py", line 22, in from dalle_mini import DalleBart, DalleBartProcessor File "/home/mr/anaconda3/envs/dalle_mini/lib/python3.6/site-packages/dalle_mini/init.py", line 3, in from .model import DalleBart, DalleBartProcessor File "/home/mr/anaconda3/envs/dalle_mini/lib/python3.6/site-packages/dalle_mini/model/init.py", line 2, in from .modeling import DalleBart File "/home/mr/anaconda3/envs/dalle_mini/lib/python3.6/site-packages/dalle_mini/model/modeling.py", line 28, in from flax.linen import partitioning as nn_partitioning ImportError: cannot import name 'partitioning'

The env version show as

_libgcc_mutex 0.1 main
_openmp_mutex 5.1 1_gnu
absl-py 1.4.0 pypi_0 pypi appdirs 1.4.4 pypi_0 pypi backcall 0.2.0 pypi_0 pypi ca-certificates 2023.05.30 h06a4308_0
certifi 2021.5.30 py36h06a4308_0
charset-normalizer 2.0.12 pypi_0 pypi chex 0.0.9 pypi_0 pypi click 8.0.4 pypi_0 pypi cycler 0.11.0 pypi_0 pypi dalle-mini 0.1.1 pypi_0 pypi dataclasses 0.8 pypi_0 pypi decorator 5.1.1 pypi_0 pypi dm-tree 0.1.8 pypi_0 pypi docker-pycreds 0.4.0 pypi_0 pypi einops 0.4.1 pypi_0 pypi emoji 2.7.0 pypi_0 pypi entrypoints 0.4 pypi_0 pypi filelock 3.4.1 pypi_0 pypi flatbuffers 23.5.26 pypi_0 pypi flax 0.3.5 pypi_0 pypi ftfy 6.0.3 pypi_0 pypi gitdb 4.0.9 pypi_0 pypi gitpython 3.1.18 pypi_0 pypi huggingface-hub 0.4.0 pypi_0 pypi idna 3.4 pypi_0 pypi importlib-metadata 4.8.3 pypi_0 pypi importlib-resources 5.4.0 pypi_0 pypi ipykernel 5.5.6 pypi_0 pypi ipython 7.16.3 pypi_0 pypi ipython-genutils 0.2.0 pypi_0 pypi jax 0.2.17 pypi_0 pypi jaxlib 0.1.65 pypi_0 pypi jedi 0.17.2 pypi_0 pypi joblib 1.1.1 pypi_0 pypi jupyter-client 7.1.2 pypi_0 pypi jupyter-core 4.9.2 pypi_0 pypi kiwisolver 1.3.1 pypi_0 pypi libedit 3.1.20221030 h5eee18b_0
libffi 3.2.1 hf484d3e_1007
libgcc-ng 11.2.0 h1234567_1
libgomp 11.2.0 h1234567_1
libstdcxx-ng 11.2.0 h1234567_1
matplotlib 3.3.4 pypi_0 pypi msgpack 1.0.5 pypi_0 pypi ncurses 6.4 h6a678d5_0
nest-asyncio 1.5.7 pypi_0 pypi numpy 1.19.5 pypi_0 pypi openssl 1.0.2u h7b6447c_0
opt-einsum 3.3.0 pypi_0 pypi optax 0.0.91 pypi_0 pypi packaging 21.3 pypi_0 pypi parso 0.7.1 pypi_0 pypi pathtools 0.1.2 pypi_0 pypi pexpect 4.8.0 pypi_0 pypi pickleshare 0.7.5 pypi_0 pypi pillow 8.4.0 pypi_0 pypi pip 21.2.2 py36h06a4308_0
promise 2.3 pypi_0 pypi prompt-toolkit 3.0.36 pypi_0 pypi protobuf 3.19.6 pypi_0 pypi psutil 5.9.5 pypi_0 pypi ptyprocess 0.7.0 pypi_0 pypi pygments 2.14.0 pypi_0 pypi pyparsing 3.0.7 pypi_0 pypi python 3.6.5 hc3d631a_2
python-dateutil 2.8.2 pypi_0 pypi pyyaml 6.0.1 pypi_0 pypi pyzmq 25.1.0 pypi_0 pypi readline 7.0 h7b6447c_5
regex 2023.8.8 pypi_0 pypi requests 2.27.1 pypi_0 pypi sacremoses 0.0.53 pypi_0 pypi scipy 1.5.4 pypi_0 pypi sentry-sdk 1.29.2 pypi_0 pypi setproctitle 1.2.3 pypi_0 pypi setuptools 58.0.4 py36h06a4308_0
shortuuid 1.0.1 pypi_0 pypi six 1.16.0 pypi_0 pypi smmap 5.0.0 pypi_0 pypi sqlite 3.33.0 h62c20be_0
tk 8.6.12 h1ccaba5_0
tokenizers 0.12.1 pypi_0 pypi toolz 0.12.0 pypi_0 pypi tornado 6.1 pypi_0 pypi tqdm 4.64.1 pypi_0 pypi traitlets 4.3.3 pypi_0 pypi transformers 4.18.0 pypi_0 pypi typing-extensions 3.10.0.2 pypi_0 pypi unidecode 1.3.6 pypi_0 pypi urllib3 1.26.16 pypi_0 pypi vqgan-jax 0.0.1 pypi_0 pypi wandb 0.13.6 pypi_0 pypi wcwidth 0.2.6 pypi_0 pypi wheel 0.37.1 pyhd3eb1b0_0
xz 5.4.2 h5eee18b_0
zipp 3.6.0 pypi_0 pypi zlib 1.2.13 h5eee18b_0

[borisdayma]

This should now work