Pennylane Lightning NotImplementedError: VJP of eigvals wrt argnums (0,) not defined

Hi all,

I’m having troubles with the lightning.qubit object of Pennylane. From my understanding, it should behave pretty much like the default.qubit object, but with a C++ backend. However, this code works with the default.qubit object and not the lightning.qubit, and furthermore I get a weird error message.

To explain exactly what I’m trying to do, I just want to define a circuit that we can apply some layers to, and then, with an optimizer, find the weights that minimize a given Hamiltonian. The problem arrises when I use an arbitrary unitary matrix for the transformation in my layer instead of a rotation. I also want to point out that the circuit with layerXX works fine if it is not executing in the optimizer.

import pennylane as qml
from pennylane import numpy as np

# Define layers
def layerR(W):
    qml.Rot(W[0, 0], W[0, 1], W[0, 2], wires=0)
    qml.Rot(W[1, 0], W[1, 1], W[1, 2], wires=1)
    qml.Rot(W[2, 0], W[2, 1], W[2, 2], wires=2)
    qml.Rot(W[3, 0], W[3, 1], W[3, 2], wires=3)

    qml.CNOT(wires=[0, 1])
    qml.CNOT(wires=[1, 2])
    qml.CNOT(wires=[2, 3])
    qml.CNOT(wires=[3, 0])


def layerXX(W):
    # circX() is a function that returns a unitary matrix as a 4x4 numpy array
    qml.QubitUnitary(circX(W[0, 0], W[0, 1], W[0, 2]), wires=[0, 1])
    qml.QubitUnitary(circX(W[0, 0], W[0, 1], W[0, 2]), wires=[2, 3])

    qml.SQISW(wires=[1, 2])
    qml.SQISW(wires=[0, 3])

# Define quantum circuits
def circuit(layer, hamiltonian, weights):
    qml.BasisState([0, 0, 0, 0], wires=[0, 1, 2, 3])

    for w in weights:
        layer(w)

    return qml.expval(hamiltonian)

ham4terms = qml.PauliZ(0) @ qml.PauliZ(1) @ qml.PauliZ(2) @ qml.PauliZ(3)
tfield = 1.0 * (qml.PauliX(0) + qml.PauliX(1) + qml.PauliX(2) + qml.PauliX(3))
H = -ham4terms + tfield

def experience(label, layer, H, num_qubits, num_layers, scale, *args):
    print(label + ":")

    # Replace by dev = qml.device("default.qubit", wires=4) if you want it to work
    dev = qml.device("lightning.qubit", wires=4)

    @qml.qnode(dev, interface="autograd")
    def current_circuit(weights):
        return circuit(layer, H, weights)

    # Initialize weights
    np.random.seed(0)
    weights_init = scale * np.random.randn(
        num_layers, num_qubits, 3, requires_grad=True
    )
    # bias_init = np.array(0.0, requires_grad=True)

    opt = qml.NesterovMomentumOptimizer(args[0], args[1])

    weights = weights_init

    for it in range(40):
        weights = opt.step(current_circuit, weights)

        if (it + 1) % 10 == 0:
            print(
                "Iter: {:5d} | Energy: {:0.7f}".format(it + 1, current_circuit(weights))
            )

    print("")

if __name__ == "__main__":
    experience("Circuit R", layerR, H, 4, 1, 0.5, 0.2, 0.1)
    experience("Circuit XX", layerXX, H, 2, 2, 0.5, 0.2, 0.1)

The error stack is:

Traceback (most recent call last):
  File "autograd\core.py", line 31, in __init__
    vjpmaker = primitive_vjps[fun]
KeyError: <function primitive.<locals>.f_wrapped at 0x000002114435A040>

