Repository: PennyLaneAI/pennylane
Pull request #3811 (2023-02-23 09:04:58)

Raise error when using `qml.probs` in the computational basis with non-commuting observables
This PR changes the method `MeasurementProcess.samples_computational_basis` to `MeasurementProcess.computational_basis`, which now returns `True` for all measurements that don't have `obs` defined (except for classical shadow measurements and `MidMeasure`).

**Note**: A lot of tests were using `qml.probs` in the computational basis with non-commuting observables! We should also make sure that none of the demos are doing the same thing.

**Note**: This change might break user code!

Fixes https://github.com/PennyLaneAI/pennylane/issues/3548

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
AlbertMitjans: [sc-30770]
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3811?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> :exclamation: No coverage uploaded for pull request base (`v0.29.0-rc0@0228133`). [Click here to learn what that means](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#section-missing-base-commit).
> The diff coverage is `n/a`.

```diff
@@              Coverage Diff               @@
##             v0.29.0-rc0    #3811   +/-   ##
==============================================
  Coverage               ?   99.72%           
==============================================
  Files                  ?      336           
  Lines                  ?    29451           
  Branches               ?        0           
==============================================
  Hits                   ?    29371           
  Misses                 ?       80           
  Partials               ?        0           
```



Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


Patch:
diff --git a/doc/releases/changelog-0.29.0.md b/doc/releases/changelog-0.29.0.md
index 211de7486..0e381983e 100644
--- a/doc/releases/changelog-0.29.0.md
+++ b/doc/releases/changelog-0.29.0.md
@@ -781,6 +781,10 @@
 
 <h3>Bug fixes</h3>
 
+* Fixed a bug where measuring ``qml.probs`` in the computational basis with non-commuting 
+  measurements returned incorrect results. Now an error is raised.
+  [(#3811)](https://github.com/PennyLaneAI/pennylane/pull/3811)
+
 * Fixed a bug in the drawer where nested controlled operations would output
   the label of the operation being controlled, rather than the control values.
   [(#3745)](https://github.com/PennyLaneAI/pennylane/pull/3745)
diff --git a/pennylane/measurements/classical_shadow.py b/pennylane/measurements/classical_shadow.py
index 2ba4071c4..0e3adb4f3 100644
--- a/pennylane/measurements/classical_shadow.py
+++ b/pennylane/measurements/classical_shadow.py
@@ -304,6 +304,10 @@ class ClassicalShadowMP(MeasurementTransform):
 
         return qml.math.cast(qml.math.stack([outcomes, recipes]), dtype=np.int8)
 
+    @property
+    def samples_computational_basis(self):
+        return False
+
     @property
     def numeric_type(self):
         return int
@@ -359,6 +363,10 @@ class ShadowExpvalMP(MeasurementTransform):
         shadow = qml.shadows.ClassicalShadow(bits, recipes, wire_map=self.wires.tolist())
         return shadow.expval(self.H, self.k)
 
+    @property
+    def samples_computational_basis(self):
+        return False
+
     @property
     def numeric_type(self):
         return float
diff --git a/pennylane/measurements/counts.py b/pennylane/measurements/counts.py
index 816c9329f..8fad907e7 100644
--- a/pennylane/measurements/counts.py
+++ b/pennylane/measurements/counts.py
@@ -179,10 +179,6 @@ class CountsMP(SampleMeasurement):
     def return_type(self):
         return AllCounts if self.all_outcomes else Counts
 
-    @property
-    def samples_computational_basis(self):
-        return self.obs is None
-
     def process_samples(
         self,
         samples: Sequence[complex],
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index ea0d9b60c..317be85a7 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -365,10 +365,8 @@ class MeasurementProcess(ABC):
 
     @property
     def samples_computational_basis(self):
-        r"""Bool: Whether or not the MeasurementProcess returns samples in the computational basis or counts of
-        computational basis states.
-        """
-        return False
+        r"""Bool: Whether or not the MeasurementProcess measures in the computational basis."""
+        return self.obs is None
 
     def expand(self):
         """Expand the measurement of an observable to a unitary
diff --git a/pennylane/measurements/mid_measure.py b/pennylane/measurements/mid_measure.py
index 6192140e2..4b3f94ad1 100644
--- a/pennylane/measurements/mid_measure.py
+++ b/pennylane/measurements/mid_measure.py
@@ -106,6 +106,10 @@ class MidMeasureMP(MeasurementProcess):
     def return_type(self):
         return MidMeasure
 
+    @property
+    def samples_computational_basis(self):
+        return False
+
     @property
     def _queue_category(self):
         return "_ops"
diff --git a/pennylane/measurements/sample.py b/pennylane/measurements/sample.py
index 7906704ab..389a4697f 100644
--- a/pennylane/measurements/sample.py
+++ b/pennylane/measurements/sample.py
@@ -149,10 +149,6 @@ class SampleMP(SampleMeasurement):
         every_term_standard = all(o.__class__ in int_eigval_obs for o in tensor_terms)
         return int if every_term_standard else float
 
-    @property
-    def samples_computational_basis(self):
-        return self.obs is None
-
     # pylint: disable=protected-access
     def shape(self, device=None):
         if qml.active_return():
diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py
index e4ae6e629..0f2c3107c 100644
--- a/pennylane/tape/qscript.py
+++ b/pennylane/tape/qscript.py
@@ -312,7 +312,7 @@ class QuantumScript:
 
     @property
     def samples_computational_basis(self):
-        """Determines if any of the measurements do sampling/counting in the computational basis."""
+        """Determines if any of the measurements are in the computational basis."""
         return any(o.samples_computational_basis for o in self.measurements)
 
     @property


Test patch:
diff --git a/tests/circuit_graph/test_circuit_graph.py b/tests/circuit_graph/test_circuit_graph.py
index 3a5817984..0318cba28 100644
--- a/tests/circuit_graph/test_circuit_graph.py
+++ b/tests/circuit_graph/test_circuit_graph.py
@@ -84,15 +84,15 @@ def circuit_measure_max_once():
 
 def circuit_measure_max_twice():
     """A fixture of a circuit that measures wire 0 twice."""
-    return qml.expval(qml.PauliX(wires=0)), qml.probs(wires=0)
+    return qml.expval(qml.PauliZ(wires=0)), qml.probs(wires=0)
 
 
 def circuit_measure_multiple_with_max_twice():
     """A fixture of a circuit that measures wire 0 twice."""
     return (
-        qml.expval(qml.PauliX(wires=0)),
+        qml.expval(qml.PauliZ(wires=0)),
         qml.probs(wires=[0, 1, 2]),
-        qml.var(qml.PauliX(wires=[1]) @ qml.PauliZ([2])),
+        qml.var(qml.PauliZ(wires=[1]) @ qml.PauliZ([2])),
     )
 
 
diff --git a/tests/measurements/test_measurements.py b/tests/measurements/test_measurements.py
index 80d7f5615..62114b5d5 100644
--- a/tests/measurements/test_measurements.py
+++ b/tests/measurements/test_measurements.py
@@ -381,6 +381,10 @@ class TestExpansion:
             CountsMP(wires=["a", 1], all_outcomes=True),
             CountsMP(),
             CountsMP(wires=["a", 1]),
+            StateMP(),
+            VnEntropyMP(wires=["a", 1]),
+            MutualInfoMP(wires=[["a", 1], ["b", 2]]),
+            ProbabilityMP(wires=["a", 1]),
         ],
     )
     def test_samples_computational_basis_true(self, m):
@@ -393,14 +397,10 @@ class TestExpansion:
             ExpectationMP(obs=qml.PauliX(2)),
             VarianceMP(obs=qml.PauliX("a")),
             ProbabilityMP(obs=qml.PauliX("b")),
-            ProbabilityMP(wires=["a", 1]),
             SampleMP(obs=qml.PauliX("a")),
             CountsMP(obs=qml.PauliX("a")),
-            StateMP(),
-            VnEntropyMP(wires=["a", 1]),
-            MutualInfoMP(wires=[["a", 1], ["b", 2]]),
-            ClassicalShadowMP(wires=[["a", 1], ["b", 2]]),
             ShadowExpvalMP(H=qml.PauliX("a")),
+            ClassicalShadowMP(wires=[["a", 1], ["b", 2]]),
         ],
     )
     def test_samples_computational_basis_false(self, m):
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index 9a38a5ce0..18ad20df8 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -669,3 +669,51 @@ class TestProbs:
         )
 
         assert np.allclose(res, expected)
+
+    def test_non_commuting_probs_raises_error(self):
+        dev = qml.device("default.qubit", wires=5)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliX(0)), qml.probs(wires=[0, 1])
+
+        with pytest.raises(
+            qml.QuantumFunctionError, match="Only observables that are qubit-wise commuting"
+        ):
+            circuit(1, 2)
+
+    def test_commuting_probs_in_computational_basis(self):
+        """Test that `qml.probs` can be used in the computational basis with other commuting observables."""
+        dev = qml.device("default.qubit", wires=5)
+
+        @qml.qnode(dev)
+        def circuit(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1])
+
+        res = circuit(1, 2)
+
+        @qml.qnode(dev)
+        def circuit2(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0))
+
+        @qml.qnode(dev)
+        def circuit3(x, y):
+            qml.RX(x, wires=[0])
+            qml.RY(y, wires=[1])
+            qml.CNOT(wires=[0, 1])
+            return qml.probs(wires=[0, 1])
+
+        res2 = circuit2(1, 2)
+        res3 = circuit3(1, 2)
+
+        assert res[0] == res2
+        assert qml.math.allequal(res[1:], res3)
diff --git a/tests/returntypes/autograd/test_autograd_qnode_new.py b/tests/returntypes/autograd/test_autograd_qnode_new.py
index 092e6f338..0df7f6a69 100644
--- a/tests/returntypes/autograd/test_autograd_qnode_new.py
+++ b/tests/returntypes/autograd/test_autograd_qnode_new.py
@@ -2153,7 +2153,7 @@ class TestReturn:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def circuit_stack(x, y):
             return anp.hstack(circuit(x, y))
@@ -2167,10 +2167,10 @@ class TestReturn:
         assert len(hess) == 2
 
         assert isinstance(hess[0], np.ndarray)
-        assert hess[0].shape == (10,)
+        assert hess[0].shape == (6,)
 
         assert isinstance(hess[1], np.ndarray)
-        assert hess[1].shape == (10,)
+        assert hess[1].shape == (6,)
 
     def test_hessian_expval_probs_multiple_param_array(self, dev_name, diff_method, mode):
         """The hessian of multiple measurements with a multiple param array return a single array."""
@@ -2187,7 +2187,7 @@ class TestReturn:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def cost(x):
             return anp.hstack(circuit(x))
@@ -2195,7 +2195,7 @@ class TestReturn:
         hess = qml.jacobian(qml.jacobian(cost))(params)
 
         assert isinstance(hess, np.ndarray)
-        assert hess.shape == (5, 2, 2)
+        assert hess.shape == (3, 2, 2)
 
     def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, mode):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
@@ -2214,7 +2214,7 @@ class TestReturn:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def circuit_stack(x, y):
             return anp.hstack(circuit(x, y))
@@ -2228,10 +2228,10 @@ class TestReturn:
         assert len(hess) == 2
 
         assert isinstance(hess[0], np.ndarray)
-        assert hess[0].shape == (10,)
+        assert hess[0].shape == (6,)
 
         assert isinstance(hess[1], np.ndarray)
-        assert hess[1].shape == (10,)
+        assert hess[1].shape == (6,)
 
     def test_hessian_var_multiple_param_array(self, dev_name, diff_method, mode):
         """The hessian of multiple measurements with a multiple param array return a single array."""
@@ -2249,7 +2249,7 @@ class TestReturn:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def cost(x):
             return anp.hstack(circuit(x))
@@ -2257,4 +2257,4 @@ class TestReturn:
         hess = qml.jacobian(qml.jacobian(cost))(params)
 
         assert isinstance(hess, np.ndarray)
-        assert hess.shape == (5, 2, 2)
+        assert hess.shape == (3, 2, 2)
diff --git a/tests/returntypes/autograd/test_autograd_qnode_shot_vector.py b/tests/returntypes/autograd/test_autograd_qnode_shot_vector.py
index 05bc79693..e67418392 100644
--- a/tests/returntypes/autograd/test_autograd_qnode_shot_vector.py
+++ b/tests/returntypes/autograd/test_autograd_qnode_shot_vector.py
@@ -501,7 +501,7 @@ class TestReturnShotVectorHessian:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def cost(x, y):
             def cost2(x, y):
@@ -516,7 +516,7 @@ class TestReturnShotVectorHessian:
         assert len(hess) == 2
         for h in hess:
             assert isinstance(h, np.ndarray)
-            assert h.shape == (2, num_copies, 5)
+            assert h.shape == (2, num_copies, 3)
 
     def test_hessian_expval_probs_multiple_param_array(
         self, dev_name, diff_method, gradient_kwargs, shots, num_copies
@@ -534,7 +534,7 @@ class TestReturnShotVectorHessian:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def cost(x):
             def cost2(x):
@@ -546,7 +546,7 @@ class TestReturnShotVectorHessian:
         hess = qml.jacobian(cost)(params)
 
         assert isinstance(hess, np.ndarray)
-        assert hess.shape == (num_copies, 5, 2, 2)
+        assert hess.shape == (num_copies, 3, 2, 2)
 
 
 shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)]
diff --git a/tests/returntypes/core/test_new_return_types.py b/tests/returntypes/core/test_new_return_types.py
index e6dd26e05..672e0bcf4 100644
--- a/tests/returntypes/core/test_new_return_types.py
+++ b/tests/returntypes/core/test_new_return_types.py
@@ -758,18 +758,18 @@ class TestSameMeasurementShotVector:
 pauliz_w2 = qml.PauliZ(wires=2)
 proj_w2 = qml.Projector([1], wires=2)
 hermitian = qml.Hermitian(np.diag([1, 2]), wires=0)
-tensor_product = qml.PauliY(wires=2) @ qml.PauliX(wires=1)
+tensor_product = qml.PauliZ(wires=2) @ qml.PauliX(wires=1)
 
 # Expval/Var with Probs
 
 scalar_probs_multi = [
     # Expval
     (qml.expval(pauliz_w2), qml.probs(wires=[2, 0])),
-    (qml.expval(proj_w2), qml.probs(wires=[2, 0])),
+    (qml.expval(proj_w2), qml.probs(wires=[1, 0])),
     (qml.expval(tensor_product), qml.probs(wires=[2, 0])),
     # Var
     (qml.var(qml.PauliZ(wires=1)), qml.probs(wires=[0, 1])),
-    (qml.var(proj_w2), qml.probs(wires=[2, 0])),
+    (qml.var(proj_w2), qml.probs(wires=[1, 0])),
     (qml.var(tensor_product), qml.probs(wires=[2, 0])),
 ]
 
diff --git a/tests/returntypes/core/test_new_return_types_qnode.py b/tests/returntypes/core/test_new_return_types_qnode.py
index bc4330212..4bd9c9ffa 100644
--- a/tests/returntypes/core/test_new_return_types_qnode.py
+++ b/tests/returntypes/core/test_new_return_types_qnode.py
@@ -2261,18 +2261,18 @@ class TestIntegrationSameMeasurementShotVector:
 pauliz_w2 = qml.PauliZ(wires=2)
 proj_w2 = qml.Projector([1], wires=2)
 hermitian = qml.Hermitian(np.diag([1, 2]), wires=0)
-tensor_product = qml.PauliY(wires=2) @ qml.PauliX(wires=1)
+tensor_product = qml.PauliZ(wires=2) @ qml.PauliX(wires=1)
 
 # Expval/Var with Probs
 
 scalar_probs_multi = [
     # Expval
     (qml.expval(pauliz_w2), qml.probs(wires=[2, 0])),
-    (qml.expval(proj_w2), qml.probs(wires=[2, 0])),
+    (qml.expval(proj_w2), qml.probs(wires=[1, 0])),
     (qml.expval(tensor_product), qml.probs(wires=[2, 0])),
     # Var
     (qml.var(qml.PauliZ(wires=1)), qml.probs(wires=[0, 1])),
-    (qml.var(proj_w2), qml.probs(wires=[2, 0])),
+    (qml.var(proj_w2), qml.probs(wires=[1, 0])),
     (qml.var(tensor_product), qml.probs(wires=[2, 0])),
 ]
 
diff --git a/tests/returntypes/jax/test_jax_jit_qnode_new.py b/tests/returntypes/jax/test_jax_jit_qnode_new.py
index 51cee13c1..8b7b52a68 100644
--- a/tests/returntypes/jax/test_jax_jit_qnode_new.py
+++ b/tests/returntypes/jax/test_jax_jit_qnode_new.py
@@ -2528,7 +2528,7 @@ class TestReturnHessian:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.jit(hessian(circuit, argnums=[0, 1]))(par_0, par_1)
 
@@ -2546,7 +2546,7 @@ class TestReturnHessian:
         for h in hess[1]:
             assert isinstance(h, tuple)
             for h_comp in h:
-                assert h_comp.shape == (4,)
+                assert h_comp.shape == (2,)
 
     def test_hessian_probs_expval_multiple_param_array(
         self, dev_name, diff_method, hessian, shots, mode, interface
@@ -2572,7 +2572,7 @@ class TestReturnHessian:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.jit(hessian(circuit))(params)
 
@@ -2582,7 +2582,7 @@ class TestReturnHessian:
         assert hess[0].shape == (2, 2)
 
         assert isinstance(hess[1], jax.numpy.ndarray)
-        assert hess[1].shape == (4, 2, 2)
+        assert hess[1].shape == (2, 2, 2)
 
     def test_hessian_probs_var_multiple_params(
         self, dev_name, diff_method, hessian, shots, mode, interface
@@ -2603,7 +2603,7 @@ class TestReturnHessian:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.jit(hessian(circuit, argnums=[0, 1]))(par_0, par_1)
 
@@ -2621,7 +2621,7 @@ class TestReturnHessian:
         for h in hess[1]:
             assert isinstance(h, tuple)
             for h_comp in h:
-                assert h_comp.shape == (4,)
+                assert h_comp.shape == (2,)
 
     def test_hessian_probs_var_multiple_param_array(
         self, dev_name, diff_method, hessian, shots, mode, interface
@@ -2641,7 +2641,7 @@ class TestReturnHessian:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.jit(hessian(circuit))(params)
 
@@ -2651,7 +2651,7 @@ class TestReturnHessian:
         assert hess[0].shape == (2, 2)
 
         assert isinstance(hess[1], jax.numpy.ndarray)
-        assert hess[1].shape == (4, 2, 2)
+        assert hess[1].shape == (2, 2, 2)
 
 
 @pytest.mark.parametrize("hessian", hessian_fn)
diff --git a/tests/returntypes/jax/test_jax_qnode_new.py b/tests/returntypes/jax/test_jax_qnode_new.py
index 5c684820f..4435d9b91 100644
--- a/tests/returntypes/jax/test_jax_qnode_new.py
+++ b/tests/returntypes/jax/test_jax_qnode_new.py
@@ -2228,7 +2228,7 @@ class TestReturn:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1)
 
@@ -2255,15 +2255,15 @@ class TestReturn:
         assert isinstance(hess[1][0], tuple)
         assert len(hess[1][0]) == 2
         assert isinstance(hess[1][0][0], jax.numpy.ndarray)
-        assert hess[1][0][0].shape == (4,)
+        assert hess[1][0][0].shape == (2,)
         assert isinstance(hess[1][0][1], jax.numpy.ndarray)
-        assert hess[1][0][1].shape == (4,)
+        assert hess[1][0][1].shape == (2,)
         assert isinstance(hess[1][1], tuple)
         assert len(hess[1][1]) == 2
         assert isinstance(hess[1][1][0], jax.numpy.ndarray)
-        assert hess[1][1][0].shape == (4,)
+        assert hess[1][1][0].shape == (2,)
         assert isinstance(hess[1][1][1], jax.numpy.ndarray)
-        assert hess[1][1][1].shape == (4,)
+        assert hess[1][1][1].shape == (2,)
 
     def test_hessian_expval_probs_multiple_param_array(
         self, dev_name, diff_method, mode, shots, interface
@@ -2290,7 +2290,7 @@ class TestReturn:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.hessian(circuit)(params)
 
@@ -2301,7 +2301,7 @@ class TestReturn:
         assert hess[0].shape == (2, 2)
 
         assert isinstance(hess[1], jax.numpy.ndarray)
-        assert hess[1].shape == (4, 2, 2)
+        assert hess[1].shape == (2, 2, 2)
 
     def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, mode, shots, interface):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
@@ -2322,7 +2322,7 @@ class TestReturn:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1)
 
@@ -2349,15 +2349,15 @@ class TestReturn:
         assert isinstance(hess[1][0], tuple)
         assert len(hess[1][0]) == 2
         assert isinstance(hess[1][0][0], jax.numpy.ndarray)
-        assert hess[1][0][0].shape == (4,)
+        assert hess[1][0][0].shape == (2,)
         assert isinstance(hess[1][0][1], jax.numpy.ndarray)
-        assert hess[1][0][1].shape == (4,)
+        assert hess[1][0][1].shape == (2,)
         assert isinstance(hess[1][1], tuple)
         assert len(hess[1][1]) == 2
         assert isinstance(hess[1][1][0], jax.numpy.ndarray)
-        assert hess[1][1][0].shape == (4,)
+        assert hess[1][1][0].shape == (2,)
         assert isinstance(hess[1][1][1], jax.numpy.ndarray)
-        assert hess[1][1][1].shape == (4,)
+        assert hess[1][1][1].shape == (2,)
 
     def test_hessian_var_probs_multiple_param_array(
         self, dev_name, diff_method, mode, shots, interface
@@ -2379,7 +2379,7 @@ class TestReturn:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.hessian(circuit)(params)
 
@@ -2390,4 +2390,4 @@ class TestReturn:
         assert hess[0].shape == (2, 2)
 
         assert isinstance(hess[1], jax.numpy.ndarray)
-        assert hess[1].shape == (4, 2, 2)
+        assert hess[1].shape == (2, 2, 2)
diff --git a/tests/returntypes/jax/test_jax_qnode_shot_vector.py b/tests/returntypes/jax/test_jax_qnode_shot_vector.py
index e8c2844ad..b32413836 100644
--- a/tests/returntypes/jax/test_jax_qnode_shot_vector.py
+++ b/tests/returntypes/jax/test_jax_qnode_shot_vector.py
@@ -603,7 +603,7 @@ class TestReturnWithShotVectors:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1)
 
@@ -635,15 +635,15 @@ class TestReturnWithShotVectors:
             assert isinstance(h[1][0], tuple)
             assert len(h[1][0]) == 2
             assert isinstance(h[1][0][0], jax.numpy.ndarray)
-            assert h[1][0][0].shape == (4,)
+            assert h[1][0][0].shape == (2,)
             assert isinstance(h[1][0][1], jax.numpy.ndarray)
-            assert h[1][0][1].shape == (4,)
+            assert h[1][0][1].shape == (2,)
             assert isinstance(h[1][1], tuple)
             assert len(h[1][1]) == 2
             assert isinstance(h[1][1][0], jax.numpy.ndarray)
-            assert h[1][1][0].shape == (4,)
+            assert h[1][1][0].shape == (2,)
             assert isinstance(h[1][1][1], jax.numpy.ndarray)
-            assert h[1][1][1].shape == (4,)
+            assert h[1][1][1].shape == (2,)
 
     def test_hessian_expval_probs_multiple_param_array(
         self, dev_name, diff_method, gradient_kwargs, shots, interface
@@ -661,7 +661,7 @@ class TestReturnWithShotVectors:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.hessian(circuit)(params)
 
@@ -677,7 +677,7 @@ class TestReturnWithShotVectors:
             assert h[0].shape == (2, 2)
 
             assert isinstance(h[1], jax.numpy.ndarray)
-            assert h[1].shape == (4, 2, 2)
+            assert h[1].shape == (2, 2, 2)
 
     def test_hessian_probs_var_multiple_params(
         self, dev_name, diff_method, gradient_kwargs, shots, interface
@@ -693,7 +693,7 @@ class TestReturnWithShotVectors:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1)
 
@@ -725,15 +725,15 @@ class TestReturnWithShotVectors:
             assert isinstance(h[1][0], tuple)
             assert len(h[1][0]) == 2
             assert isinstance(h[1][0][0], jax.numpy.ndarray)
-            assert h[1][0][0].shape == (4,)
+            assert h[1][0][0].shape == (2,)
             assert isinstance(h[1][0][1], jax.numpy.ndarray)
-            assert h[1][0][1].shape == (4,)
+            assert h[1][0][1].shape == (2,)
             assert isinstance(h[1][1], tuple)
             assert len(h[1][1]) == 2
             assert isinstance(h[1][1][0], jax.numpy.ndarray)
-            assert h[1][1][0].shape == (4,)
+            assert h[1][1][0].shape == (2,)
             assert isinstance(h[1][1][1], jax.numpy.ndarray)
-            assert h[1][1][1].shape == (4,)
+            assert h[1][1][1].shape == (2,)
 
     def test_hessian_var_probs_multiple_param_array(
         self, dev_name, diff_method, gradient_kwargs, shots, interface
@@ -748,7 +748,7 @@ class TestReturnWithShotVectors:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         hess = jax.hessian(circuit)(params)
 
@@ -764,7 +764,7 @@ class TestReturnWithShotVectors:
             assert h[0].shape == (2, 2)
 
             assert isinstance(h[1], jax.numpy.ndarray)
-            assert h[1].shape == (4, 2, 2)
+            assert h[1].shape == (2, 2, 2)
 
 
 shots = [(1, 20, 100), (1, (20, 1), 100), (1, (5, 4), 100)]
diff --git a/tests/returntypes/tf/test_tensorflow_qnode_new.py b/tests/returntypes/tf/test_tensorflow_qnode_new.py
index f8d1932fc..4910ba69e 100644
--- a/tests/returntypes/tf/test_tensorflow_qnode_new.py
+++ b/tests/returntypes/tf/test_tensorflow_qnode_new.py
@@ -2252,7 +2252,7 @@ class TestReturn:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
@@ -2268,10 +2268,10 @@ class TestReturn:
         assert len(hess) == 2
 
         assert isinstance(hess[0], tf.Tensor)
-        assert hess[0].shape == (10,)
+        assert hess[0].shape == (6,)
 
         assert isinstance(hess[1], tf.Tensor)
-        assert hess[1].shape == (10,)
+        assert hess[1].shape == (6,)
 
     def test_hessian_probs_expval_multiple_param_array(
         self, dev_name, diff_method, mode, interface
@@ -2295,7 +2295,7 @@ class TestReturn:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
@@ -2307,7 +2307,7 @@ class TestReturn:
         hess = tape1.jacobian(grad, params)
 
         assert isinstance(hess, tf.Tensor)
-        assert hess.shape == (5, 2, 2)
+        assert hess.shape == (3, 2, 2)
 
     def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, mode, interface):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
@@ -2326,7 +2326,7 @@ class TestReturn:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
@@ -2342,10 +2342,10 @@ class TestReturn:
         assert len(hess) == 2
 
         assert isinstance(hess[0], tf.Tensor)
-        assert hess[0].shape == (10,)
+        assert hess[0].shape == (6,)
 
         assert isinstance(hess[1], tf.Tensor)
-        assert hess[1].shape == (10,)
+        assert hess[1].shape == (6,)
 
     def test_hessian_probs_var_multiple_param_array(self, dev_name, diff_method, mode, interface):
         """The hessian of multiple measurements with a multiple param array return a single array."""
@@ -2363,7 +2363,7 @@ class TestReturn:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
@@ -2375,4 +2375,4 @@ class TestReturn:
         hess = tape1.jacobian(grad, params)
 
         assert isinstance(hess, tf.Tensor)
-        assert hess.shape == (5, 2, 2)
+        assert hess.shape == (3, 2, 2)
diff --git a/tests/returntypes/tf/test_tensorflow_qnode_shot_vector.py b/tests/returntypes/tf/test_tensorflow_qnode_shot_vector.py
index f5a0ce104..05a4ff767 100644
--- a/tests/returntypes/tf/test_tensorflow_qnode_shot_vector.py
+++ b/tests/returntypes/tf/test_tensorflow_qnode_shot_vector.py
@@ -406,7 +406,7 @@ class TestReturnShotVectorHessian:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
@@ -438,7 +438,7 @@ class TestReturnShotVectorHessian:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         with tf.GradientTape() as tape1:
             with tf.GradientTape(persistent=True) as tape2:
diff --git a/tests/returntypes/torch/test_torch_qnode_new.py b/tests/returntypes/torch/test_torch_qnode_new.py
index d9622f424..8481fb198 100644
--- a/tests/returntypes/torch/test_torch_qnode_new.py
+++ b/tests/returntypes/torch/test_torch_qnode_new.py
@@ -2264,7 +2264,7 @@ class TestReturn:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def circuit_stack(x, y):
             return torch.hstack(circuit(x, y))
@@ -2279,16 +2279,16 @@ class TestReturn:
         assert isinstance(hess[0], tuple)
         assert len(hess[0]) == 2
         assert isinstance(hess[0][0], torch.Tensor)
-        assert tuple(hess[0][0].shape) == (5,)
+        assert tuple(hess[0][0].shape) == (3,)
         assert isinstance(hess[0][1], torch.Tensor)
-        assert tuple(hess[0][1].shape) == (5,)
+        assert tuple(hess[0][1].shape) == (3,)
 
         assert isinstance(hess[1], tuple)
         assert len(hess[1]) == 2
         assert isinstance(hess[1][0], torch.Tensor)
-        assert tuple(hess[1][0].shape) == (5,)
+        assert tuple(hess[1][0].shape) == (3,)
         assert isinstance(hess[1][1], torch.Tensor)
-        assert tuple(hess[1][1].shape) == (5,)
+        assert tuple(hess[1][1].shape) == (3,)
 
     def test_hessian_expval_probs_multiple_param_array(self, dev_name, diff_method, mode, shots):
         """The hessian of multiple measurements with a multiple param array return a single array."""
@@ -2308,7 +2308,7 @@ class TestReturn:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def circuit_stack(x):
             return torch.hstack(circuit(x))
@@ -2318,7 +2318,7 @@ class TestReturn:
         hess = jacobian(jac_fn, par)
 
         assert isinstance(hess, torch.Tensor)
-        assert tuple(hess.shape) == (5, 2, 2)
+        assert tuple(hess.shape) == (3, 2, 2)
 
     def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, mode, shots):
         """The hessian of multiple measurements with multiple params return a tuple of arrays."""
@@ -2339,7 +2339,7 @@ class TestReturn:
             qml.RX(x, wires=[0])
             qml.RY(y, wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def circuit_stack(x, y):
             return torch.hstack(circuit(x, y))
@@ -2354,16 +2354,16 @@ class TestReturn:
         assert isinstance(hess[0], tuple)
         assert len(hess[0]) == 2
         assert isinstance(hess[0][0], torch.Tensor)
-        assert tuple(hess[0][0].shape) == (5,)
+        assert tuple(hess[0][0].shape) == (3,)
         assert isinstance(hess[0][1], torch.Tensor)
-        assert tuple(hess[0][1].shape) == (5,)
+        assert tuple(hess[0][1].shape) == (3,)
 
         assert isinstance(hess[1], tuple)
         assert len(hess[1]) == 2
         assert isinstance(hess[1][0], torch.Tensor)
-        assert tuple(hess[1][0].shape) == (5,)
+        assert tuple(hess[1][0].shape) == (3,)
         assert isinstance(hess[1][1], torch.Tensor)
-        assert tuple(hess[1][1].shape) == (5,)
+        assert tuple(hess[1][1].shape) == (3,)
 
     def test_hessian_var_probs_multiple_param_array(self, dev_name, diff_method, mode, shots):
         """The hessian of multiple measurements with a multiple param array return a single array."""
@@ -2383,7 +2383,7 @@ class TestReturn:
             qml.RX(x[0], wires=[0])
             qml.RY(x[1], wires=[1])
             qml.CNOT(wires=[0, 1])
-            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0, 1])
+            return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0])
 
         def circuit_stack(x):
             return torch.hstack(circuit(x))
@@ -2393,4 +2393,4 @@ class TestReturn:
         hess = jacobian(jac_fn, par)
 
         assert isinstance(hess, torch.Tensor)
-        assert tuple(hess.shape) == (5, 2, 2)
+        assert tuple(hess.shape) == (3, 2, 2)
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index 1fe0d7a68..efe48510d 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -1074,7 +1074,7 @@ class TestExpand:
             # expands into RY on wire b
             qml.expval(qml.PauliZ("a") @ qml.Hadamard("b"))
             # expands into QubitUnitary on wire 0
-            qml.var(qml.Hermitian(np.array([[1, 2], [2, 4]]), wires=[0]))
+            qml.var(qml.Hermitian(np.array([[1, 2], [2, 4]]), wires=[1]))
 
         new_tape = tape.expand(expand_measurements=True)