Repository: PennyLaneAI/pennylane (python)

Issue:
Drawing fails for (a) qubit-wise commuting observables, and (b) QNodes with no operations
Consider the following minimal example in which a QNode returns multiple qubit wise commuting obeservables:
```python
import networkx
import pennylane as qml

from pennylane import expval, PauliX, PauliY, PauliZ

dev = qml.device('default.qubit', wires=4)
@qml.qnode(dev)
def qnode(x, y):
    qml.RY(x, wires=0)
    qml.Hadamard(0)
    qml.RZ(y, wires=0)
    return [expval(PauliX(wires=[0]) @ PauliX(wires=[1]) @ PauliX(wires=[2])), expval(PauliX(wires=[0]) @ PauliZ(wires=[3]))]
    #return [expval(PauliX(wires=[0]) @ PauliX(wires=[1]) @ PauliX(wires=[2]))]
    #return [expval(PauliX(wires=[0]) @ PauliZ(wires=[3]))]

qnode(1., 2.)
print(qnode.draw())
```

Problem 1: While executing the circuit works, drawing it leads to the error 
```python
.../pennylane/circuit_drawer/grid.py in __init__(self, raw_grid)
     47             self.raw_grid = np.array(raw_grid, dtype=object)
     48             if len(self.raw_grid.shape) != 2:
---> 49                 raise ValueError(
     50                     "The entered raw grid was not parsed as two-dimensional array: {}".format(
     51                         raw_grid

ValueError: The entered raw grid was not parsed as two-dimensional array: [[expval(PauliZ(wires=[0]) @ PauliZ(wires=[1]) @ PauliZ(wires=[2])), expval(PauliZ(wires=[0]) @ PauliZ(wires=[3]))], [expval(PauliZ(wires=[0]) @ PauliZ(wires=[1]) @ PauliZ(wires=[2]))], [expval(PauliZ(wires=[0]) @ PauliZ(wires=[1]) @ PauliZ(wires=[2]))], [expval(PauliZ(wires=[0]) @ PauliZ(wires=[3]))]]
```
Commenting out the active return line line of the QNode and instead uncommenting one of the other two shows that PennyLane can draw with both observables individually, but not with both.

Problem 2: Trying to draw a circuit with no Operations (and only a return statement) results in the error:
```python
.../pennylane/pennylane/circuit_drawer/grid.py in copy(self)
    177             Grid: A copy of the Grid
    178         """
--> 179         return Grid(self.raw_grid.copy())
    180 
    181     def append_grid_by_layers(self, other_grid):

AttributeError: 'NoneType' object has no attribute 'copy'
```

It would be nice if instead
```python
import networkx
import pennylane as qml

from pennylane import expval, PauliX, PauliY, PauliZ

dev = qml.device('default.qubit', wires=4)
@qml.qnode(dev)
def qnode(x, y):
    return [expval(PauliX(wires=[0]) @ PauliX(wires=[1]) @ PauliX(wires=[2]))]

qnode(1., 2.)
print(qnode.draw())
```
would simply draw
```
 0: ────╭┤ ⟨X ⊗ X ⊗ X⟩ 
 1: ────├┤ ⟨X ⊗ X ⊗ X⟩ 
 2: ────╰┤ ⟨X ⊗ X ⊗ X⟩ 
```

Pull request description:
Fix CircuitDrawer single wire multiple observables
**Context**
QNodes can have multiple measurements on the same wire so long as they are part of qubit-wise commuting observables. Drawing such QNodes results in an error, unfortunately.

This boils down to how the `CircuitGraph` processes observables. It assumes that there is a single observable for every QNode or alternatively it merges observables for separate return measurement processes. When there are more than one observable per wire, it causes an error

**Changes**

Changes the logic in the `greedy_layers` method of the `CircuitGraph` class to account for multiple cases where there was a wire that was measured more than once.

*Note*: in such cases the circuit includes diagonalizations and the observables will be drawn as PauliZ operators.

**Ralated Issues**
Closes #1326 

Patch:
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
index c0a68ec87..a2d4b3f1f 100644
--- a/.github/CHANGELOG.md
+++ b/.github/CHANGELOG.md
@@ -322,6 +322,9 @@ random_mat2 = rng.standard_normal(3, requires_grad=False)
 
 <h3>Bug fixes</h3>
 