During handling of the above exception, another exception occurred:
  [...]
  File "pennylane\optimize\gradient_descent.py", line 93, in step
    g, _ = self.compute_grad(objective_fn, args, kwargs, grad_fn=grad_fn)
  File "pennylane\optimize\nesterov_momentum.py", line 76, in compute_grad
    grad = g(*shifted_args, **kwargs)
  File "pennylane\_grad.py", line 165, in __call__
    grad_value, ans = grad_fn(*args, **kwargs)  # pylint: disable=not-callable
  File "autograd\wrap_util.py", line 20, in nary_f
    return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
  File "pennylane\_grad.py", line 183, in _grad_with_forward
    vjp, ans = _make_vjp(fun, x)  # pylint: disable=redefined-outer-name
  File "autograd\core.py", line 10, in make_vjp
    end_value, end_node =  trace(start_node, fun, x)
  File "autograd\tracer.py", line 10, in trace
    end_box = fun(start_box)
  File "autograd\wrap_util.py", line 15, in unary_f
    return fun(*subargs, **kwargs)
  File "pennylane\workflow\qnode.py", line 1095, in __call__
    self._update_gradient_fn(shots=override_shots, tape=self._tape)
  File "pennylane\workflow\qnode.py", line 604, in _update_gradient_fn
    self.gradient_fn, self.gradient_kwargs, self.device = self.get_gradient_fn(
  File "pennylane\workflow\qnode.py", line 650, in get_gradient_fn
    if device.supports_derivatives(config, circuit=tape):
  File "pennylane_lightning\lightning_qubit\lightning_qubit.py", line 627, in supports_derivatives
    return _supports_adjoint(circuit=circuit)
  File "pennylane_lightning\lightning_qubit\lightning_qubit.py", line 363, in _supports_adjoint
    prog((circuit,))
  File "pennylane\transforms\core\transform_program.py", line 509, in __call__
    new_tapes, fn = transform(tape, *targs, **tkwargs)
  File "pennylane\devices\preprocess.py", line 340, in decompose
    new_ops = [
  File "pennylane\devices\preprocess.py", line 343, in <listcomp>
    for final_op in _operator_decomposition_gen(
  File "pennylane\devices\preprocess.py", line 62, in _operator_decomposition_gen
    decomp = decomposer(op)
  File "pennylane\devices\preprocess.py", line 328, in decomposer
    return op.decomposition()
  File "pennylane\operation.py", line 1285, in decomposition
    return self.compute_decomposition(
  File "pennylane\ops\qubit\matrix_ops.py", line 223, in compute_decomposition
    return qml.ops.two_qubit_decomposition(U, Wires(wires))
  File "pennylane\ops\op_math\decompositions\two_qubit_unitary.py", line 615, in two_qubit_decomposition
    num_cnots = _compute_num_cnots(U)
  File "pennylane\ops\op_math\decompositions\two_qubit_unitary.py", line 126, in _compute_num_cnots
    evs = math.linalg.eigvals(gammaU)
  File "autoray.py", line 81, in do
    return func(*args, **kwargs)
  File "pennylane\numpy\wrapper.py", line 117, in _wrapped
    res = obj(*args, **kwargs)
  File "autograd\tracer.py", line 45, in f_wrapped
    node = node_constructor(ans, f_wrapped, argvals, kwargs, argnums, parents)
  File "autograd\core.py", line 34, in __init__
    raise NotImplementedError("VJP of {} wrt argnums {} not defined"
NotImplementedError: VJP of eigvals wrt argnums (0,) not defined

Here is the full package version:

Name: PennyLane
Version: 0.36.0
Summary: 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.
Home-page: https://github.com/PennyLaneAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: ###

Platform info:           Windows-10-10.0.19045-SP0
Python version:          3.9.19
Numpy version:           1.21.5
Scipy version:           1.10.1
Installed devices:
- default.clifford (PennyLane-0.36.0)
- default.gaussian (PennyLane-0.36.0)
- default.mixed (PennyLane-0.36.0)
- default.qubit (PennyLane-0.36.0)
- default.qubit.autograd (PennyLane-0.36.0)
- default.qubit.jax (PennyLane-0.36.0)
- default.qubit.legacy (PennyLane-0.36.0)
- default.qubit.tf (PennyLane-0.36.0)
- default.qubit.torch (PennyLane-0.36.0)
- default.qutrit (PennyLane-0.36.0)
- default.qutrit.mixed (PennyLane-0.36.0)
- null.qubit (PennyLane-0.36.0)
- lightning.qubit (PennyLane-Lightning-0.36.0)

Hi @DSlap0 , welcome to the Forum!

It does look like something is not implemented in lightning.qubit. We released a new version of PennyLane today. If you run
python -m pip install pennylane --upgrade
does it solve your problem?

It actually broke it, now I can’t seem to be able to load a lightning.qubit device at all! There are no error messages or anything of the sort, the code will just stop at the following line

dev = qml.device("lightning.qubit", wires=4)

Here is my qml.about() output so you know I have correctly updated the library:

Name: PennyLane
Version: 0.37.0
Summary: 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.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: ###
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Windows-10-10.0.19045-SP0
Python version:          3.9.19
Numpy version:           1.21.5
Scipy version:           1.10.1
Installed devices:
- default.clifford (PennyLane-0.37.0)
- default.gaussian (PennyLane-0.37.0)
- default.mixed (PennyLane-0.37.0)
- default.qubit (PennyLane-0.37.0)
- default.qubit.autograd (PennyLane-0.37.0)
- default.qubit.jax (PennyLane-0.37.0)
- default.qubit.legacy (PennyLane-0.37.0)
- default.qubit.tf (PennyLane-0.37.0)
- default.qubit.torch (PennyLane-0.37.0)
- default.qutrit (PennyLane-0.37.0)
- default.qutrit.mixed (PennyLane-0.37.0)
- default.tensor (PennyLane-0.37.0)
- null.qubit (PennyLane-0.37.0)
- lightning.qubit (PennyLane-Lightning-0.37.0)

Hi @DSlap0,

This confirms that there’s some package mismatch in the environment you’re using.

I recommended that you create a new virtual environment to avoid any installation issues. 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.

You don’t need to install lightning.qubit in addition to this, it already comes together with PennyLane.

Let me know if this solves your issue!

Hi,

So I did exactly what you told me, and yet I still can’t make a lightning.qubit anymore. Here is my qml.about() output now:

Name: PennyLane
Version: 0.37.0
Summary: 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.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: ###
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane_Lightning

Platform info:           Windows-10-10.0.19045-SP0
Python version:          3.10.14
Numpy version:           1.26.4
Scipy version:           1.14.0
Installed devices:
- default.clifford (PennyLane-0.37.0)
- default.gaussian (PennyLane-0.37.0)
- default.mixed (PennyLane-0.37.0)
- default.qubit (PennyLane-0.37.0)
- default.qubit.autograd (PennyLane-0.37.0)
- default.qubit.jax (PennyLane-0.37.0)
- default.qubit.legacy (PennyLane-0.37.0)
- default.qubit.tf (PennyLane-0.37.0)
- default.qubit.torch (PennyLane-0.37.0)
- default.qutrit (PennyLane-0.37.0)
- default.qutrit.mixed (PennyLane-0.37.0)
- default.tensor (PennyLane-0.37.0)
- null.qubit (PennyLane-0.37.0)
- lightning.qubit (PennyLane_Lightning-0.37.0)

I’m sorry to hear that @DSlap0 .

This seems to be one of those rare cases where Conda doesn’t work. In those cases the recommendation is to use venv instead.

You can create a virtual environment with venv and install PennyLane as follows (Mac & Windows):

  1. Download Python 3.10 or newer here.
  2. Double-click on the .pkg file (Mac) or .exe file (Windows).
  3. Double-click on the Install Certificates.command file. It will open a terminal/command
    line window. You can close it after it says “process completed”.
  4. Open your terminal/command line.
  5. Create the directory for your virtual environment with: mkdir ~/.virtualenvs
  6. Create a new virtual environment with:
    python3.X -m venv ~/.virtualenvs/new_venv
    where X is the version of Python you installed. E.g. if you installed Python 3.10 then start with python3.10
    You can change new_venv to whatever name you want for your environment.
  7. Activate the environment with:
    source ~/.virtualenvs/new_venv/bin/activate
  8. Install PennyLane with: python -m pip install pennylane
  9. Install other useful packages such as Jupyter and Matplotlib with:
    python -m pip install jupyter matplotlib
  10. Write jupyter notebook in your terminal/command line and you will be ready to
    create programs using PennyLane.
    Every time you want to use PennyLane make sure to open a new terminal/command line, activate the environment with step 7, and open a Jupyter notebook with step 10.

Let me know if this solves your issue!