Repository: PennyLaneAI/pennylane

Issue:
Multi-threaded QNodeCollection does not work
#### Issue description
Thread-parallel execution of `QNodeCollection` fails by the following error in tape mode.

```
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/collections/qnode_collection.py", line 271, in __call__
    results = self.evaluate(args, kwargs)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/collections/qnode_collection.py", line 229, in evaluate
    return dask.compute(*results, scheduler=_scheduler)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/base.py", line 452, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/threaded.py", line 84, in get
    **kwargs
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/local.py", line 486, in get_async
    raise_exception(exc, tb)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/local.py", line 316, in reraise
    raise exc
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/local.py", line 222, in execute_task
    result = _execute_task(task, data)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/core.py", line 121, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/utils.py", line 29, in apply
    return func(*args, **kwargs)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/tape/qnode.py", line 433, in __call__
    self.construct(args, kwargs)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/tape/qnode.py", line 395, in construct
    "All measurements must be returned in the order they are measured."
pennylane.qnodes.base.QuantumFunctionError: All measurements must be returned in the order they are measured.
```
Reproduction code is:

```
import pennylane as qml


import numpy as np

n_qubits = 4
n_batches = 5
dev = qml.device("default.qubit", wires=n_qubits)

qml.enable_tape()

@qml.qnode(dev)
def circuit(inputs, weights):
    qml.templates.AngleEmbedding(inputs, wires=range(n_qubits))
    qml.templates.BasicEntanglerLayers(weights, wires=range(n_qubits))
    return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]


n_layers = 1
weight_shapes = {"weights": (n_layers, n_qubits)}

qnode = QNodeCollection([circuit for _ in range(n_batches)])
x = np.random.rand((n_qubits)).astype(np.float64)#.to(torch.device('cuda'))
p = np.random.rand(*weight_shapes['weights']).astype(np.float64)#.to(torch.device('cuda'))
res = qnode(x, p, parallel=True)

```
The reason for this seems to be the empty wires in the `qtape` in `QNode`. I could not yet find the actual reason for this .

* *Expected behavior:* (What you expect to happen)
Thread-parallel execution is possible in tape mode.

* *Actual behavior:* (What actually happens)
Thread-parallel execution fails in tape mode.

* *Reproduces how often:* (What percentage of the time does it reproduce?)
~50%

* *System information:* (post the output of `import pennylane as qml; qml.about()`)
```
Name: PennyLane
Version: 0.13.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/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg
Requires: numpy, scipy, networkx, autograd, toml, appdirs, semantic-version
Required-by: PennyLane-SF, pennylane-qulacs, PennyLane-qiskit
Platform info:           Linux-5.4.0-53-generic-x86_64-with-debian-bullseye-sid
Python version:          3.7.9
Numpy version:           1.19.2
Scipy version:           1.5.2
Installed devices:
- strawberryfields.fock (PennyLane-SF-0.12.0)
- strawberryfields.gaussian (PennyLane-SF-0.12.0)
- strawberryfields.gbs (PennyLane-SF-0.12.0)
- strawberryfields.remote (PennyLane-SF-0.12.0)
- strawberryfields.tf (PennyLane-SF-0.12.0)
- qulacs.simulator (pennylane-qulacs-0.12.0)
- qiskit.aer (PennyLane-qiskit-0.12.0)
- qiskit.basicaer (PennyLane-qiskit-0.12.0)
- qiskit.ibmq (PennyLane-qiskit-0.12.0)
- default.gaussian (PennyLane-0.13.0.dev0)
- default.mixed (PennyLane-0.13.0.dev0)
- default.qubit (PennyLane-0.13.0.dev0)
- default.qubit.autograd (PennyLane-0.13.0.dev0)
- default.qubit.tf (PennyLane-0.13.0.dev0)
- default.tensor (PennyLane-0.13.0.dev0)
- default.tensor.tf (PennyLane-0.13.0.dev0)

```



Discussion:
I also find non tape-mode parallel does not work. Reproduction code is

```
import pennylane as qml


import numpy as np
from pennylane import QNodeCollection

n_qubits = 4
n_batches = 5
dev = qml.device("default.qubit", wires=n_qubits)

# qml.enable_tape()

@qml.qnode(dev)
def circuit(inputs, weights):
    qml.templates.AngleEmbedding(inputs, wires=range(n_qubits))
    qml.templates.BasicEntanglerLayers(weights, wires=range(n_qubits))
    return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]


n_layers = 1
weight_shapes = {"weights": (n_layers, n_qubits)}

qnode = QNodeCollection([circuit, circuit, circuit, circuit])
x = np.random.rand(n_qubits).astype(np.float64)#.to(torch.device('cuda'))
p = np.random.rand(*weight_shapes['weights']).astype(np.float64)#.to(torch.device('cuda'))
res = qnode(x, p, parallel=True)
```

