Repository: PennyLaneAI/pennylane

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⟩ 
```


Discussion:
Thanks @cvjjm for catching this. The first is definitely a bug --- when we added support for qubit-wise commuting observables within a single QNode, we did not update the drawing code, which was an oversight.

Re: the second error, I am surprised we never considered this edge case! This is indeed something we should add.

Resolving patch:
diff --git a/pennylane/circuit_graph.py b/pennylane/circuit_graph.py
--- a/pennylane/circuit_graph.py
+++ b/pennylane/circuit_graph.py
@@ -21,6 +21,7 @@
 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 @@ def __init__(self, ops, obs, wires, par_info=None, trainable_params=None):
         # 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 @@ def greedy_layers(self, wire_order=None, show_all_wires=False):
             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 @@ def has_path(self, a, b):
             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
--- a/tests/circuit_drawer/test_circuit_drawer.py
+++ b/tests/circuit_drawer/test_circuit_drawer.py
@@ -736,8 +736,6 @@ def test_qubit_circuit_with_state(
         """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 @@ def qfunc(a, w):
             + " 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
--- a/tests/circuit_graph/test_circuit_graph.py
+++ b/tests/circuit_graph/test_circuit_graph.py
@@ -74,6 +74,25 @@ def qfunc(a, b, c, d, e, f):
     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 @@ def test_iterate_layers(self, parameterized_circuit, wires):
         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