Repository: PennyLaneAI/pennylane (python)

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.

Pull request description:
Return expected wire order with qml.probs(), use same logic in QSV
**Context:**
Bug #3741 explains the background, but essentially qml.probs() does not order wires correctly. The desired wire order was given to an array, essentially saying "rearrange your axes to look like this", but that's not exactly how array indexing works. The arbitrary rearrangement of axes should instead be done with the `transpose(array, axes)`.

**Description of the Change:**
- Fix the wire rearranging logic in `qml.probs()` by using `np.transpose()`
- Update `QubitStateVector.state_vector()` to use this same logic.
  - I can do this in another PR but it feels related enough here. Context: I wrote that code a couple weeks ago and I thought to myself "hmmm why can't I rearrange my current axes to some arbitrary target axes? roll and swap just don't cut it! I'll just write my own basic axis-shuffler." Because the root of this bug is the same axis challenge, I figured I should update QSV as well

**Benefits:**
- `qml.probs(wires=new_order)` behaves as expected
- the code in `QubitStateVector.state_vector()` is prettier 😄 

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
Fixes #3741 

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