Error trace is
```
  File "/home/shukob/Research/***/***/scripts/debug_parallel.py", line 27, in <module>
    res = qnode(x, p, parallel=True)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/collections/qnode_collection.py", line 271, in __call__
    results = self.evaluate(args, kwargs)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/collections/qnode_collection.py", line 229, in evaluate
    return dask.compute(*results, scheduler=_scheduler)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/base.py", line 452, in compute
    results = schedule(dsk, keys, **kwargs)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/threaded.py", line 84, in get
    **kwargs
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/local.py", line 486, in get_async
    raise_exception(exc, tb)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/local.py", line 316, in reraise
    raise exc
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/local.py", line 222, in execute_task
    result = _execute_task(task, data)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/dask/core.py", line 121, in _execute_task
    return func(*(_execute_task(a, cache) for a in args))
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/interfaces/autograd.py", line 69, in __call__
    return self.evaluate(args, kwargs)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/autograd/tracer.py", line 48, in f_wrapped
    return f_raw(*args, **kwargs)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/qnodes/base.py", line 840, in evaluate
    self._construct(args, kwargs)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/qnodes/jacobian.py", line 90, in _construct
    self.par_to_grad_method = {k: self._best_method(k) for k in self.variable_deps}
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/qnodes/jacobian.py", line 90, in <dictcomp>
    self.par_to_grad_method = {k: self._best_method(k) for k in self.variable_deps}
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/qnodes/qubit.py", line 76, in _best_method
    S = self.circuit.nodes_between(op, ob)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/circuit_graph.py", line 463, in nodes_between
    A = self.descendants([a])
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/circuit_graph.py", line 408, in descendants
    return set().union(*(nx.dag.descendants(self._graph, o) for o in ops)) - set(ops)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/PennyLane-0.13.0.dev0-py3.7.egg/pennylane/circuit_graph.py", line 408, in <genexpr>
    return set().union(*(nx.dag.descendants(self._graph, o) for o in ops)) - set(ops)
  File "/home/shukob/anaconda3/envs/***/lib/python3.7/site-packages/networkx/algorithms/dag.py", line 59, in descendants
    raise nx.NetworkXError(f"The node {source} is not in the graph.")
networkx.exception.NetworkXError: The node RX(0.942264978721866, wires=[0]) is not in the graph.

```
Thanks for catching this @shukob! 

I haven't encountered this before, so I'm unsure of the cause - my best guess could be a race condition during the queuing of the operations into the circuit?
@josh146 Thank you! 

It seems it does not always happen. I will check into the queueing operation.
@josh146  It seems to be a problem of instantiation of the `QNode`s in my code. If I explicitly instantiate them by constructors like below,
```
qnode = QNodeCollection([QNode(circuit, dev) for _ in range(n_batches)])
```
the error disappears. I misunderstood the `qnode` annotation to instantiate the `QNode` class, but it does not, so the same `QNode` instances are used in `QNodeCollection`. This was the problem. Sorry about this.
I am confused, the above seems not to be the case. I need thread locking in `__enter__` / `__exit__`  of `QuantumTape` to make it work in tape mode.

For non-tape mode, explicit instantiation above seems to be working. 

I will check more because it does not always happen.
Thanks for delving into this in more detail @shukob, this is really good to know. Can I ask how you are doing the thread locking in `__enter__` and `__exit__`?
Thanks, @josh146! I have modified `QuantumTape` like

```
    from threading import Lock
    _lock = Lock()

    def __enter__(self):
        QuantumTape._lock.acquire()
        if not QueuingContext.recording():
            # if the tape is the first active queuing context
            # monkeypatch the operations to support the new queuing context
            with contextlib.ExitStack() as stack:
                for mock in mock_operations():
                    stack.enter_context(mock)
                self._stack = stack.pop_all()

        QueuingContext.append(self)
        return super().__enter__()

    def __exit__(self, exception_type, exception_value, traceback):
        super().__exit__(exception_type, exception_value, traceback)

        if not QueuingContext.recording():
            # remove the monkeypatching
            self._stack.__exit__(exception_type, exception_value, traceback)

        self._process_queue()
        QuantumTape._lock.release()
```
Thanks @shukob! We appreciate your work uncovering (and potentially fixing) this issue
@shukob, would you like to make a small PR to add the thread locking?
@co9olguy Thanks!
@josh146 Yes! Willingly.

Resolving patch:
diff --git a/pennylane/tape/tapes/tape.py b/pennylane/tape/tapes/tape.py
--- a/pennylane/tape/tapes/tape.py
+++ b/pennylane/tape/tapes/tape.py
@@ -18,6 +18,7 @@
 import contextlib
 import copy
 from collections import Counter