+* Fixes drawing QNodes that contain multiple measurements on a single wire.
+  [(#1353)](https://github.com/PennyLaneAI/pennylane/pull/1353)
+
 * Fixes drawing QNodes with no operations.
   [(#1354)](https://github.com/PennyLaneAI/pennylane/pull/1354)
 
diff --git a/pennylane/circuit_graph.py b/pennylane/circuit_graph.py
index 2bc47173a..f57da751f 100644
--- a/pennylane/circuit_graph.py
+++ b/pennylane/circuit_graph.py
@@ -21,6 +21,7 @@ from collections import Counter, OrderedDict, namedtuple
 import networkx as nx
 
 import pennylane as qml
+import numpy as np
 
 from .circuit_drawer import CHARSETS, CircuitDrawer
 
@@ -67,6 +68,18 @@ def _list_at_index_or_none(ls, idx):
     return None
 
 
+def is_returned_observable(op):
+    """Helper for the condition of having an observable or
+    measurement process in the return statement.
+
+    Returns:
+        bool: whether or not the observable or measurement process is in the
+        return statement
+    """
+    is_obs = isinstance(op, (qml.operation.Observable, qml.measure.MeasurementProcess))
+    return is_obs and op.return_type is not None
+
+
 Layer = namedtuple("Layer", ["ops", "param_inds"])
 """Parametrized layer of the circuit.
 
@@ -168,6 +181,10 @@ class CircuitGraph:
         # including the observables
         self._operation_graph = None
 
+        # Required to keep track if we need to handle multiple returned
+        # observables per wire
+        self._max_simultaneous_measurements = None
+
     def print_contents(self):
         """Prints the contents of the quantum circuit."""
 
@@ -480,19 +497,35 @@ class CircuitGraph:
             l += 1
 
         observables = OrderedDict()
-        for wire in sorted(self._grid):
-            observables[wire] = list(
-                filter(
-                    lambda op: isinstance(
-                        op, (qml.operation.Observable, qml.measure.MeasurementProcess)
+
+        if self.max_simultaneous_measurements == 1:
+
+            # There is a single measurement
+            for wire in sorted(self._grid):
+                observables[wire] = list(
+                    filter(
+                        lambda op: isinstance(
+                            op, (qml.operation.Observable, qml.measure.MeasurementProcess)
+                        )
+                        and op.return_type is not None,
+                        self._grid[wire],
                     )
-                    and op.return_type is not None,
-                    self._grid[wire],
                 )
-            )
+                if not observables[wire]:
+                    observables[wire] = [None]
+        else:
+
+            # There are multiple measurements
+            for wire in sorted(self._grid):
+                mp_map = dict(zip(self.observables, range(self.max_simultaneous_measurements)))
 
-            if not observables[wire]:
-                observables[wire] = [None]
+                # Initialize to None everywhere
+                observables[wire] = [None] * self.max_simultaneous_measurements
+
+                for op in self._grid[wire]:
+                    if is_returned_observable(op):
+                        obs_idx = mp_map[op]
+                        observables[wire][obs_idx] = op
 
         if wire_order is not None:
             temp_op_grid = OrderedDict()
@@ -608,3 +641,43 @@ class CircuitGraph:
             bool: returns ``True`` if a path exists
         """
         return nx.has_path(self._graph, a, b)
+
+    @property
+    def max_simultaneous_measurements(self):
+        """Returns the maximum number of measurements on any wire in the circuit graph.
+
+        This method counts the number of measurements for each wire and returns
+        the maximum.
+
+        **Examples**
+
+
+        >>> dev = qml.device('default.qubit', wires=3)
+        >>> def circuit_measure_max_once():
+        ...     return qml.expval(qml.PauliX(wires=0))
+        >>> qnode = qml.QNode(circuit_measure_max_once, dev)
+        >>> qnode()
+        >>> qnode.qtape.graph.max_simultaneous_measurements
+        1
+        >>> def circuit_measure_max_twice():
+        ...     return qml.expval(qml.PauliX(wires=0)), qml.probs(wires=0)
+        >>> qnode = qml.QNode(circuit_measure_max_twice, dev)
+        >>> qnode()
+        >>> qnode.qtape.graph.max_simultaneous_measurements
+        2
+
+        Returns:
+            int: the maximum number of measurements
+        """
+        if self._max_simultaneous_measurements is None:
+            all_wires = []
+
+            for obs in self.observables:
+                all_wires.extend(obs.wires.tolist())
+
+            a = np.array(all_wires)
+            _, counts = np.unique(a, return_counts=True)
+            self._max_simultaneous_measurements = (
+                counts.max() if counts.size != 0 else 1
+            )  # qml.state() will result in an empty array
+        return self._max_simultaneous_measurements


Test patch:
diff --git a/tests/circuit_drawer/test_circuit_drawer.py b/tests/circuit_drawer/test_circuit_drawer.py
index 473d9c301..f96e0afd5 100644
--- a/tests/circuit_drawer/test_circuit_drawer.py
+++ b/tests/circuit_drawer/test_circuit_drawer.py
@@ -736,8 +736,6 @@ class TestCircuitDrawerIntegration:
         """Test that a qubit circuit with unused wires renders correctly."""
         output = qubit_circuit_with_state.draw()
 
-        print(output)
-        print(drawn_qubit_circuit_with_state)
         assert output == drawn_qubit_circuit_with_state
 
     def test_direct_qnode_integration(self):
@@ -765,6 +763,30 @@ class TestCircuitDrawerIntegration:
             + " 1: -----+RX(2.3)--Rot(1.2, 3.2, 0.7)--+RX(-2.3)--+| <Z @ Z> \n"
         )
 
+    def test_same_wire_multiple_measurements(self):
+        """Test that drawing a QNode with multiple measurements on certain wires works correctly."""
+        dev = qml.device("default.qubit", wires=4)
+
+        @qml.qnode(dev)
+        def qnode(x, y):
+            qml.RY(x, wires=0)
+            qml.Hadamard(0)
+            qml.RZ(y, wires=0)
+            return [
+                qml.expval(qml.PauliX(wires=[0]) @ qml.PauliX(wires=[1]) @ qml.PauliX(wires=[2])),
+                qml.expval(qml.PauliX(wires=[0]) @ qml.PauliX(wires=[3])),
+            ]
+
+        qnode(1.0, 2.0)
+
+        expected = (
+            " 0: ──RY(1)──────H──RZ(2)──RY(-1.57)──╭┤ ⟨Z ⊗ Z ⊗ Z⟩ ╭┤ ⟨Z ⊗ Z⟩ \n"
+            + " 1: ──RY(-1.57)───────────────────────├┤ ⟨Z ⊗ Z ⊗ Z⟩ │┤         \n"
+            + " 2: ──RY(-1.57)───────────────────────╰┤ ⟨Z ⊗ Z ⊗ Z⟩ │┤         \n"
+            + " 3: ──RY(-1.57)────────────────────────┤             ╰┤ ⟨Z ⊗ Z⟩ \n"
+        )
+        assert qnode.draw() == expected
+
 
 class TestWireOrdering:
     """Tests for wire ordering functionality"""
diff --git a/tests/circuit_graph/test_circuit_graph.py b/tests/circuit_graph/test_circuit_graph.py
index 736482f33..ea811689b 100644
--- a/tests/circuit_graph/test_circuit_graph.py
+++ b/tests/circuit_graph/test_circuit_graph.py
@@ -74,6 +74,25 @@ def parameterized_circuit(wires):
     return qfunc
 
 
+def circuit_measure_max_once():
+    """A fixture of a circuit that measures wire 0 once."""
+    return qml.expval(qml.PauliX(wires=0))
+
+
+def circuit_measure_max_twice():
+    """A fixture of a circuit that measures wire 0 twice."""
+    return qml.expval(qml.PauliX(wires=0)), qml.probs(wires=0)
+
+
+def circuit_measure_multiple_with_max_twice():
+    """A fixture of a circuit that measures wire 0 twice."""
+    return (
+        qml.expval(qml.PauliX(wires=0)),
+        qml.probs(wires=[0, 1, 2]),
+        qml.var(qml.PauliX(wires=[1]) @ qml.PauliZ([2])),
+    )
+
+
 class TestCircuitGraph:
     """Test conversion of queues to DAGs"""
 
@@ -209,3 +228,21 @@ class TestCircuitGraph:
         assert set(result[2][1]) == set(circuit.operations[5:])
         assert result[2][2] == (6, 7)
         assert set(result[2][3]) == set(circuit.observables[1:])
+
+    @pytest.mark.parametrize(
+        "circ, expected",
+        [
+            (circuit_measure_max_once, 1),
+            (circuit_measure_max_twice, 2),
+            (circuit_measure_multiple_with_max_twice, 2),
+        ],
+    )
+    def test_max_simultaneous_measurements(self, circ, expected):
+        """A test for getting the maximum number of measurements on any wire in
+        the circuit graph."""
+
+        dev = qml.device("default.qubit", wires=3)
+        qnode = qml.QNode(circ, dev)
+        qnode()
+        circuit = qnode.qtape.graph
+        assert circuit.max_simultaneous_measurements == expected