[BUG]Computing Jacobian for multiple measurement outcomes

[pratibharhegde]

Expected behavior

I want to compute the Jacobian of a QNode which has a tuple of measurement results as output. I am referring to the example shown in qml.jacobian — PennyLane 0.38.0 documentation. The code is expected to return a tuple with Jacobian corresponding to each measurement outcome.

Actual behavior

The code returns
TypeError: 'ArrayVSpace' object cannot be interpreted as an integer

Additional information

No response

Source code

import pennylane as qml
from pennylane import numpy as np


dev = qml.device("default.qubit", wires=2)

print(qml.about())
@qml.qnode(dev)
def circuit(x, y, z):
    qml.RX(x, wires=0)
    qml.RY(y, wires=1)
    qml.RZ(z, wires=0)
    return tuple(qml.expval(qml.Z(w)) for w in dev.wires)

x = np.array(0.2, requires_grad=True)
y = np.array(0.9, requires_grad=True)
z = np.array(-1.4, requires_grad=True)

jac = qml.jacobian(circuit)(x, y, z)

Tracebacks

UserWarning: Output seems independent of input.
  warnings.warn("Output seems independent of input.")
Traceback (most recent call last):
  File "C:\Users\hedge\PycharmProjects\venv\Lib\site-packages\numpy\core\fromnumeric.py", line 59, in _wrapfunc
    return bound(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^^
TypeError: 'ArrayVSpace' object cannot be interpreted as an integer

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "C:\Users\hedge\AppData\Roaming\JetBrains\PyCharmEdu2022.1\scratches\scratch.py", line 19, in <module>
    jac = qml.jacobian(circuit)(x, y, z)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hedge\PycharmProjects\venv\Lib\site-packages\pennylane\_grad.py", line 456, in _jacobian_function
    jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hedge\PycharmProjects\venv\Lib\site-packages\pennylane\_grad.py", line 456, in <genexpr>
    jac = tuple(_jacobian(func, arg)(*args, **kwargs) for arg in _argnum)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hedge\PycharmProjects\venv\Lib\site-packages\autograd\wrap_util.py", line 20, in nary_f
    return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hedge\PycharmProjects\venv\Lib\site-packages\autograd\differential_operators.py", line 64, in jacobian
    return np.reshape(np.stack(grads), jacobian_shape)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hedge\PycharmProjects\venv\Lib\site-packages\autograd\tracer.py", line 48, in f_wrapped
    return f_raw(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hedge\PycharmProjects\venv\Lib\site-packages\numpy\core\fromnumeric.py", line 285, in reshape
    return _wrapfunc(a, 'reshape', newshape, order=order)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hedge\PycharmProjects\venv\Lib\site-packages\numpy\core\fromnumeric.py", line 68, in _wrapfunc
    return _wrapit(obj, method, *args, **kwds)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "C:\Users\hedge\PycharmProjects\venv\Lib\site-packages\numpy\core\fromnumeric.py", line 45, in _wrapit
    result = getattr(asarray(obj), method)(*args, **kwds)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
TypeError: 'ArrayVSpace' object cannot be interpreted as an integer

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: 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: C:\Users\hedge\PycharmProjects\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.22631-SP0
Python version: 3.11.9
Numpy version: 1.26.4
Scipy version: 1.14.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)
None

Existing GitHub issues

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

[albi3ro]

Thanks for bringing this to our attention @pratibharhegde . This is a case of our documentation being very out of date.

One option would be to switch to using a different ML framework, like jax, that supports functions with complicated inputs and outputs.

To continue using autograd, you will need to stack the results into a numpy array. This can be done in the qnode itself:

dev = qml.device("default.qubit", wires=3)

@qml.qnode(dev)
def circuit(x, y):
    qml.RX(x[0], wires=0)
    qml.RY(y[0, 3], wires=1)
    qml.RX(x[1], wires=2)
    return qml.numpy.array([qml.expval(qml.Z(w)) for w in [0, 1, 2]])

Or after the qnode:

qml.jacobian(lambda *args, **kwargs: qml.math.hstack(circuit(*args, **kwargs)) )(x, y)