+from threading import RLock
 
 import numpy as np
 
@@ -220,6 +221,9 @@ class QuantumTape(AnnotatedQueue):
     [0.56, 0.543, 0.133]
     """
 
+    _lock = RLock()
+    """threading.RLock: Used to synchronize appending to/popping from global QueueingContext."""
+
     def __init__(self, name=None):
         super().__init__()
         self.name = name
@@ -262,25 +266,33 @@ def __repr__(self):
         return f"<{self.__class__.__name__}: wires={self.wires.tolist()}, params={self.num_params}>"
 
     def __enter__(self):
-        if not QueuingContext.recording():
-            # if the tape is the first active queuing context
-            # monkeypatch the operations to support the new queuing context
-            with contextlib.ExitStack() as stack:
-                for mock in mock_operations():
-                    stack.enter_context(mock)
-                self._stack = stack.pop_all()
-
-        QueuingContext.append(self)
-        return super().__enter__()
+        QuantumTape._lock.acquire()
+        try:
+            if not QueuingContext.recording():
+                # if the tape is the first active queuing context
+                # monkeypatch the operations to support the new queuing context
+                with contextlib.ExitStack() as stack:
+                    for mock in mock_operations():
+                        stack.enter_context(mock)
+                    self._stack = stack.pop_all()
+
+            QueuingContext.append(self)
+            return super().__enter__()
+        except Exception as _:
+            QuantumTape._lock.release()
+            raise
 
     def __exit__(self, exception_type, exception_value, traceback):
-        super().__exit__(exception_type, exception_value, traceback)
+        try:
+            super().__exit__(exception_type, exception_value, traceback)
 
-        if not QueuingContext.recording():
-            # remove the monkeypatching
-            self._stack.__exit__(exception_type, exception_value, traceback)
+            if not QueuingContext.recording():
+                # remove the monkeypatching
+                self._stack.__exit__(exception_type, exception_value, traceback)
 
-        self._process_queue()
+            self._process_queue()
+        finally:
+            QuantumTape._lock.release()
 
     @property
     def interface(self):


Test patch:
diff --git a/tests/interfaces/test_torch.py b/tests/interfaces/test_torch.py
--- a/tests/interfaces/test_torch.py
+++ b/tests/interfaces/test_torch.py
@@ -643,7 +643,7 @@ def div(a, b):
 
         div(a, b).backward()
         assert a.grad == 1 / b
-        assert b.grad == -a / b ** 2
+        assert torch.isclose(b.grad, -a / b ** 2)
 
     @pytest.mark.parametrize("x, y", gradient_test_data)
     def test_composition_qnodes_gradient(self, qnodes, x, y):
diff --git a/tests/tape/tapes/test_qnode.py b/tests/tape/tapes/test_qnode.py
--- a/tests/tape/tapes/test_qnode.py
+++ b/tests/tape/tapes/test_qnode.py
@@ -16,6 +16,7 @@
 import numpy as np
 
 import pennylane as qml
+from pennylane import QNodeCollection
 from pennylane.tape import JacobianTape, QNode, qnode, QubitParamShiftTape, CVParamShiftTape
 
 
@@ -605,3 +606,43 @@ def func(x, y):
 
         assert np.allclose(res, res2, atol=tol, rtol=0)
         assert func.qtape is not old_tape
+
+
+class TestQNodeCollection:
+    """Unittests for the QNodeCollection"""
+
+    @pytest.fixture
+    def enable_tape_mode(self):
+        qml.enable_tape()
+        yield
+        qml.disable_tape()
+
+    def test_multi_thread(self, enable_tape_mode):
+        """Test that multi-threaded queuing in tape mode works correctly"""
+        n_qubits = 4
+        n_batches = 5
+        dev = qml.device("default.qubit", wires=n_qubits)
+
+
+        def circuit(inputs, weights):
+            for index, input in enumerate(inputs):
+                qml.RY(input, wires=index)
+            for index in range(n_qubits - 1):
+                qml.CNOT(wires=(index, index + 1))
+            for index, weight in enumerate(weights):
+                qml.RX(weight, wires=index)
+            return [qml.expval(qml.PauliZ(wires=i)) for i in range(n_qubits)]
+
+        weight_shapes = {"weights": (n_qubits)}
+
+        try:
+            qnode = QNodeCollection([QNode(circuit, dev) for _ in range(n_batches)])
+        except Exception as e:
+            pytest.fail("QNodeCollection cannot be instantiated")
+        x = np.random.rand(n_qubits).astype(np.float64)
+        p = np.random.rand(weight_shapes["weights"]).astype(np.float64)
+        try:
+            for _ in range(10):
+                qnode(x, p, parallel=True)
+        except:
+            pytest.fail("Multi-threading on QuantumTape failed")