Issue setting up-to-date environment with Catalyst

Catalyst
1

Hi everyone,

I would like to ask for you help about the catalyst initial setup. I am relatively new to Pennylane, and so far I have been using it either on its own, or with lightning plugins and jax workflows, with everything working great. But, now I wanted to start a new project with Catalyst. Initially, I had some troubles setting it up on my existing conda env, as it was clashing with a number of old dependencies, so I decided to set up a new env altogether since this is for a new project anyways.

Still, despite not being restricted to a specific version of, e.g., Pennylane or even Python, I have not been able to easily find a config that has every of my requirements working together. The only other restriction that I am imposing is that my env should contain the plugin for cuda 12.8 drivers (which somehow limits also the compatible jax version). I wanted to ask: is there a specific combination of dependencies that would allow me to run pennylane, cuda-enabled jax, and catalyst at the same time?

System: Ubuntu 24.04.2
CUDA version: 12.8
Python version: 3.12.11
Conda version: 25.5.1

Modules installed in env:
Package Version

absl-py 2.1.0
appdirs 1.4.4
asttokens 3.0.0
astunparse 1.6.3
autograd 1.8.0
autoray 0.7.2
Brotli 1.1.0
cachetools 6.2.1
certifi 2025.10.5
cffi 1.17.1
charset-normalizer 3.4.3
chex 0.1.89
comm 0.2.3
debugpy 1.8.17
decorator 5.2.1
diastatic-malt 2.15.2
dm-tree 0.1.7
etils 1.12.2
exceptiongroup 1.3.0
executing 2.2.1
gast 0.6.0
h2 4.3.0
hpack 4.1.0
hyperframe 6.1.0
idna 3.11
importlib_metadata 8.7.0
ipykernel 6.30.1
ipython 9.6.0
ipython_pygments_lexers 1.1.1
jax 0.4.28
jaxlib 0.4.28.dev20240710
jedi 0.19.2
jupyter_client 8.6.3
jupyter_core 5.8.1
matplotlib-inline 0.1.7
ml-dtypes 0.5.1
nest_asyncio 1.6.0
networkx 3.5
numpy 1.26.4
opt_einsum 3.4.0
optax 0.2.3
packaging 25.0
parso 0.8.5
PennyLane 0.37.0
PennyLane_Lightning 0.37.0
PennyLane_Lightning_GPU 0.37.0
pexpect 4.9.0
pickleshare 0.7.5
pip 25.2
platformdirs 4.5.0
prompt_toolkit 3.0.52
psutil 7.1.0
ptyprocess 0.7.0
pure_eval 0.2.3
pycparser 2.22
Pygments 2.19.2
PySocks 1.7.1
python-dateutil 2.9.0.post0
pyzmq 27.1.0
requests 2.32.5
rustworkx 0.17.1
scipy 1.16.2
scipy-openblas32 0.3.30.0.2
semantic-version 2.10.0
setuptools 80.9.0
six 1.17.0
stack_data 0.6.3
termcolor 3.1.0
toml 0.10.2
tomlkit 0.13.3
toolz 1.0.0
tornado 6.5.2
traitlets 5.14.3
typing_extensions 4.15.0
urllib3 2.5.0
wcwidth 0.2.14
wheel 0.45.1
zipp 3.23.0
zstandard 0.25.0

Thanks for anyone willing to give some hint!

2

Hi @salvozamm, thank you for giving Catalyst a spin!

Regarding your issue, it’s true that each Catalyst version relies on a specific JAX version because JAX internals frequently change. The dependency management for the jax and jaxlib packages should automatically be handled by Catalyst. At the moment, v0.12.0 of Catalyst requires both JAX libraries to be at v0.6.0.
Similarly, PennyLane and Catalyst are released in lockstep and should also always be compatible with each other when using their latest release versions (currently pennylane-catalyst==0.12.0 and pennylane==0.42.3, but older versions work too as long as the middle number matches).

The tricky part may the CUDA packages then. You mention that you are looking for cuda-enabled jax, but I don’t see any cuda-related packages in your list, am I correct in assuming you are talking about the jax cuda plugin (installed via jax[cuda12])? Are there any version constraints for you regarding the plugin (e.g. because you require a specific cuda driver version)? Have you tried installing the plugin after installing catalyst, and if so did you get any errors?

3

Hi @David_Ittah , thanks for the hints!

