Default.qubit.jax not working with newer Jax version?

Dear Pennylane Team,

I’am currentlx trying to combine pennylane with Jax to speedup my computions.
I have Pennylane==0.33., Jax==0.4.34 installed.
When trying to initialize a device:

dev_qnpu_state = qml.device(“default.qubit.jax”, wires=n_wires_qnpu_state, shots=None)

I get the error:
ImportError: default.qubit.jax device requires installing jax>0.3.20
which I clearly have. Trying older versions of jax also didn’t resolve the issue for me. I tested that the correct version is loaded in the code with print(jax.version).
Do you know this issue, or could you reccomend a jax version that works for sure?

Thanks and best regards,
Pia

Hi @Pia !

To me it looks like an environment issue. I would recommend creating a new environment if possible.

On the other hand, is there a reason why you need to use PennyLane 0.33? I generally recommend to stay up to date with the latest version, which is v0.38 at the moment.

Below are the steps you can follow to create a new environment with the latest version of PennyLane. You will need to adapt it to install Jax too, in addition to the specific version of PennyLane in case you prefer not to use the latest one.

Let me know if this works for you. If it doesn’t please post the output of qml.about()
I hope this helps!


You can create a virtual environment with Conda and install PennyLane as follows:

  1. Install Miniconda following the instructions here.
  2. Open your terminal (mac) or command line (Windows).
  3. Create a new Conda environment with: conda create --name <name_of_your_environment> python=3.10
  4. Activate the environment with: conda activate <name_of_your_environment>
  5. Install PennyLane with: python -m pip install pennylane
  6. Install other useful packages with: python -m pip install jupyter matplotlib

Note that you will be installing 3 packages here: PennyLane, Jupyter, and Matplotlib. Also, note that where it says <name_of_your_environment> you can choose any name that you want.

Thank you :slight_smile: that helped and it runs now on my cluster.
I have two follow-up questions:

  1. I would like to run a large-circuit (approx. 30 qubit) variational quantum algorithm.
    There are various backends discussed that are supposed to speed-up the simulations or make those large runs possible as e.g. bluequbit or Jax.
    Is there anything you can reccomend for training of a large quantum circuit?

  2. I further tried to run the first example using bluequbit ( Using the BlueQubit (CPU) device with PennyLane | PennyLane Demos). I have installed bluequbit, and have my personalized token. When I call qml.about, the bluequbit device shows. However, when I try to initialize a device with

bq_dev = qml.device(“bluequbit.cpu”, wires=2, token=“myToken”)

I get the following error, that I don’t understand:

Traceback (most recent call last):
File “”, line 1, in
File “/data/homes/psiegl/.conda/envs/QcOpt/lib/python3.12/site-packages/pennylane/init.py”, line 346, in device
plugin_device_class = plugin_devices[name].load()
^^^^^^^^^^^^^^^^^^^^^^^^^^^
File “/data/homes/psiegl/.conda/envs/QcOpt/lib/python3.12/site-packages/pkg_resources/init.py”, line 2781, in load
return self.resolve()
^^^^^^^^^^^^^^
File “/data/homes/psiegl/.conda/envs/QcOpt/lib/python3.12/site-packages/pkg_resources/init.py”, line 2787, in resolve
module = import(self.module_name, fromlist=[‘name’], level=0)
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File “/data/homes/psiegl/.conda/envs/QcOpt/lib/python3.12/site-packages/bluequbit/pennylane_plugin/init.py”, line 1, in
from .bluequbit_cpu import BluequbitCPU # noqa: F401
^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
File “/data/homes/psiegl/.conda/envs/QcOpt/lib/python3.12/site-packages/bluequbit/pennylane_plugin/bluequbit_cpu.py”, line 1, in
from .bluequbit_device import BluequbitDevice
File “/data/homes/psiegl/.conda/envs/QcOpt/lib/python3.12/site-packages/bluequbit/pennylane_plugin/bluequbit_device.py”, line 17, in
from pennylane.devices import QubitDevice
ImportError: cannot import name ‘QubitDevice’ from ‘pennylane.devices’ (/data/homes/psiegl/.conda/envs/QcOpt/lib/python3.12/site-packages/pennylane/devices/init.py)

I would be happy if you have an idea what the issue might be.

Best regards,
Pia

Hi @Pia ,

  1. 30 qubits is a lot. It would require over 17 GiB of RAM just to store a state, and even more to do computations. It’s good that you have a cluster but I would still suggest that you start with slightly lower numbers and then see how much you can increase it.
    If you’re running on CPUs then you’d probably want to use lightning.qubit. You can try using the default.tensor device for tensor network simulations. I encourage you to check out the simulators info in our Performance page, and see our demos focused on performance in the News and announcements section of that page.

  2. Can you please post the output of qml.about() ? This looks like an installation issue but we may be able to confirm with the qml.about() output.

I hope this helps!

Dear Catalina,

I will have a look at those devices! As far as I see here ( Introducing Catalyst: quantum just-in-time compilation | PennyLane Blog) they are even compatible with jax :slight_smile:

Here is the output of the qml.about() call:

Name: PennyLane
Version: 0.33.1
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: GitHub - PennyLaneAI/pennylane: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Author:
Author-email:
License: Apache License 2.0
Location: /pathname/lib/python3.12/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Lightning

Platform info: Linux-5.15.0-122-generic-x86_64-with-glibc2.35
Python version: 3.12.5
Numpy version: 2.0.2
Scipy version: 1.14.1
Installed devices:

  • default.gaussian (PennyLane-0.33.1)
  • default.mixed (PennyLane-0.33.1)
  • default.qubit (PennyLane-0.33.1)
  • default.qubit.autograd (PennyLane-0.33.1)
  • default.qubit.jax (PennyLane-0.33.1)
  • default.qubit.legacy (PennyLane-0.33.1)
  • default.qubit.tf (PennyLane-0.33.1)
  • default.qubit.torch (PennyLane-0.33.1)
  • default.qutrit (PennyLane-0.33.1)
  • null.qubit (PennyLane-0.33.1)
  • lightning.qubit (PennyLane-Lightning-0.34.0)
  • bluequbit.cpu (bluequbit-0.9.3b1)

Thanks already,
Best,
Pia

Hi @Pia,

I see that there’s a mismatch between your versions of PennyLane and PennyLane-Lightning. Are you able to create a new virtual environment where you install PennyLane v0.40?

You should be able to install or upgrade PennyLane with python -m pip install pennylane --upgrade

Lightning-Qubit comes pre-installed with PennyLane so you shouldn’t need to install it separately.

Let us know if this works for you or if you have any issues.