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

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

This should now work

I think it doesn't work again in 2025:

My logs when running the first part of the code:

# Required only for colab environments + GPU
!pip install jax==0.3.25 jaxlib==0.3.25 -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Install required libraries
!pip install -q dalle-mini
!pip install -q git+https://github.com/patil-suraj/vqgan-jax.git

Looking in links: https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
Collecting jax==0.3.25
  Downloading jax-0.3.25.tar.gz (1.1 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 1.1/1.1 MB 21.2 MB/s eta 0:00:00
  Preparing metadata (setup.py) ... done
ERROR: Ignored the following yanked versions: 0.4.32
ERROR: Could not find a version that satisfies the requirement jaxlib==0.3.25 (from versions: 0.4.17, 0.4.17+cuda11.cudnn86, 0.4.17+cuda12.cudnn89, 0.4.18, 0.4.18+cuda11.cudnn86, 0.4.18+cuda12.cudnn89, 0.4.19, 0.4.19+cuda11.cudnn86, 0.4.19+cuda12.cudnn89, 0.4.20, 0.4.20+cuda11.cudnn86, 0.4.20+cuda12.cudnn89, 0.4.21, 0.4.21+cuda11.cudnn86, 0.4.21+cuda12.cudnn89, 0.4.22, 0.4.22+cuda11.cudnn86, 0.4.22+cuda12.cudnn89, 0.4.23, 0.4.23+cuda11.cudnn86, 0.4.23+cuda12.cudnn89, 0.4.24, 0.4.24+cuda11.cudnn86, 0.4.24+cuda12.cudnn89, 0.4.25, 0.4.25+cuda11.cudnn86, 0.4.25+cuda12.cudnn89, 0.4.26, 0.4.26+cuda12.cudnn89, 0.4.27, 0.4.27+cuda12.cudnn89, 0.4.28, 0.4.28+cuda12.cudnn89, 0.4.29, 0.4.29+cuda12.cudnn91, 0.4.30, 0.4.31, 0.4.33, 0.4.34, 0.4.35, 0.4.36, 0.4.38, 0.5.0, 0.5.1, 0.5.3, 0.6.0, 0.6.1, 0.6.2, 0.7.0, 0.7.1, 0.7.2, 0.8.0)
ERROR: No matching distribution found for jaxlib==0.3.25
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 93.9/93.9 kB 5.2 MB/s eta 0:00:00
  Preparing metadata (setup.py) ... done
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 314.9/314.9 kB 19.5 MB/s eta 0:00:00
  Installing build dependencies ... done
  Getting requirements to build wheel ... done
  Preparing metadata (pyproject.toml) ... done
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 197.4/197.4 kB 12.4 MB/s eta 0:00:00
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 66.6/66.6 kB 3.2 MB/s eta 0:00:00
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 5.8/5.8 MB 77.0 MB/s eta 0:00:00
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 608.4/608.4 kB 24.1 MB/s eta 0:00:00
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 44.8/44.8 kB 2.9 MB/s eta 0:00:00
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 235.8/235.8 kB 15.2 MB/s eta 0:00:00
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 154.1/154.1 kB 14.9 MB/s eta 0:00:00
   ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 87.9/87.9 kB 6.6 MB/s eta 0:00:00
  Building wheel for jax (setup.py) ... done
  error: subprocess-exited-with-error
  
  × Building wheel for tokenizers (pyproject.toml) did not run successfully.
  │ exit code: 1
  ╰─> See above for output.
  
  note: This error originates from a subprocess, and is likely not a problem with pip.
  Building wheel for tokenizers (pyproject.toml) ... error
  ERROR: Failed building wheel for tokenizers
ERROR: ERROR: Failed to build installable wheels for some pyproject.toml based projects (tokenizers)
  Preparing metadata (setup.py) ... done
  Building wheel for vqgan-jax (setup.py) ... done

Navigation Menu

Repository navigation

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

The output errors

The env version show as

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

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

Description

The output errors

The env version show as

Activity

borisdayma commented on Aug 22, 2023

empeje commented on Oct 18, 2025

Metadata

Metadata

Assignees

Labels

Projects

Milestone

Relationships

Development

Issue actions