Feature request: simultaneous measure of observables using `qml.map`

[SamMugel]

---

Issue description

It is not possible to evaluate various expectation values simultaneously when using the qml.map function. At present, the number of times the circuit is simulated grows linearly with the number of observables.

Feature request

At present, the function qml.map expects an argument observables with type (Iterable[Observable]). I would like this to have type (Iterable[Observable or list]), i.e: one item of observables can be a list of objects Observable to be measured simultaneously.

Desired behaviour example

In this example, I would like to measure Pauli_z on wires 0 and 1, using qml.map , without having to rebuild the entire circuit.

def my_template(params, wires, **kwargs):
    qml.RX(params[0], wires=wires[0])
    qml.RX(params[1], wires=wires[1])
    qml.CNOT(wires=wires)
obs_list = [[qml.PauliZ(0), qml.PauliZ(1)], qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliX(1)]
dev = qml.device("default.qubit", wires=2)
qnodes = qml.map(my_template, obs_list, dev, measure="expval")
params = [0.54, 0.12]
qnodes(params)

At present, this code snippet raises

ValueError: Could not create QNodes. Some or all observables are not valid.

Previous discussion

This issue was first raised on slack:

https://xanadu-quantum.slack.com/archives/CRLS0EW07/p1584749829082500

[josh146]

Thanks @SamMugel! This is something we can very easily support. However, one slight issue is deciding what the QNode collection should return if this is allowed.

Currently, QNode collections assume that each observable in the observable list is a 'scalar', allowing the collection to return a 1D array:

import pennylane as qml

def my_template(params, wires, **kwargs):
    qml.RX(params[0], wires=wires[0])
    qml.RX(params[1], wires=wires[1])
    qml.CNOT(wires=wires)

obs_list = [
    qml.PauliX(0) @ qml.PauliZ(1),
    qml.PauliZ(0) @ qml.PauliX(1)
]

dev = qml.device("default.qubit", wires=2)
qnodes = qml.map(my_template, obs_list, dev, measure="expval")

Executing the collection,

>>> params = [0.54, 0.12]
>>> qnodes(params)
array([-0.06154835440747172, 0.0])

However, it's less clear what the QNode collection should return if the observable list has a nested structure. E.g, using your example above,

obs_list = [
    [qml.PauliZ(0), qml.PauliZ(1)],
    qml.PauliX(0) @ qml.PauliZ(1),
    qml.PauliZ(0) @ qml.PauliX(1)
]
qnodes = qml.map(my_template, obs_list, dev, measure="expval")

This could return the following:

1. A list of arrays:

   >>> qnode(params)
   [array([0.85770868, 0.85154059]), array([-0.06154835]), array([0.])]

2. A flattened array:

   >>> qnode(params)
   array([0.85770868, 0.85154059, -0.06154835, 0.])]

3. A jagged array:

   >>> qnode(params)
   array([array([0.8577, 0.8515]), array([-0.0615]), array([0.])],
     dtype=object)

4. A padded stacked array:

   >>> qnode(params)
   array([[ 0.8577,  0.8515],
      [-0.0615,  0.    ],
      [ 0.    ,  0.    ]])

In terms of matching the existing behaviour (where a 1D numpy array is returned), option (2) is the most attractive, but also confusing.

In terms of matching the observable list indexing, I think option (1) makes most sense, but results in the annoying case that the existing behaviour (a list of single observables) will result in nested output:

>>> obs_list = [qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliX(1)]
>>> qnodes = qml.map(my_template, obs_list, dev, measure="expval")
>>> params = [0.54, 0.12]
>>> qnodes(params)
[array([-0.06154835440747172]), array([0.0])]

[co9olguy]

What about returning a nested structure, with the same structure as the original list?

[SamMugel]

I understand, this is an issue indeed. Though it looks like option 3 solves this issue, as

>>> np.array([np.array([0.8577]), np.array([-0.0615]), np.array([0.])], dtype=object)
array([[0.8577],
       [-0.0615],
       [0.0]], dtype=object)

[josh146]

A quick update: we have been working on improving the PennyLane core (see https://pennylane.readthedocs.io/en/latest/code/qml_tape.html), and as a result it should solve some of the blockers we had above. In particular:

However, it's less clear what the QNode collection should return if the observable list has a nested structure

It turned out that QNodes returning ragged output was a much more common use-case than we originally expected, and would show up in a lot of places. For example, return qml.probs(wires=0), qml.probs(wires=[1, 2]).

In the new experimental tape mode, ragged QNodes return a flat array/tensor. This was chosen to maximize support and ensure identical behaviour for both black box and passthru QNodes.

[timmysilv]

I'm going to close this issue, because the relevant code has been replaced. QNodes can now return multiple measurements, and the parallelization is being handled in an all-new way (some bonus details pending with the upcoming device API! 🎆 )