Repository: PennyLaneAI/pennylane

Issue:
[BUG] KeyError while drawing a QNode
### Expected behavior

Drawing the circuit should work normally

### Actual behavior

A KeyError is raised.

### Additional information

Please see the sample program to reproduce the problem.

### Source code

```python
import pennylane as qml
import numpy as np

dev = qml.device('default.qubit', wires=2)
params = np.random.rand(2, 3)

@qml.qnode(dev)
def qnode(params, foo=None):
    qml.Rot(*params[0], wires=[0])
    qml.Rot(*params[1], wires=[1])
    qml.CNOT(wires=[0, 1])
    return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)), qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))]

qnode.construct((params,), {'foo':2})
print(qnode.draw())
```


### Tracebacks

```shell
Traceback (most recent call last):
  File "draw_test.py", line 15, in <module>
    print(qnode.draw())
  File "/home/cvjjm/src/bar/foo/pennylane/pennylane/qnode.py", line 841, in draw
    return self.qtape.draw(
  File "/home/cvjjm/src/bar/foo/pennylane/pennylane/tape/tape.py", line 1119, in draw
    return self.graph.draw(
  File "/home/cvjjm/src/bar/foo/pennylane/pennylane/circuit_graph.py", line 612, in draw
    grid, obs = self.greedy_layers(wire_order=wire_order, show_all_wires=show_all_wires)
  File "/home/cvjjm/src/bar/foo/pennylane/pennylane/circuit_graph.py", line 534, in greedy_layers
    obs_idx = mp_map[op]
KeyError: expval(PauliZ(wires=[0]) @ PauliZ(wires=[1]))
```


### System information

```shell
WARNING: pip is being invoked by an old script wrapper. This will fail in a future version of pip.
Please see https://github.com/pypa/pip/issues/5599 for advice on fixing the underlying issue.
To avoid this problem you can invoke Python with '-m pip' instead of running pip directly.
Name: PennyLane
Version: 0.19.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: /home/cvjjm/src/covqcstack/qcware/pennylane
Requires: numpy, scipy, networkx, autograd, toml, appdirs, semantic-version, autoray, cachetools, pennylane-lightning
Required-by: covvqetools, pytket-pennylane, PennyLane-Qchem, PennyLane-Lightning
Platform info:           Linux-5.10.16.3-microsoft-standard-WSL2-x86_64-with-glibc2.10
Python version:          3.8.8
Numpy version:           1.20.1
Scipy version:           1.6.1
Installed devices:
- cov.caching.qubit (covvqetools-0.1.0)
- cov.caching.qulacs (covvqetools-0.1.0)
- cov.parsec.casbox (covvqetools-0.1.0)
- cov.parsec.qubit (covvqetools-0.1.0)
- pytket.pytketdevice (pytket-pennylane-0.1.0)
- lightning.qubit (PennyLane-Lightning-0.18.0)
- default.gaussian (PennyLane-0.19.0.dev0)
- default.mixed (PennyLane-0.19.0.dev0)
- default.qubit (PennyLane-0.19.0.dev0)
- default.qubit.autograd (PennyLane-0.19.0.dev0)
- default.qubit.jax (PennyLane-0.19.0.dev0)
- default.qubit.tf (PennyLane-0.19.0.dev0)
- default.qubit.torch (PennyLane-0.19.0.dev0)
- default.tensor (PennyLane-0.19.0.dev0)
- default.tensor.tf (PennyLane-0.19.0.dev0)

Compilation finished at Wed Sep 29 10:51:53
```


### 

- [X] I have searched exisisting GitHub issues to make sure the issue does not already exist.


Discussion:
Thanks @cvjjm! This is indeed strange, seems to be something about the combination of observables. I've reduced it down to an even more minimal non-working case:

```python
dev = qml.device('default.qubit', wires=2)

@qml.qnode(dev)
def qnode():
    return [
        qml.expval(qml.PauliZ(0)),
        qml.expval(qml.PauliZ(1)),
        qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
    ]

print(qml.draw(qnode)())  # errors
```

Strangely enough, commenting out at least one of the observables in the return list causes it to work:

```python
@qml.qnode(dev)
def qnode():
    return [
        qml.expval(qml.PauliZ(0)),
        # qml.expval(qml.PauliZ(1)),
        qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
    ]

>>> print(qml.draw(qnode)())
 0: ──┤ ⟨Z⟩  ╭┤ ⟨Z ⊗ Z⟩
 1: ──┤     ╰┤ ⟨Z ⊗ Z⟩
```

```python
@qml.qnode(dev)
def qnode():
    return [
        # qml.expval(qml.PauliZ(0)),
        qml.expval(qml.PauliZ(1)),
        qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
    ]

>>> print(qml.draw(qnode)())
 0: ──┤    ╭┤ ⟨Z ⊗ Z⟩
 1: ──┤ ⟨Z⟩ ╰┤ ⟨Z ⊗ Z⟩
```


Yes, indeed. The error seems to always be raised on the last of the observables when the order is changed. Must have something to do with how `mp_map` is built.

Resolving patch:
diff --git a/pennylane/circuit_graph.py b/pennylane/circuit_graph.py
--- a/pennylane/circuit_graph.py
+++ b/pennylane/circuit_graph.py
@@ -507,7 +507,7 @@ def greedy_layers(self, wire_order=None, show_all_wires=False):
 
         if self.max_simultaneous_measurements == 1:
 
-            # There is a single measurement
+            # There is a single measurement for every wire
             for wire in sorted(self._grid):
                 observables[wire] = list(
                     filter(
@@ -522,12 +522,16 @@ def greedy_layers(self, wire_order=None, show_all_wires=False):
                     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)))
+            # There are wire(s) with multiple measurements.
+            # We are creating a separate "visual block" at the end of the
+            # circuit for each observable and mapping observables with block
+            # indices.
+            num_observables = len(self.observables)
+            mp_map = dict(zip(self.observables, range(num_observables)))
 
+            for wire in sorted(self._grid):
                 # Initialize to None everywhere
-                observables[wire] = [None] * self.max_simultaneous_measurements
+                observables[wire] = [None] * num_observables
 
                 for op in self._grid[wire]:
                     if _is_returned_observable(op):


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
@@ -806,6 +806,32 @@ def qnode(x, y):
         )
         assert qnode.draw() == expected
 
+    def test_same_wire_multiple_measurements_many_obs(self):
+        """Test that drawing a QNode with multiple measurements on certain
+        wires works correctly when there are more observables than the number of
+        observables for any wire.
+        """
+        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.PauliZ(0)),
+                qml.expval(qml.PauliZ(1)),
+                qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)),
+            ]
+
+        qnode(0.3, 0.2)
+
+        expected = (
+            " 0: ──RY(0.3)──H──RZ(0.2)──┤ ⟨Z⟩ ┤     ╭┤ ⟨Z ⊗ Z⟩ \n"
+            + " 1: ───────────────────────┤     ┤ ⟨Z⟩ ╰┤ ⟨Z ⊗ Z⟩ \n"
+        )
+        assert qnode.draw() == expected
+
 
 class TestWireOrdering:
     """Tests for wire ordering functionality"""