[WIP] MPI support of mid-measurement

.github/CHANGELOG.md

@@ -2,6 +2,9 @@
 
 ### New features since last release
 
+* Add `mid-circuit measurements` support to `lightning.gpu`'s single-GPU backend.
+  [(#931)](https://github.com/PennyLaneAI/pennylane-lightning/pull/931)
+
 * Add Matrix Product Operator (MPO) for all gates support to `lightning.tensor`. Note current C++ implementation only works for MPO sites data provided by users.
   [(#859)](https://github.com/PennyLaneAI/pennylane-lightning/pull/859)
 

mpitests/conftest.py

@@ -15,9 +15,10 @@
 Pytest configuration file for PennyLane-Lightning-GPU test suite.
 """
 # pylint: disable=missing-function-docstring,wrong-import-order,unused-import
-
 import itertools
 import os
+from functools import reduce
+from typing import Sequence
 
 import pennylane as qml
 import pytest
@@ -125,3 +126,79 @@ def _device(wires):
         )
 
     return _device
+
+
+#######################################################################
+
+
+def validate_counts(shots, results1, results2):
+    """Compares two counts.
+    If the results are ``Sequence``s, loop over entries.
+    Fails if a key of ``results1`` is not found in ``results2``.
+    Passes if counts are too low, chosen as ``100``.
+    Otherwise, fails if counts differ by more than ``20`` plus 20 percent.
+    """
+    if isinstance(results1, Sequence):
+        assert isinstance(results2, Sequence)
+        assert len(results1) == len(results2)
+        for r1, r2 in zip(results1, results2):
+            validate_counts(shots, r1, r2)
+        return
+    for key1, val1 in results1.items():
+        val2 = results2[key1]
+        if abs(val1 + val2) > 100:
+            assert np.allclose(val1, val2, rtol=20, atol=0.2)
+
+
+def validate_samples(shots, results1, results2):
+    """Compares two samples.
+    If the results are ``Sequence``s, loop over entries.
+    Fails if the results do not have the same shape, within ``20`` entries plus 20 percent.
+    This is to handle cases when post-selection yields variable shapes.
+    Otherwise, fails if the sums of samples differ by more than ``20`` plus 20 percent.
+    """
+    if isinstance(shots, Sequence):
+        assert isinstance(results1, Sequence)
+        assert isinstance(results2, Sequence)
+        assert len(results1) == len(results2)
+        for s, r1, r2 in zip(shots, results1, results2):
+            validate_samples(s, r1, r2)
+    else:
+        sh1, sh2 = results1.shape[0], results2.shape[0]
+        assert np.allclose(sh1, sh2, rtol=20, atol=0.2)
+        assert results1.ndim == results2.ndim
+        if results2.ndim > 1:
+            assert results1.shape[1] == results2.shape[1]
+        np.allclose(np.sum(results1), np.sum(results2), rtol=20, atol=0.2)
+
+
+def validate_others(shots, results1, results2):
+    """Compares two expval, probs or var.
+    If the results are ``Sequence``s, validate the average of items.
+    If ``shots is None``, validate using ``np.allclose``'s default parameters.
+    Otherwise, fails if the results do not match within ``0.01`` plus 20 percent.
+    """
+    if isinstance(results1, Sequence):
+        assert isinstance(results2, Sequence)
+        assert len(results1) == len(results2)
+        results1 = reduce(lambda x, y: x + y, results1) / len(results1)
+        results2 = reduce(lambda x, y: x + y, results2) / len(results2)
+        validate_others(shots, results1, results2)
+        return
+    if shots is None:
+        assert np.allclose(results1, results2)
+        return
+    assert np.allclose(results1, results2, atol=0.01, rtol=0.2)
+
+
+def validate_measurements(func, shots, results1, results2):
+    """Calls the correct validation function based on measurement type."""
+    if func is qml.counts:
+        validate_counts(shots, results1, results2)
+        return
+
+    if func is qml.sample:
+        validate_samples(shots, results1, results2)
+        return
+
+    validate_others(shots, results1, results2)