NotImplementedError: VJP of eigvals wrt argnums (0,) not defined[BUG]

[dslap0]

Expected behavior

With default.qubit, the code works without any issues, so you can just change lightning.qubit by default.qubit to get the expected behaviour.

Actual behavior

The code crashes while raising a NotImplementedError when the lightning.qubit device is used.

Additional information

No response

Source code

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)))

Tracebacks

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

System information

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: c:\Users\Nicolas Levasseur\.virtualenvs\new_venv\Lib\site-packages
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.12.4
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)

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.

[CatalinaAlbornoz]

Thanks again for the making this bug report @dslap0 !

Note for the person working on this bug:
The output of the circuit is a tensor when using default.qubit but it’s a float when using lightning.qubit. Could this be related to the error?

This bug report originated from this Forum thread.

[albi3ro]

Note that this issue is strictly a backprop versus non-backprop issue, and can be replicated with default.qubit by setting either diff_method="adjoint" or diff_method="parameter-shift". The issue is that two_qubit_decomposition, our way of decomposing a QubitUnitary into rotation gates, is not fully differentiable.

For example, this also shows the issue:

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 f(W):
    # circX() is a function that returns a unitary matrix as a 4x4 numpy array
    op = qml.QubitUnitary(circX(W[0], W[1], W[2]), wires=[0, 1])
    return qml.matrix(op.decomposition, wire_order=(0,1))()

weights = np.random.randn(3, requires_grad=True)

qml.grad(f)(weights)