Fix Lightning fallbacks to support PL 0.29 changes

pennylane_lightning/lightning_qubit.py

@@ -744,6 +744,11 @@ def probability(self, wires=None, shot_range=None, bin_size=None):
         Returns:
             array[float]: list of the probabilities
         """
+        if wires and len(wires) > 1 and (not np.all(list(wires)[:-1] <= list(wires)[1:])):
+            raise RuntimeError(
+                "Lightning does not currently support out-of-order indices for probabilities"
+            )
+
         if self.shots is not None:
             return self.estimate_probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
 
@@ -756,6 +761,7 @@ def probability(self, wires=None, shot_range=None, bin_size=None):
         # To support np.complex64 based on the type of self._state
         dtype = self._state.dtype
         ket = np.ravel(self._state)
+
         state_vector = StateVectorC64(ket) if self.use_csingle else StateVectorC128(ket)
         M = MeasuresC64(state_vector) if self.use_csingle else MeasuresC128(state_vector)
         return M.probs(device_wires)

pennylane_lightning/src/bindings/Bindings.cpp

@@ -87,11 +87,12 @@ void lightning_class_bindings(py::module_ &m) {
         .def(py::init<const StateVectorRawCPU<PrecisionT> &>())
         .def("probs",
              [](Measures<PrecisionT> &M, const std::vector<size_t> &wires) {
-                 if (wires.empty()) {
-                     return py::array_t<ParamT>(py::cast(M.probs()));
-                 }
                  return py::array_t<ParamT>(py::cast(M.probs(wires)));
              })
+        .def("probs",
+             [](Measures<PrecisionT> &M) {
+                 return py::array_t<ParamT>(py::cast(M.probs()));
+             })
         .def("expval",
              static_cast<PrecisionT (Measures<PrecisionT>::*)(
                  const std::string &, const std::vector<size_t> &)>(

pennylane_lightning/src/simulator/Measures.hpp

@@ -80,7 +80,9 @@ class Measures {
      * @return Floating point std::vector with probabilities.
      * The basis columns are rearranged according to wires.
      */
-    std::vector<fp_t> probs(const std::vector<size_t> &wires) {
+    std::vector<fp_t>
+    probs(const std::vector<size_t> &wires,
+          [[maybe_unused]] const std::vector<size_t> &device_wires = {}) {
         // Determining index that would sort the vector.
         // This information is needed later.
         const auto sorted_ind_wires = Util::sorting_indices(wires);

tests/test_measures.py

@@ -97,7 +97,6 @@ def circuit():
         "cases",
         [
             [[0, 1], [0.9165164490394898, 0.0, 0.08348355096051052, 0.0]],
-            [[1, 0], [0.9165164490394898, 0.08348355096051052, 0.0, 0.0]],
             [0, [0.9165164490394898, 0.08348355096051052]],
             [[0], [0.9165164490394898, 0.08348355096051052]],
         ],
@@ -116,6 +115,31 @@ def circuit():
 
         assert np.allclose(circuit(), cases[1], atol=tol, rtol=0)
 
+    @pytest.mark.parametrize(
+        "cases",
+        [
+            [[1, 0], [0.9165164490394898, 0.08348355096051052, 0.0, 0.0]],
+            [["a", "0"], [0.9165164490394898, 0.08348355096051052]],
+        ],
+    )
+    def test_fail_probs_tape_unordered_wires(self, cases, tol, dev):
+        """Test probs with a circuit on wires=[0]"""
+
+        x, y, z = [0.5, 0.3, -0.7]
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.4, wires=[0])
+            qml.Rot(x, y, z, wires=[0])
+            qml.RY(-0.2, wires=[0])
+            return qml.probs(wires=cases[0])
+
+        with pytest.raises(
+            RuntimeError,
+            match="Lightning does not currently support out-of-order indices for probabilities",
+        ):
+            assert np.allclose(circuit(), cases[1], atol=tol, rtol=0)
+
     @pytest.mark.parametrize(
         "cases",
         [
@@ -128,6 +152,24 @@ def circuit():
                     0.0013668981445561978,
                 ],
             ],
+            [0, [0.938791280945186, 0.061208719054813635]],
+            [[0], [0.938791280945186, 0.061208719054813635]],
+        ],
+    )
+    def test_probs_tape_wire01(self, cases, tol, dev):
+        """Test probs with a circuit on wires=[0,1]"""
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.RX(0.5, wires=[0])
+            qml.RY(0.3, wires=[1])
+            return qml.probs(wires=cases[0])
+
+        assert np.allclose(circuit(), cases[1], atol=tol, rtol=0)
+
+    @pytest.mark.parametrize(
+        "cases",
+        [
             [
                 [1, 0],
                 [
@@ -137,11 +179,9 @@ def circuit():
                     0.0013668981445561978,
                 ],
             ],
-            [0, [0.938791280945186, 0.061208719054813635]],
-            [[0], [0.938791280945186, 0.061208719054813635]],
         ],
     )
-    def test_probs_tape_wire01(self, cases, tol, dev):
+    def test_fail_probs_tape_wire01(self, cases, tol, dev):
         """Test probs with a circuit on wires=[0,1]"""
 
         @qml.qnode(dev)
@@ -150,7 +190,11 @@ def circuit():
             qml.RY(0.3, wires=[1])
             return qml.probs(wires=cases[0])
 
-        assert np.allclose(circuit(), cases[1], atol=tol, rtol=0)
+        with pytest.raises(
+            RuntimeError,
+            match="Lightning does not currently support out-of-order indices for probabilities",
+        ):
+            assert np.allclose(circuit(), cases[1], atol=tol, rtol=0)
 
 
 class TestExpval: