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!

Hi Catalina,

Thank you for your help, I was now able to run the Pennylane code in the virtual environment, but the code gives me the same error message as before.

Hi @DSlap0 ,

Are you able to share the code for your circX() function? I’m not being able to replicate your error without it.

Hi Catalina,

I have this issue with 2 different functions that generate unitary matrices. Here is one of them:

def circX(a, b, g):
    # Matrix implementation in logical basis.
    # Obtained from Eqn.33 of paper.

    # Eigenvalues lambda_0,...,lambda_3. Note lambda_0=lambda_4 from paper.
    l0 = 2 * a + g
    l1 = -2 * b - g
    l2 = -2 * a + g
    l3 = 2 * b - g

    e0 = 0.25 * np.exp(-1.0j * l0)
    e1 = 0.25 * np.exp(-1.0j * l1)
    e2 = 0.25 * np.exp(-1.0j * l2)
    e3 = 0.25 * np.exp(-1.0j * l3)

    e11 = e0 + e1 + e2 + e3
    e12 = e0 - 1.0j * e1 - e2 + 1.0j * e3
    e13 = e0 - e1 + e2 - e3
    e14 = e0 + 1.0j * e1 - e2 - 1.0j * e3

    return np.array(
        [
            [e11, e12, e13, e14],
            [e14, e11, e12, e13],
            [e13, e14, e11, e12],
            [e12, e13, e14, e11],
        ]
    )

Hi @DSlap0 ,

Thanks for sharing this code.

There does seem to be some kind of bug here with lightning.qubit being unable to decompose QubitUnitary. Would you like to open the bug report yourself? If not I can open it.

In the meantime, you can try using qml.ArbitraryUnitary instead. The weights need to have length 4^n - 1, where n is the number of wires the ArbitraryUnitary acts upon. Here I’ve made a quick test where I use the 6 parameters you have in every layer and add 9 trailing zeros in order to complete the required length for the weights.

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])
    weights_test = np.reshape(W,(6,)) # Reshape the weights into a vector
    weights_test = np.pad(weights_test,(0,9),mode='constant',constant_values=0) # Add trailing zeros to make the length 15
    qml.ArbitraryUnitary(weights_test, wires=[0, 1])
    qml.ArbitraryUnitary(weights_test, wires=[2, 3])

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

With this change the code no longer shows an error for me.

Let me know if this works for you too.

Hi, and thank you again for all your help!

I would like to open the support ticket (if you could point me where I could do that). For the qml.ArbitraryUnitary, my only issue is that the code is made to test the 2 circX matrices, and so by using an arbitrary matrix I lose the main purpose of the code. However, it seems to be working for me too.

Hi @DSlap0 ,

No problem! You can open a bug report here. It’s not really a support ticket, it’s more a report so that the core team developing PennyLane is aware of this issue and can look into fixing it.

It’s important to share a minimal example so that they can focus on what’s actually causing the problem.

Below is a minimal working example for your problem.

import pennylane as qml
from pennylane import numpy as np

def circX(a, b, g):
    # Matrix implementation in logical basis.
    # Obtained from Eqn.33 of paper.

    # Eigenvalues lambda_0,...,lambda_3. Note lambda_0=lambda_4 from paper.
    l0 = 2 * a + g
    l1 = -2 * b - g
    l2 = -2 * a + g
    l3 = 2 * b - g

    e0 = 0.25 * np.exp(-1.0j * l0)
    e1 = 0.25 * np.exp(-1.0j * l1)
    e2 = 0.25 * np.exp(-1.0j * l2)
    e3 = 0.25 * np.exp(-1.0j * l3)

    e11 = e0 + e1 + e2 + e3
    e12 = e0 - 1.0j * e1 - e2 + 1.0j * e3
    e13 = e0 - e1 + e2 - e3
    e14 = e0 + 1.0j * e1 - e2 - 1.0j * e3

    return np.array(
        [
            [e11, e12, e13, e14],
            [e14, e11, e12, e13],
            [e13, e14, e11, e12],
            [e12, e13, e14, e11],
        ]
    )

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

# Define quantum circuit
dev = qml.device("lightning.qubit", wires=4)
# When the device is changed to "default.qubit" it works

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

    return qml.expval(qml.PauliZ(0))

# Initialize weights
np.random.seed(0)
weights = np.random.randn(3, requires_grad=True)

print(type(circuit(weights)))

# Optimization
opt = qml.NesterovMomentumOptimizer(0.2, 0.1)

for it in range(40):
  weights = opt.step(circuit, weights)
  if (it + 1) % 10 == 0:
    print("Iter: {:5d} | Energy: {:0.7f}".format(it + 1, circuit(weights)))

By making this minimal example I noticed that the output of the circuit is a tensor when using default.qubit but it’s a float when using lightning.qubit. This may be related to the error so it’s a good clue to mention in the bug report.

Once you make the bug report make sure to add the link here so that anyone having a similar problem can find the bug report.

Let me know if you have any questions about this!

Thank you again for your help! Here is the link to the bug report for posteriority: NotImplementedError: VJP of eigvals wrt argnums (0,) not defined[BUG] · Issue #6045 · PennyLaneAI/pennylane · GitHub

Awesome @DSlap0 !

Thank you for making the bug report!