Error for gradient with tensor observables QiskitDevice

[antalszava]

Describe the bug An error is raised when the gradient of a QNode with tensor observables is computed when using the Qiskit device. (Thanks @trbromley for catching this!)

To Reproduce

import pennylane as qml
import numpy as np

n_qubits = 2
depth = 2
noise = 0.1

def ansatz(weights):
    weights = weights.reshape(depth, n_qubits, 3)
    qml.templates.StronglyEntanglingLayers(weights, wires=list(range(n_qubits)))
#     return qml.expval(qml.PauliZ(0))  # The error does NOT occur when using a single qubit observable
    return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

from qiskit.providers.aer.noise import NoiseModel
from qiskit.providers.aer.noise.errors.standard_errors import (
    depolarizing_error,
    amplitude_damping_error,
)

noise_model = NoiseModel()
noise_model.add_all_qubit_quantum_error(depolarizing_error(noise, 1), ["u1", "u2", "u3"])

dev_qsk = qml.device(
                "qiskit.aer",
                wires=n_qubits,
                shots=10000,
                noise_model=noise_model,
                backend="qasm_simulator",
            )

weights = np.random.random((depth, n_qubits, 3)).flatten()

exp_sampled = qml.QNode(ansatz, dev_qsk, diff_method="parameter-shift") # Want to get expectation value and gradient
grad_shift = qml.grad(exp_sampled, argnum=0)
print(exp_sampled(weights))
print(grad_shift(weights))

Expected behavior The gradient is computed without any errors and the correct result is returned.

Screenshots

~/pennylane_qiskit/qiskit_device.py in expval(self, observable, wires, par)
    353 
    354         # estimate the ev
--> 355         return np.mean(self.sample(observable, wires, par))
    356 
    357     def var(self, observable, wires, par):

~/pennylane_qiskit/qiskit_device.py in sample(self, observable, wires, par)
    404 
    405         for w, b in zip(eigvals, itertools.product([0, 1], repeat=len(wires))):
--> 406             samples = np.where(np.all(res == b, axis=1), w, samples)
    407 
    408         return samples

<__array_function__ internals> in where(*args, **kwargs)

ValueError: operands could not be broadcast together with shapes (9999,) () (10000,) 

qml.about() output

Name: PennyLane
Version: 0.9.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: None
Author-email: None
License: Apache License 2.0
Location: /pennylane
Requires: numpy, scipy, networkx, autograd, toml, appdirs, semantic-version
Required-by: 
Platform info:           Linux-4.19.11-041911-generic-x86_64-with-debian-buster-sid
Python version:          3.7.6
Numpy version:           1.18.1
Scipy version:           1.3.2
Installed devices:
- qiskit.aer (PennyLane-qiskit-0.9.0.dev0)
- qiskit.basicaer (PennyLane-qiskit-0.9.0.dev0)
- qiskit.ibmq (PennyLane-qiskit-0.9.0.dev0)
- strawberryfields.fock (PennyLane-SF-0.8.0)
- strawberryfields.gaussian (PennyLane-SF-0.8.0)
- forest.numpy_wavefunction (PennyLane-Forest-0.8.0)
- forest.qpu (PennyLane-Forest-0.8.0)
- forest.qvm (PennyLane-Forest-0.8.0)
- forest.wavefunction (PennyLane-Forest-0.8.0)
- projectq.classical (PennyLane-PQ-0.7.0.dev0)
- projectq.ibm (PennyLane-PQ-0.7.0.dev0)
- projectq.simulator (PennyLane-PQ-0.7.0.dev0)
- microsoft.QuantumSimulator (PennyLane-qsharp-0.6.0)
- cirq.simulator (PennyLane-Cirq-0.8.0)
- default.gaussian (PennyLane-0.9.0.dev0)
- default.qubit (PennyLane-0.9.0.dev0)
- default.tensor (PennyLane-0.9.0.dev0)
- default.tensor.tf (PennyLane-0.9.0.dev0)

Additional context This bug is related to how the tensor observable part of QiskitDevice.sample is used: https://github.com/XanaduAI/pennylane-qiskit/blob/269802a7c58dd951c2a183e071e31ed93e3c906d/pennylane_qiskit/qiskit_device.py#L400

When changing this to the logic used in QubitDevice, the error seems to be resolved.

Also, on first look, it gave the impression that the error is raised upon consecutive calls to sample.

Possible ways to resolve this bug:

1. Modify the tensor observable part of QiskitDevice.sample.
2. Port QiskitDevice to QubitDevice (bigger undertaking).

In either case, an integration test needs to be added testing gradient computation for the tensor observable case.

[josh146]

I'm in favour of porting to QubitDevice, as this will reduce maintenance load in the long run.

[antalszava]

Closing as solved by #83