Repository: PennyLaneAI/pennylane

Issue:
[BUG] Confusing behavior with the `wires` argument of `qml.probs`
### Expected behavior

Consider a case where where a device's wires are ordered as `[2,0,1]` and we'd like to get the probabilities associated with these wires in the order `[0,1,2]`. That is, if a `qml.PauliX(wires=[0])` operation is applied, I would expect `qml.probs(wires=[0,1,2])` to return the probabilities `[0,0,0,0,1,0,0,0]` because `0` is the ordered as the most significant bit in the wire ordering.

In the two following code snippets, I expect the `[0,0,0,0,1,0,0,0]` probability vector to be returned when `qml.PauliX` is applied to wire `[0]`.

1. Explicitly specifying the operation's wire:

```python
dev = qml.device("default.qubit", wires=[2,0,1])
@qml.qnode(dev)
def circ():
    qml.PauliX(wires=[0])
    return qml.probs(wires=[0,1,2])

circ()
```

2. Specifying the operation's wire based on its relative position in `dev.wires`:

```python
dev = qml.device("default.qubit", wires=[2,0,1])
@qml.qnode(dev)
def circ():
    qml.PauliX(wires=dev.wires[1])
    return qml.probs(wires=[0,1,2])

circ()
```

I believe that this expectation should hold because I specify that the `qml.PauliX` operates upon the `[0]` wire and that `qml.probs(wires=[0,1,2])` should return the probabilities ordered with the `[0]` wire being the most significant bit.

### Actual behavior

## Case 1.

```python
dev = qml.device("default.qubit", wires=[2,0,1])
@qml.qnode(dev)
def circ():
    qml.PauliX(wires=[0])
    return qml.probs(wires=[0,1,2])

circ()
```

Returns:

```
tensor([0., 1., 0., 0., 0., 0., 0., 0.], requires_grad=True)
```

## Case 2:

```python
dev = qml.device("default.qubit", wires=[2,0,1])
@qml.qnode(dev)
def circ():
    qml.PauliX(wires=dev.wires[1])
    return qml.probs(wires=[0,1,2])

circ()
```

Returns:

```
tensor([0., 1., 0., 0., 0., 0., 0., 0.], requires_grad=True)
```

## Problem

In both cases, it appears that the `qml.PauliX` operation is applied to wire `[2]` instead of wire `[0]` like I intended.

### Additional information

_No response_

### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.28.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/brian/opt/anaconda3/lib/python3.9/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning, qNetVO

Platform info:           macOS-10.16-x86_64-i386-64bit
Python version:          3.9.13
Numpy version:           1.21.5
Scipy version:           1.9.1
Installed devices:
- lightning.qubit (PennyLane-Lightning-0.28.1)
- default.gaussian (PennyLane-0.28.0)
- default.mixed (PennyLane-0.28.0)
- default.qubit (PennyLane-0.28.0)
- default.qubit.autograd (PennyLane-0.28.0)
- default.qubit.jax (PennyLane-0.28.0)
- default.qubit.tf (PennyLane-0.28.0)
- default.qubit.torch (PennyLane-0.28.0)
- default.qutrit (PennyLane-0.28.0)
- null.qubit (PennyLane-0.28.0)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Discussion:
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.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3753?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#3753](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3753?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (aeb1dfa) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/08511eaa8ae13953d2e5464217864bf2ab028333?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (08511ea) will **not change** coverage.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #3753   +/-   ##
=======================================
  Coverage   99.82%   99.82%           
=======================================
  Files         330      330           
  Lines       28909    28908    -1     
=======================================
- Hits        28858    28857    -1     
  Misses         51       51           
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3753?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/measurements/probs.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3753?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9wcm9icy5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/qubit/state\_preparation.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/3753?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9xdWJpdC9zdGF0ZV9wcmVwYXJhdGlvbi5weQ==) | `100.00% <100.00%> (ø)` | |

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)

