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