Update probs in Lightning.GPU (#941)

### Before submitting Please complete the following checklist when submitting a PR: - [ ] All new features must include a unit test. If you've fixed a bug or added code that should be tested, add a test to the [`tests`](../tests) directory! - [ ] All new functions and code must be clearly commented and documented. If you do make documentation changes, make sure that the docs build and render correctly by running `make docs`. - [ ] Ensure that the test suite passes, by running `make test`. - [x] Add a new entry to the `.github/CHANGELOG.md` file, summarizing the change, and including a link back to the PR. - [x] Ensure that code is properly formatted by running `make format`. When all the above are checked, delete everything above the dashed line and fill in the pull request template. ------------------------------------------------------------------------------------------------------------ **Context:** [SC-75562] & [SC-73560] **Description of the Change:** **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:** --------- Co-authored-by: ringo-but-quantum <[email protected]>
PennyLaneAI · Oct 10, 2024 · 4945ed0 · 4945ed0
1 parent 6062a75
commit 4945ed0
Show file tree

Hide file tree

Showing 13 changed files with 46 additions and 90 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -43,6 +43,9 @@
 
 ### Improvements
 
+* The `prob` data return `lightning.gpu` C++ layer is aligned with other state-vector backends and `lightning.gpu` supports out-of-order `qml.prob`.
+    [(#941)](https://github.com/PennyLaneAI/pennylane-lightning/pull/941)
+
 * Add `setStateVector(state, wire)` support to the `lightning.gpu` C++ layer.
   [(#930)](https://github.com/PennyLaneAI/pennylane-lightning/pull/930)
 

diff --git a/pennylane_lightning/core/_measurements_base.py b/pennylane_lightning/core/_measurements_base.py
@@ -131,15 +131,6 @@ def expval(self, measurementprocess: MeasurementProcess):
             measurementprocess.obs.name, measurementprocess.obs.wires
         )
 
-    def _probs_retval_conversion(self, probs_results: Any) -> np.ndarray:
-        """Convert the data structure from the C++ backend to a common structure through lightning devices.
-        Args:
-            probs_result (Any): Result provided by C++ backend.
-        Returns:
-            np.ndarray with probabilities of the supplied observable or wires.
-        """
-        return probs_results
-
     def probs(self, measurementprocess: MeasurementProcess):
         """Probabilities of the supplied observable or wires contained in the MeasurementProcess.
 
@@ -161,7 +152,7 @@ def probs(self, measurementprocess: MeasurementProcess):
                 [qml.adjoint(g, lazy=False) for g in reversed(diagonalizing_gates)]
             )
 
-        return self._probs_retval_conversion(results)
+        return results
 
     def var(self, measurementprocess: MeasurementProcess):
         """Variance of the supplied observable contained in the MeasurementProcess.

diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.39.0-dev41"
+__version__ = "0.39.0-dev42"
diff --git a/pennylane_lightning/core/src/measurements/tests/Test_MeasurementsBase.cpp b/pennylane_lightning/core/src/measurements/tests/Test_MeasurementsBase.cpp
@@ -85,44 +85,33 @@ template <typename TypeList> void testProbabilities() {
         // Expected results calculated with Pennylane default.qubit:
         std::vector<
             std::pair<std::vector<std::size_t>, std::vector<PrecisionT>>>
-            input = {
-#if defined(_ENABLE_PLGPU)
-                // Bit index reodering conducted in the python layer
-                // for L-GPU. Also L-GPU backend doesn't support
-                // out of order wires for probability calculation
-                {{2, 1, 0},
-                 {0.67078706, 0.03062806, 0.0870997, 0.00397696, 0.17564072,
-                  0.00801973, 0.02280642, 0.00104134}}
-#else
-                // LightningQubit currently supports arbitrary wire index
-                // ordering.
-                {{0, 2, 1},
-                 {0.67078706, 0.0870997, 0.03062806, 0.00397696, 0.17564072,
-                  0.02280642, 0.00801973, 0.00104134}},
-                {{1, 0, 2},
-                 {0.67078706, 0.03062806, 0.17564072, 0.00801973, 0.0870997,
-                  0.00397696, 0.02280642, 0.00104134}},
-                {{1, 2, 0},
-                 {0.67078706, 0.17564072, 0.03062806, 0.00801973, 0.0870997,
-                  0.02280642, 0.00397696, 0.00104134}},
-                {{2, 0, 1},
-                 {0.67078706, 0.0870997, 0.17564072, 0.02280642, 0.03062806,
-                  0.00397696, 0.00801973, 0.00104134}},
-                {{2, 1, 0},
-                 {0.67078706, 0.17564072, 0.0870997, 0.02280642, 0.03062806,
-                  0.00801973, 0.00397696, 0.00104134}},
-                {{2, 1}, {0.84642778, 0.10990612, 0.0386478, 0.0050183}},
-                {{0, 1, 2},
-                 {0.67078706, 0.03062806, 0.0870997, 0.00397696, 0.17564072,
-                  0.00801973, 0.02280642, 0.00104134}},
-                {{0, 1}, {0.70141512, 0.09107666, 0.18366045, 0.02384776}},
-                {{0, 2}, {0.75788676, 0.03460502, 0.19844714, 0.00906107}},
-                {{1, 2}, {0.84642778, 0.0386478, 0.10990612, 0.0050183}},
-                {{0}, {0.79249179, 0.20750821}},
-                {{1}, {0.88507558, 0.11492442}},
-                {{2}, {0.9563339, 0.0436661}}
-#endif
-            };
+            input = {// LightningQubit currently supports arbitrary wire index
+                     // ordering.
+                     {{0, 2, 1},
+                      {0.67078706, 0.0870997, 0.03062806, 0.00397696,
+                       0.17564072, 0.02280642, 0.00801973, 0.00104134}},
+                     {{1, 0, 2},
+                      {0.67078706, 0.03062806, 0.17564072, 0.00801973,
+                       0.0870997, 0.00397696, 0.02280642, 0.00104134}},
+                     {{1, 2, 0},
+                      {0.67078706, 0.17564072, 0.03062806, 0.00801973,
+                       0.0870997, 0.02280642, 0.00397696, 0.00104134}},
+                     {{2, 0, 1},
+                      {0.67078706, 0.0870997, 0.17564072, 0.02280642,
+                       0.03062806, 0.00397696, 0.00801973, 0.00104134}},
+                     {{2, 1, 0},
+                      {0.67078706, 0.17564072, 0.0870997, 0.02280642,
+                       0.03062806, 0.00801973, 0.00397696, 0.00104134}},
+                     {{2, 1}, {0.84642778, 0.10990612, 0.0386478, 0.0050183}},
+                     {{0, 1, 2},
+                      {0.67078706, 0.03062806, 0.0870997, 0.00397696,
+                       0.17564072, 0.00801973, 0.02280642, 0.00104134}},
+                     {{0, 1}, {0.70141512, 0.09107666, 0.18366045, 0.02384776}},
+                     {{0, 2}, {0.75788676, 0.03460502, 0.19844714, 0.00906107}},
+                     {{1, 2}, {0.84642778, 0.0386478, 0.10990612, 0.0050183}},
+                     {{0}, {0.79249179, 0.20750821}},
+                     {{1}, {0.88507558, 0.11492442}},
+                     {{2}, {0.9563339, 0.0436661}}};
 
         // Defining the Statevector that will be measured.
         auto statevector_data = createNonTrivialState<StateVectorT>();
@@ -404,11 +393,7 @@ template <typename TypeList> void testProbabilitiesObsShots() {
             std::size_t num_shots = 10000;
             auto prob_obs_shots = Measurer_obs_shots.probs(*obs, num_shots);
 
-#ifdef _ENABLE_PLGPU
-            auto prob = Measurer.probs(std::vector<std::size_t>({2, 1, 0}));
-#else
             auto prob = Measurer.probs(std::vector<std::size_t>({0, 1, 2}));
-#endif
 
             REQUIRE_THAT(prob_obs_shots, Catch::Approx(prob).margin(5e-2));
         }
@@ -434,11 +419,7 @@ template <typename TypeList> void testProbabilitiesObsShots() {
 
             std::size_t num_shots = 10000;
             auto prob_obs_shots = Measurer_obs_shots.probs(*obs, num_shots);
-#ifdef _ENABLE_PLGPU
-            auto prob = Measurer.probs(std::vector<std::size_t>({2, 1, 0}));
-#else
             auto prob = Measurer.probs(std::vector<std::size_t>({0, 1, 2}));
-#endif
 
             REQUIRE_THAT(prob_obs_shots, Catch::Approx(prob).margin(5e-2));
         }

diff --git a/pennylane_lightning/core/src/measurements/tests/mpi/Test_MeasurementsBaseMPI.cpp b/pennylane_lightning/core/src/measurements/tests/mpi/Test_MeasurementsBaseMPI.cpp
@@ -58,7 +58,7 @@ template <typename TypeList> void testProbabilities() {
             input = {// Bit index reodering conducted in the python layer
                      // for L-GPU. Also L-GPU backend doesn't support
                      // out of order wires for probability calculation
-                     {{2, 1, 0},
+                     {{0, 1, 2},
                       {0.67078706, 0.03062806, 0.0870997, 0.00397696,
                        0.17564072, 0.00801973, 0.02280642, 0.00104134}}};
 
@@ -386,7 +386,7 @@ template <typename TypeList> void testProbabilitiesObsShots() {
 
             std::size_t num_shots = 10000;
             auto prob_obs_shots = Measurer_obs_shots.probs(*obs, num_shots);
-            auto prob = Measurer.probs(std::vector<std::size_t>({2, 1, 0}));
+            auto prob = Measurer.probs(std::vector<std::size_t>({0, 1, 2}));
             auto prob_all = mpi_manager.allgather(prob);
             REQUIRE_THAT(prob_obs_shots, Catch::Approx(prob_all).margin(5e-2));
         }

diff --git a/pennylane_lightning/core/src/simulators/lightning_gpu/measurements/MeasurementsGPU.hpp b/pennylane_lightning/core/src/simulators/lightning_gpu/measurements/MeasurementsGPU.hpp
@@ -94,16 +94,10 @@ class Measurements final
      */
     auto probs(const std::vector<std::size_t> &wires)
         -> std::vector<PrecisionT> {
-        PL_ABORT_IF_NOT(std::is_sorted(wires.cbegin(), wires.cend()) ||
-                            std::is_sorted(wires.rbegin(), wires.rend()),
-                        "LightningGPU does not currently support out-of-order "
-                        "wire indices with probability calculations");
-
         // Data return type fixed as double in custatevec function call
         std::vector<double> probabilities(Pennylane::Util::exp2(wires.size()));
         // this should be built upon by the wires not participating
-        int maskLen =
-            0; // static_cast<int>(BaseType::getNumQubits() - wires.size());
+        int maskLen = 0;
         int *maskBitString = nullptr; //
         int *maskOrdering = nullptr;
 
@@ -125,6 +119,8 @@ class Measurements final
                                this->_statevector.getNumQubits() - 1 - x);
                        });
 
+        std::reverse(wires_int.begin(), wires_int.end());
+
         PL_CUSTATEVEC_IS_SUCCESS(custatevecAbs2SumArray(
             /* custatevecHandle_t */ this->_statevector.getCusvHandle(),
             /* const void* */ this->_statevector.getData(),

diff --git a/pennylane_lightning/core/src/simulators/lightning_gpu/measurements/MeasurementsGPUMPI.hpp b/pennylane_lightning/core/src/simulators/lightning_gpu/measurements/MeasurementsGPUMPI.hpp
@@ -130,6 +130,8 @@ class MeasurementsMPI final
             }
         }
 
+        std::reverse(wires_local.begin(), wires_local.end());
+
         std::vector<double> local_probabilities(
             Pennylane::Util::exp2(wires_local.size()));
 

diff --git a/...e/src/simulators/lightning_gpu/measurements/tests/Test_StateVectorCudaManaged_Measure.cpp b/...e/src/simulators/lightning_gpu/measurements/tests/Test_StateVectorCudaManaged_Measure.cpp
@@ -257,7 +257,7 @@ TEMPLATE_TEST_CASE("Probabilities", "[Measures]", float, double) {
     using StateVectorT = StateVectorCudaManaged<TestType>;
     // Probabilities calculated with Pennylane default.qubit:
     std::vector<std::pair<std::vector<std::size_t>, std::vector<TestType>>>
-        input = {{{2, 1, 0},
+        input = {{{0, 1, 2},
                   {0.67078706, 0.03062806, 0.0870997, 0.00397696, 0.17564072,
                    0.00801973, 0.02280642, 0.00104134}}};
 

diff --git a/...e/src/simulators/lightning_gpu/measurements/tests/mpi/Test_StateVectorCudaMPI_Measure.cpp b/...e/src/simulators/lightning_gpu/measurements/tests/mpi/Test_StateVectorCudaMPI_Measure.cpp
@@ -411,7 +411,7 @@ TEMPLATE_TEST_CASE("Probabilities", "[MeasuresMPI]", double) {
     using StateVectorT = StateVectorCudaMPI<TestType>;
     // Probabilities calculated with Pennylane default.qubit:
     std::vector<std::pair<std::vector<std::size_t>, std::vector<TestType>>>
-        input = {{{2, 1, 0},
+        input = {{{0, 1, 2},
                   {0.67078706, 0.03062806, 0.0870997, 0.00397696, 0.17564072,
                    0.00801973, 0.02280642, 0.00104134}}};
 

diff --git a/pennylane_lightning/core/src/utils/Util.hpp b/pennylane_lightning/core/src/utils/Util.hpp
@@ -21,12 +21,12 @@
 #include <algorithm>
 #include <cmath>
 #include <complex>
+#include <concepts> // integral, floating_point
 #include <numbers>
 #include <numeric> // transform_reduce
 #include <set>
 #include <type_traits> // is_same_v
 #include <vector>
-#include <concepts> // integral, floating_point
 
 #include "Error.hpp"
 #include "TypeTraits.hpp" // remove_complex_t

diff --git a/pennylane_lightning/lightning_gpu/_measurements.py b/pennylane_lightning/lightning_gpu/_measurements.py
@@ -34,7 +34,7 @@
 except ImportError as error_import:
     warn(str(error_import), UserWarning)
 
-from typing import Any, List
+from typing import List
 
 import numpy as np
 import pennylane as qml
@@ -200,20 +200,3 @@ def expval(self, measurementprocess: MeasurementProcess):
         return self._measurement_lightning.expval(
             measurementprocess.obs.name, measurementprocess.obs.wires
         )
-
-    def _probs_retval_conversion(self, probs_results: Any) -> np.ndarray:
-        """Convert the data structure from the C++ backend to a common structure through lightning devices.
-
-        Args:
-            probs_result (Any): Result provided by C++ backend.
-
-        Returns:
-            np.ndarray with probabilities of the supplied observable or wires.
-        """
-
-        # Device returns as col-major orderings, so perform transpose on data for bit-index shuffle for now.
-        if len(probs_results) > 0:
-            num_local_wires = len(probs_results).bit_length() - 1 if len(probs_results) > 0 else 0
-            return probs_results.reshape([2] * num_local_wires).transpose().reshape(-1)
-
-        return probs_results
diff --git a/tests/lightning_qubit/test_measurements_class.py b/tests/lightning_qubit/test_measurements_class.py
@@ -669,7 +669,7 @@ def test_double_return_value(self, shots, measurement, obs0_, obs1_, lightning_s
             assert np.allclose(r, e, atol=dtol, rtol=dtol)
 
     @pytest.mark.skipif(
-        device_name in ("lightning.gpu", "lightning.tensor"),
+        device_name in ("lightning.tensor"),
         reason=f"{device_name} does not support out of order probs.",
     )
     @pytest.mark.parametrize(

diff --git a/tests/test_measurements.py b/tests/test_measurements.py
@@ -151,8 +151,8 @@ def circuit():
             _ = circuit()
 
     @pytest.mark.skipif(
-        device_name in ("lightning.gpu", "lightning.tensor"),
-        reason="lightning.gpu/lightning.tensor does not support out of order prob.",
+        device_name in ("lightning.tensor"),
+        reason="lightning.tensor does not support out of order prob.",
     )
     @pytest.mark.parametrize(
         "cases",