Indeed, at the time of writing the post, it was the optax installation that was breaking things, which would also clash with the installation of CUDA plugins. I should have specified that.

Still, as you suggested, I tried first getting catalyst, and, only after that, forcing the reinstall of the right version ofjax[cuda12] with

pip install –force-reinstall -v “jax[cuda12]=0.6.0”

but while that worked momentarily, it would break again when installing optax right after, as some dependencies would get updated to non-compatible versions.

What fixed it for me was installing everything together with

pip install –force-reinstall -v “jax[cuda12]=0.6.0” optax

I’m guessing that specifying everything to pip in one swoop makes it not override optax dependencies.

For anyone else that may come across this in the future, the order of installations was (but at this point, I guess installing everything in one command would do the job as well), on a fresh environment:

  1. install pennylane-catalyst
  2. install CUDA-enabled jax + optax
  3. install custatevec-cu12
  4. install pennylane-lightning-gpu

This should result in a pip list showing something like the following:

Package Version

absl-py 2.3.1
appdirs 1.4.4
asttokens 3.0.0
astunparse 1.6.3
autograd 1.8.0
autoray 0.7.2
cachetools 6.2.1
certifi 2025.10.5
charset-normalizer 3.4.4
chex 0.1.90
comm 0.2.3
custatevec-cu12 1.10.1
debugpy 1.8.16
decorator 5.2.1
diastatic-malt 2.15.2
exceptiongroup 1.3.0
executing 2.2.1
gast 0.6.0
idna 3.11
importlib_metadata 8.7.0
ipykernel 7.0.1
ipython 9.6.0
ipython_pygments_lexers 1.1.1
jax 0.6.0
jax-cuda12-pjrt 0.6.0
jax-cuda12-plugin 0.6.0
jaxlib 0.6.0
jedi 0.19.2
jupyter_client 8.6.3
jupyter_core 5.8.1
matplotlib-inline 0.1.7
ml_dtypes 0.5.3
nest_asyncio 1.6.0
networkx 3.5
numpy 2.3.3
nvidia-cublas-cu12 12.9.1.4
nvidia-cuda-cupti-cu12 12.9.79
nvidia-cuda-nvcc-cu12 12.9.86
nvidia-cuda-runtime-cu12 12.9.79
nvidia-cudnn-cu12 9.14.0.64
nvidia-cufft-cu12 11.4.1.4
nvidia-cusolver-cu12 11.7.5.82
nvidia-cusparse-cu12 12.5.10.65
nvidia-nccl-cu12 2.28.3
nvidia-nvjitlink-cu12 12.9.86
opt_einsum 3.4.0
optax 0.2.6
packaging 25.0
parso 0.8.5
pennylane 0.42.3
pennylane_catalyst 0.12.0
pennylane_lightning 0.42.0
pennylane_lightning_gpu 0.42.0
pexpect 4.9.0
pickleshare 0.7.5
pip 25.2
platformdirs 4.5.0
prompt_toolkit 3.0.52
psutil 7.0.0
ptyprocess 0.7.0
pure_eval 0.2.3
Pygments 2.19.2
python-dateutil 2.9.0.post0
pyzmq 27.1.0
requests 2.32.5
rustworkx 0.17.1
scipy 1.16.2
scipy-openblas32 0.3.30.0.2
setuptools 80.9.0
six 1.17.0
stack_data 0.6.3
termcolor 3.1.0
tomlkit 0.13.3
toolz 1.0.0
tornado 6.5.1
traitlets 5.14.3
typing_extensions 4.15.0
urllib3 2.5.0
wcwidth 0.2.14
wheel 0.45.1
zipp 3.23.0

Thanks again for the help!

4

Great, I’m glad you could work out the problem!

I’m guessing that specifying everything to pip in one swoop makes it not override optax dependencies.

Yes, providing multiple packages to pip at once will it try and satisfy all their dependencies together, so that can be helpful in tricky situations :slight_smile:

Out of curiosity, why did you use the –force-reinstall flag? I think that might also lead to upgrading dependencies that you’ve previously pinned to a particular version (see pip --force-reinstall - pip documentation v25.2).

5

Indeed, good to know!

No particular reason. I was not aware of the default behavior you mentioned, and was only referencing the pip install -h manual, where the flag is only described as Reinstall all packages even if they are already up-to-date. and therefore I though it was yet another way to make sure that re-installation indeed occurs.

Thanks again for the tips!

1 Like