timmysilv: fwiw, I just found out that `pennylane_cirq.SimulatorDevice._expand_state` [expanded and re-arranged nicely](https://github.com/PennyLaneAI/pennylane-cirq/blob/4c09d20cb4c386d6912aa1d8c58253747accd0a7/pennylane_cirq/simulator_device.py#L117-L132) (but specifically for state-prep ops). `tranpose` is more concise, but still thought this was cool to find (I am deleting it in another PR though : D )
timmysilv: Tbh I wouldn't unless it's too burdensome. I could do something in this PR, but I hesitated because the qubit/qutrit devices should disappear and be replaced with the new Device API, and they will use the `qml.probs` implementation. That leaves only `probs.process_state` and `QSV.state_vector`, and ultimately the logic is just `np.transpose` and that's already bundled all nice for us

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 9e8f3e39f..560da79ca 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -718,6 +718,10 @@
 * Redirect `qml.math.ndim` to `jnp.ndim` when using it on a jax tensor.
   [(#3730)](https://github.com/PennyLaneAI/pennylane/pull/3730)
 
+* Implementations of `marginal_prob` (and subsequently, `qml.probs`) now return
+  probabilities with the expected wire order.
+  [(#3753)](https://github.com/PennyLaneAI/pennylane/pull/3753)
+
 <h3>Contributors</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index b6f6146fd..a3609f306 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -141,67 +141,6 @@ class QubitDevice(Device):
         """Data type preserving multiply operation"""
         return qmlmul(constant, array, dtype=array.dtype)
 
-    def _permute_wires(self, observable):
-        r"""Given an observable which acts on multiple wires, permute the wires to
-          be consistent with the device wire order.
-
-          Suppose we are given an observable :math:`\hat{O} = \Identity \otimes \Identity \otimes \hat{Z}`.
-          This observable can be represented in many ways:
-
-        .. code-block:: python
-
-              O_1 = qml.Identity(wires=0) @ qml.Identity(wires=1) @ qml.PauliZ(wires=2)
-              O_2 = qml.PauliZ(wires=2) @ qml.Identity(wires=0) @ qml.Identity(wires=1)
-
-          Notice that while the explicit tensor product matrix representation of :code:`O_1` and :code:`O_2` is
-          different, the underlying operator is identical due to the wire labelling (assuming the labels in
-          ascending order are {0,1,2}). If we wish to compute the expectation value of such an observable, we must
-          ensure it is identical in both cases. To facilitate this, we permute the wires in our state vector such
-          that they are consistent with this swapping of order in the tensor observable.
-
-        .. code-block:: python
-
-              >>> print(O_1.wires)
-              <Wires = [0, 1, 2]>
-              >>> print(O_2.wires)
-              <Wires = [2, 0, 1]>
-
-          We might naively think that we must permute our state vector to match the wire order of our tensor observable.
-          We must be careful and realize that the wire order of the terms in the tensor observable DOES NOT match the
-          permutation of the terms themselves. As an example we directly compare :code:`O_1` and :code:`O_2`:
-
-          The first term in :code:`O_1` (:code:`qml.Identity(wires=0)`) became the second term in :code:`O_2`.
-          By similar comparison we see that each term in the tensor product was shifted one position forward
-          (i.e 0 --> 1, 1 --> 2, 2 --> 0). The wires in our permuted quantum state should follow their respective
-          terms in the tensor product observable.
-
-          Thus, the correct wire ordering should be :code:`permuted_wires = <Wires = [1, 2, 0]>`. But if we had
-          taken the naive approach we would have permuted our state according to
-          :code:`permuted_wires = <Wires = [2, 0, 1]>` which is NOT correct.
-
-          This function uses the observable wires and the global device wire ordering in order to determine the
-          permutation of the wires in the observable required such that if our quantum state vector is
-          permuted accordingly then the amplitudes of the state will match the matrix representation of the observable.
-
-          Args:
-              observable (Observable): the observable whose wires are to be permuted.
-
-          Returns:
-              permuted_wires (Wires): permuted wires object
-        """
-        ordered_obs_wire_lst = self.order_wires(
-            observable.wires
-        ).tolist()  # order according to device wire order
-
-        mapped_wires = self.map_wires(observable.wires)
-        if isinstance(mapped_wires, Wires):
-            # by default this should be a Wires obj, but it is overwritten to list object in default.qubit
-            mapped_wires = mapped_wires.tolist()
-
-        permutation = np.argsort(mapped_wires)  # extract permutation via argsort
-
-        return Wires([ordered_obs_wire_lst[index] for index in permutation])
-
     observables = {
         "PauliX",
         "PauliY",
@@ -1670,35 +1609,29 @@ class QubitDevice(Device):
         device_wires = self.map_wires(wires)
         inactive_device_wires = self.map_wires(inactive_wires)
 
-        # reshape the probability so that each axis corresponds to a wire
-        shape = [2] * self.num_wires
-        if batch_size is not None:
-            shape.insert(0, batch_size)
-        # prob now is reshaped to have self.num_wires+1 axes in the case of broadcasting
-        prob = self._reshape(prob, shape)
-
-        # sum over all inactive wires
         # hotfix to catch when default.qubit uses this method
         # since then device_wires is a list
         if isinstance(inactive_device_wires, Wires):
             inactive_device_wires = inactive_device_wires.labels
 
+        # reshape the probability so that each axis corresponds to a wire
+        shape = [2] * self.num_wires
+        desired_axes = np.argsort(np.argsort(device_wires))
+        flat_shape = (-1,)
         if batch_size is not None:
+            # prob now is reshaped to have self.num_wires+1 axes in the case of broadcasting
+            shape.insert(0, batch_size)
             inactive_device_wires = [idx + 1 for idx in inactive_device_wires]
-        flat_shape = (-1,) if batch_size is None else (batch_size, -1)
-        prob = self._reshape(self._reduce_sum(prob, inactive_device_wires), flat_shape)
-
-        # The wires provided might not be in consecutive order (i.e., wires might be [2, 0]).
-        # If this is the case, we must permute the marginalized probability so that
-        # it corresponds to the orders of the wires passed.
-        num_wires = len(device_wires)
-        basis_states = self.generate_basis_states(num_wires)
-        basis_states = basis_states[:, np.argsort(np.argsort(device_wires))]
+            desired_axes = np.insert(desired_axes + 1, 0, 0)
+            flat_shape = (batch_size, -1)
 
-        powers_of_two = 2 ** np.arange(len(device_wires))[::-1]
-        perm = basis_states @ powers_of_two
-        # The permutation happens on the last axis both with and without broadcasting
-        return self._gather(prob, perm, axis=1 if batch_size is not None else 0)
+        prob = self._reshape(prob, shape)
+        # sum over all inactive wires
+        prob = self._reduce_sum(prob, inactive_device_wires)
+        # rearrange wires to desired order
+        prob = self._transpose(prob, desired_axes)
+        # flatten and return probabilities
+        return self._reshape(prob, flat_shape)
 
     def expval(self, observable, shot_range=None, bin_size=None):
         if observable.name == "Projector":
@@ -1718,11 +1651,7 @@ class QubitDevice(Device):
                     f"Cannot compute analytic expectations of {observable.name}."
                 ) from e
 
-            # the probability vector must be permuted to account for the permuted
-            # wire order of the observable
-            permuted_wires = self._permute_wires(observable)
-
-            prob = self.probability(wires=permuted_wires)
+            prob = self.probability(wires=observable.wires)
             # In case of broadcasting, `prob` has two axes and this is a matrix-vector product
             return self._dot(prob, eigvals)
 
@@ -1753,10 +1682,7 @@ class QubitDevice(Device):
                     f"Cannot compute analytic variance of {observable.name}."
                 ) from e
 
-            # the probability vector must be permuted to account for the permuted wire order of the observable
-            permuted_wires = self._permute_wires(observable)
-
-            prob = self.probability(wires=permuted_wires)
+            prob = self.probability(wires=observable.wires)
             # In case of broadcasting, `prob` has two axes and these are a matrix-vector products
             return self._dot(prob, (eigvals**2)) - self._dot(prob, eigvals) ** 2
 
diff --git a/pennylane/_qutrit_device.py b/pennylane/_qutrit_device.py
index ea03c4052..3787e5a3f 100644
--- a/pennylane/_qutrit_device.py
+++ b/pennylane/_qutrit_device.py
@@ -313,18 +313,9 @@ class QutritDevice(QubitDevice):  # pylint: disable=too-many-public-methods
         else:
             wires = inactive_device_wires
 
-        prob = self._flatten(self._reduce_sum(prob, wires))
-
-        # The wires provided might not be in consecutive order (i.e., wires might be [2, 0]).
-        # If this is the case, we must permute the marginalized probability so that
-        # it corresponds to the orders of the wires passed.
-        num_wires = len(device_wires)
-        basis_states = self.generate_basis_states(num_wires)
-        basis_states = basis_states[:, np.argsort(np.argsort(device_wires))]
-
-        powers_of_three = 3 ** np.arange(len(device_wires))[::-1]
-        perm = basis_states @ powers_of_three
-        return self._gather(prob, perm)
+        prob = self._reduce_sum(prob, wires)
+        prob = self._transpose(prob, np.argsort(np.argsort(device_wires)))
+        return self._flatten(prob)
 
     def sample(self, observable, shot_range=None, bin_size=None, counts=False):
         def _samples_to_counts(samples, no_observable_provided):
diff --git a/pennylane/devices/default_mixed.py b/pennylane/devices/default_mixed.py
index ce0e70daa..aba03ed3f 100644
--- a/pennylane/devices/default_mixed.py
+++ b/pennylane/devices/default_mixed.py
@@ -132,6 +132,7 @@ class DefaultMixed(QubitDevice):
 
     _reshape = staticmethod(qnp.reshape)
     _flatten = staticmethod(qnp.flatten)
+    _transpose = staticmethod(qnp.transpose)
     # Allow for the `axis` keyword argument for integration with broadcasting-enabling
     # code in QubitDevice. However, it is not used as DefaultMixed does not support broadcasting
     # pylint: disable=unnecessary-lambda
diff --git a/pennylane/measurements/probs.py b/pennylane/measurements/probs.py
index 8d644667f..95d411e21 100644
--- a/pennylane/measurements/probs.py
+++ b/pennylane/measurements/probs.py
@@ -17,6 +17,7 @@ This module contains the qml.probs measurement.
 from typing import Sequence, Tuple
 
 import pennylane as qml
+from pennylane import numpy as np
 from pennylane.wires import Wires
 
 from .measurements import MeasurementShapeError, Probability, SampleMeasurement, StateMeasurement
@@ -299,23 +300,19 @@ class ProbabilityMP(SampleMeasurement, StateMeasurement):
         # reshape the probability so that each axis corresponds to a wire
         num_device_wires = len(wire_order)
         shape = [2] * num_device_wires
+        desired_axes = np.argsort(np.argsort(mapped_wires))
+        flat_shape = (-1,)
         if batch_size is not None:
+            # prob now is reshaped to have self.num_wires+1 axes in the case of broadcasting
             shape.insert(0, batch_size)
             inactive_wires = [idx + 1 for idx in inactive_wires]
-        # prob now is reshaped to have self.num_wires+1 axes in the case of broadcasting
-        prob = qml.math.reshape(prob, shape)
+            desired_axes = np.insert(desired_axes + 1, 0, 0)
+            flat_shape = (batch_size, -1)
 
+        prob = qml.math.reshape(prob, shape)
         # sum over all inactive wires
-        flat_shape = (-1,) if batch_size is None else (batch_size, -1)
-        prob = qml.math.reshape(qml.math.sum(prob, axis=tuple(inactive_wires)), flat_shape)
-
-        # The wires provided might not be in consecutive order (i.e., wires might be [2, 0]).
-        # If this is the case, we must permute the marginalized probability so that
-        # it corresponds to the orders of the wires passed.
-        num_wires = len(mapped_wires)
-        basis_states = qml.QubitDevice.generate_basis_states(num_wires)
-        basis_states = basis_states[:, qml.math.argsort(qml.math.argsort(mapped_wires))]
-
-        powers_of_two = 2 ** qml.math.arange(num_wires)[::-1]
-        perm = basis_states @ powers_of_two
-        return prob[:, perm] if batch_size is not None else qml.math.gather(prob, perm)
+        prob = qml.math.sum(prob, axis=tuple(inactive_wires))
+        # rearrange wires if necessary
+        prob = qml.math.transpose(prob, desired_axes)
+        # flatten and return probabilities
+        return qml.math.reshape(prob, flat_shape)
diff --git a/pennylane/ops/qubit/state_preparation.py b/pennylane/ops/qubit/state_preparation.py
index 22b8905fd..41c6feae8 100644
--- a/pennylane/ops/qubit/state_preparation.py
+++ b/pennylane/ops/qubit/state_preparation.py
@@ -194,7 +194,8 @@ class QubitStateVector(StatePrep):
         if wire_order is None or Wires(wire_order) == self.wires:
             return op_vector
 
-        if not Wires(wire_order).contains_wires(self.wires):
+        wire_order = Wires(wire_order)
+        if not wire_order.contains_wires(self.wires):
             raise WireError("Custom wire_order must contain all QubitStateVector wires")
 
         num_total_wires = len(wire_order)
@@ -203,15 +204,10 @@ class QubitStateVector(StatePrep):
         ket[indices] = op_vector
 
         # unless wire_order is [*self.wires, *rest_of_wire_order], need to rearrange
-        if (op_wires := list(self.wires)) != wire_order[:num_op_wires]:
-            new_wires = set(wire_order) - set(op_wires)
-            new_order = op_wires + list(new_wires)
-            for i, wire in enumerate(wire_order):
-                # after each loop iteration, the i'th axis will represent the correct wire
-                i_wire_pos = new_order.index(wire)
-                if i_wire_pos != i:
-                    ket = np.swapaxes(ket, i, i_wire_pos)
-                    new_order[i], new_order[i_wire_pos] = new_order[i_wire_pos], new_order[i]
+        if self.wires != wire_order[:num_op_wires]:
+            current_order = self.wires + list(Wires.unique_wires([wire_order, self.wires]))
+            desired_order = [current_order.index(w) for w in wire_order]
+            ket = ket.transpose(desired_order)
 
         return math.convert_like(ket, op_vector)
 


Test patch:
diff --git a/tests/devices/test_default_qutrit.py b/tests/devices/test_default_qutrit.py
index 98361b10a..8de8eaa4d 100644
--- a/tests/devices/test_default_qutrit.py
+++ b/tests/devices/test_default_qutrit.py
@@ -970,6 +970,20 @@ class TestProbabilityIntegration:
 
         assert dev.analytic_probability() is None
 
+    def test_marginal_prob_wire_order(self):
+        """Tests that marginal_prob rearranges wires as expected."""
+        dev = qml.device("default.qutrit", wires=3)
+
+        @qml.qnode(dev)
+        def circuit():
+            qml.QutritUnitary(U_x_02, wires=[1])  # second wire ("1") set to 2-state
+            return qml.probs(wires=[2, 0, 1])  # third wire ("1") should be in 2-state here
+
+        probs = qml.math.reshape(circuit(), (3, 3, 3))
+        assert probs[0, 0, 2] == 1
+        probs[0, 0, 2] = 0
+        assert qml.math.allequal(probs, 0)
+
 
 class TestWiresIntegration:
     """Test that the device integrates with PennyLane's wire management."""
diff --git a/tests/measurements/test_probs.py b/tests/measurements/test_probs.py
index 316971901..9a38a5ce0 100644
--- a/tests/measurements/test_probs.py
+++ b/tests/measurements/test_probs.py
@@ -58,8 +58,8 @@ def custom_measurement_process(device, spy):
 np.random.seed(42)
 
 
-@pytest.fixture
-def init_state():
+@pytest.fixture(name="init_state")
+def fixture_init_state():
     """Fixture that creates an initial state"""
 
     def _init_state(n):
@@ -74,6 +74,8 @@ def init_state():
 class TestProbs:
     """Tests for the probs function"""
 
+    # pylint:disable=too-many-public-methods
+
     def test_queue(self):
         """Test that the right measurement class is queued."""
         dev = qml.device("default.qubit", wires=2)
@@ -194,12 +196,10 @@ class TestProbs:
         spy_probs = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
 
-        spy = mocker.spy(qml.QubitDevice, "states_to_binary")
-
         @qml.qnode(dev)
         def circuit():
             qml.QubitStateVector(state, wires=list(range(4)))
-            return qml.probs(wires=[1, 0, 3])  # <--- more than 2 wires: states_to_binary used
+            return qml.probs(wires=[1, 0, 3])
 
         res = circuit()
 
@@ -207,10 +207,50 @@ class TestProbs:
         expected = np.einsum("ijkl->jil", expected).flatten()
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
-        spy.assert_called_once()
-
         custom_measurement_process(dev, spy_probs)
 
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    @pytest.mark.parametrize(
+        "subset_wires,expected",
+        [
+            ([0, 1], [0.25, 0.25, 0.25, 0.25]),
+            ([1, 2], [0.5, 0, 0.5, 0]),
+            ([0, 2], [0.5, 0, 0.5, 0]),
+            ([2, 0], [0.5, 0.5, 0, 0]),
+            ([2, 1], [0.5, 0.5, 0, 0]),
+            ([1, 2, 0], [0.25, 0.25, 0, 0, 0.25, 0.25, 0, 0]),
+        ],
+    )
+    def test_process_state(self, interface, subset_wires, expected):
+        """Tests that process_state functions as expected with all interfaces."""
+        state = qml.math.array([1 / 2, 0] * 4, like=interface)
+        wires = qml.wires.Wires(range(3))
+        subset_probs = qml.probs(wires=subset_wires).process_state(state, wires)
+        assert subset_probs.shape == (len(expected),)
+        assert qml.math.allclose(subset_probs, expected)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    @pytest.mark.parametrize(
+        "subset_wires,expected",
+        [
+            ([1, 2], [[0.5, 0, 0.5, 0], [0, 0.5, 0, 0.5]]),
+            ([2, 0], [[0.5, 0.5, 0, 0], [0, 0, 0.5, 0.5]]),
+            (
+                [1, 2, 0],
+                [[0.25, 0.25, 0, 0, 0.25, 0.25, 0, 0], [0, 0, 0.25, 0.25, 0, 0, 0.25, 0.25]],
+            ),
+        ],
+    )
+    def test_process_state_batched(self, interface, subset_wires, expected):
+        """Tests that process_state functions as expected with all interfaces with batching."""
+        states = qml.math.array([[1 / 2, 0] * 4, [0, 1 / 2] * 4], like=interface)
+        wires = qml.wires.Wires(range(3))
+        subset_probs = qml.probs(wires=subset_wires).process_state(states, wires)
+        assert subset_probs.shape == qml.math.shape(expected)
+        assert qml.math.allclose(subset_probs, expected)
+
     def test_integration(self, tol, mocker):
         """Test the probability is correct for a known state preparation."""
         dev = qml.device("default.qubit", wires=2)
@@ -411,6 +451,7 @@ class TestProbs:
     @pytest.mark.parametrize("wire", [0, 1, 2, 3])
     def test_prob_generalize_initial_state(self, hermitian, wire, init_state, tol, mocker):
         """Test that the correct probability is returned."""
+        # pylint:disable=too-many-arguments
         dev = qml.device("default.qubit", wires=4)
         spy = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
@@ -454,6 +495,7 @@ class TestProbs:
     @pytest.mark.parametrize("wire", [0, 1, 2, 3])
     def test_operation_prob(self, operation, wire, init_state, tol, mocker):
         "Test the rotated probability with different wires and rotating operations."
+        # pylint:disable=too-many-arguments
         dev = qml.device("default.qubit", wires=4)
         spy = mocker.spy(qml.QubitDevice, "probability")
         state = init_state(4)
diff --git a/tests/ops/op_math/test_controlled_ops.py b/tests/ops/op_math/test_controlled_ops.py
index d6d8a241e..2b058a487 100644
--- a/tests/ops/op_math/test_controlled_ops.py
+++ b/tests/ops/op_math/test_controlled_ops.py
@@ -228,8 +228,8 @@ class TestControlledQubitUnitary:
         U = unitary_group.rvs(2 ** len(target_wires), random_state=1967)
 
         # Pick random starting state for the control and target qubits
-        control_state_weights = np.random.normal(size=(2 ** (len(control_wires) + 1) - 2))
-        target_state_weights = np.random.normal(size=(2 ** (len(target_wires) + 1) - 2))
+        control_state_weights = np.random.normal(size=2 ** (len(control_wires) + 1) - 2)
+        target_state_weights = np.random.normal(size=2 ** (len(target_wires) + 1) - 2)
 
         @qml.qnode(dev)
         def circuit_mixed_polarity():
diff --git a/tests/ops/qubit/test_state_prep.py b/tests/ops/qubit/test_state_prep.py
index 4e0c97271..1e5909fbd 100644
--- a/tests/ops/qubit/test_state_prep.py
+++ b/tests/ops/qubit/test_state_prep.py
@@ -108,6 +108,7 @@ class TestStateVector:
             (3, [1, 2, 0], (1, 0, 0)),
             (3, [1, 2, "a"], (1, 0, 0)),
             (3, [2, 1, 0], (0, 1, 0)),
+            (4, [3, 2, 0, 1], (0, 0, 0, 1)),
         ],
     )
     def test_QubitStateVector_state_vector(self, num_wires, wire_order, one_position):
diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py
index 9cd7695f1..3d6bed8d9 100644
--- a/tests/test_qubit_device.py
+++ b/tests/test_qubit_device.py
@@ -1043,6 +1043,16 @@ class TestMarginalProb:
             np.array([0.3970422, 0.28090525, 0.11116351, 0.21088904]),
             [2, 0],
         ),
+        (
+            np.array([0.05, 0.07, 0.11, 0.13, 0.17, 0.19, 0.23, 0.05]),
+            np.array([0.05, 0.11, 0.17, 0.23, 0.07, 0.13, 0.19, 0.05]),
+            [2, 0, 1],
+        ),
+        (
+            np.arange(1, 17) / 136,
+            np.array([3, 11, 19, 27, 7, 15, 23, 31]) / 136,
+            [2, 0, 1],
+        ),
     ]
 
     @pytest.mark.parametrize("probs, marginals, wires", marginal_test_data)
@@ -1086,6 +1096,22 @@ class TestMarginalProb:
             [2, 0],
             3,
         ),
+        (
+            np.array(
+                [
+                    [0.05, 0.07, 0.11, 0.13, 0.17, 0.19, 0.23, 0.05],
+                    np.arange(1, 9) / 36,
+                ]
+            ),
+            np.array(
+                [
+                    [0.05, 0.11, 0.17, 0.23, 0.07, 0.13, 0.19, 0.05],
+                    np.array([1, 3, 5, 7, 2, 4, 6, 8]) / 36,
+                ],
+            ),
+            [2, 0, 1],
+            3,
+        ),
     ]
 
     @pytest.mark.parametrize("probs, marginals, wires, num_wires", broadcasted_marginal_test_data)
diff --git a/tests/tests_passing_pylint.txt b/tests/tests_passing_pylint.txt
index 963cd01fc..e317ed80f 100644
--- a/tests/tests_passing_pylint.txt
+++ b/tests/tests_passing_pylint.txt
@@ -2,6 +2,7 @@ tests/data/**
 tests/devices/qubit/test_apply_operation.py
 tests/devices/test_execution_config.py
 tests/drawer/test_tape_mpl.py
+tests/measurements/test_probs.py
 tests/qchem/test_givens_rotations.py
 tests/templates/test_subroutines/test_basis_rotation.py
 tests/ops/op_math/test_controlled_ops.py