Use wire order specified on instantiation (#705)

* use map wires instead of map_to_standard_wires

* oops

* Update .github/CHANGELOG.md

* black

* Trigger CIs

* add test

* black on tests

* Update .github/CHANGELOG.md

---------

Co-authored-by: AmintorDusko <[email protected]>
Co-authored-by: Ali Asadi <[email protected]>

Author: 3 people

.github/CHANGELOG.md

@@ -108,6 +108,9 @@
 
 ### Bug fixes
 
+* Lightning Qubit once again respects the wire order specified on device instantiation.
+  [(#705)](https://github.com/PennyLaneAI/pennylane-lightning/pull/705)
+
 * `dynamic_one_shot` was refactored to use `SampleMP` measurements as a way to return the mid-circuit measurement samples. `LightningQubit`'s `simulate` is modified accordingly.
   [(#694)](https://github.com/PennyLaneAI/pennylane-lightning/pull/694)
 

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -87,7 +87,7 @@ def simulate(circuit: QuantumScript, state: LightningStateVector, mcmc: dict = N
     return LightningMeasurements(final_state, **mcmc).measure_final_state(circuit)
 
 
-def jacobian(circuit: QuantumTape, state: LightningStateVector, batch_obs=False):
+def jacobian(circuit: QuantumTape, state: LightningStateVector, batch_obs=False, wire_map=None):
     """Compute the Jacobian for a single quantum script.
 
     Args:
@@ -96,17 +96,21 @@ def jacobian(circuit: QuantumTape, state: LightningStateVector, batch_obs=False)
         batch_obs (bool): Determine whether we process observables in parallel when
             computing the jacobian. This value is only relevant when the lightning
             qubit is built with OpenMP. Default is False.
+        wire_map (Optional[dict]): a map from wire labels to simulation indices
 
     Returns:
         TensorLike: The Jacobian of the quantum script
     """
-    circuit = circuit.map_to_standard_wires()
+    if wire_map is not None:
+        [circuit], _ = qml.map_wires(circuit, wire_map)
     state.reset_state()
     final_state = state.get_final_state(circuit)
     return LightningAdjointJacobian(final_state, batch_obs=batch_obs).calculate_jacobian(circuit)
 
 
-def simulate_and_jacobian(circuit: QuantumTape, state: LightningStateVector, batch_obs=False):
+def simulate_and_jacobian(
+    circuit: QuantumTape, state: LightningStateVector, batch_obs=False, wire_map=None
+):
     """Simulate a single quantum script and compute its Jacobian.
 
     Args:
@@ -115,20 +119,26 @@ def simulate_and_jacobian(circuit: QuantumTape, state: LightningStateVector, bat
         batch_obs (bool): Determine whether we process observables in parallel when
             computing the jacobian. This value is only relevant when the lightning
             qubit is built with OpenMP. Default is False.
+        wire_map (Optional[dict]): a map from wire labels to simulation indices
 
     Returns:
         Tuple[TensorLike]: The results of the simulation and the calculated Jacobian
 
     Note that this function can return measurements for non-commuting observables simultaneously.
     """
-    circuit = circuit.map_to_standard_wires()
+    if wire_map is not None:
+        [circuit], _ = qml.map_wires(circuit, wire_map)
     res = simulate(circuit, state)
     jac = LightningAdjointJacobian(state, batch_obs=batch_obs).calculate_jacobian(circuit)
     return res, jac
 
 
 def vjp(
-    circuit: QuantumTape, cotangents: Tuple[Number], state: LightningStateVector, batch_obs=False
+    circuit: QuantumTape,
+    cotangents: Tuple[Number],
+    state: LightningStateVector,
+    batch_obs=False,
+    wire_map=None,
 ):
     """Compute the Vector-Jacobian Product (VJP) for a single quantum script.
     Args:
@@ -141,10 +151,13 @@ def vjp(
         batch_obs (bool): Determine whether we process observables in parallel when
             computing the VJP. This value is only relevant when the lightning
             qubit is built with OpenMP.
+        wire_map (Optional[dict]): a map from wire labels to simulation indices
+
     Returns:
         TensorLike: The VJP of the quantum script
     """
-    circuit = circuit.map_to_standard_wires()
+    if wire_map is not None:
+        [circuit], _ = qml.map_wires(circuit, wire_map)
     state.reset_state()
     final_state = state.get_final_state(circuit)
     return LightningAdjointJacobian(final_state, batch_obs=batch_obs).calculate_vjp(
@@ -153,7 +166,11 @@ def vjp(
 
 
 def simulate_and_vjp(
-    circuit: QuantumTape, cotangents: Tuple[Number], state: LightningStateVector, batch_obs=False
+    circuit: QuantumTape,
+    cotangents: Tuple[Number],
+    state: LightningStateVector,
+    batch_obs=False,
+    wire_map=None,
 ):
     """Simulate a single quantum script and compute its Vector-Jacobian Product (VJP).
     Args:
@@ -166,11 +183,14 @@ def simulate_and_vjp(
         batch_obs (bool): Determine whether we process observables in parallel when
             computing the jacobian. This value is only relevant when the lightning
             qubit is built with OpenMP.
+        wire_map (Optional[dict]): a map from wire labels to simulation indices
+
     Returns:
         Tuple[TensorLike]: The results of the simulation and the calculated VJP
     Note that this function can return measurements for non-commuting observables simultaneously.
     """
-    circuit = circuit.map_to_standard_wires()
+    if wire_map is not None:
+        [circuit], _ = qml.map_wires(circuit, wire_map)
     res = simulate(circuit, state)
     _vjp = LightningAdjointJacobian(state, batch_obs=batch_obs).calculate_vjp(circuit, cotangents)
     return res, _vjp
@@ -449,6 +469,11 @@ def __init__(  # pylint: disable=too-many-arguments
 
         super().__init__(wires=wires, shots=shots)
 
+        if isinstance(wires, int):
+            self._wire_map = None  # should just use wires as is
+        else:
+            self._wire_map = {w: i for i, w in enumerate(self.wires)}
+
         self._statevector = LightningStateVector(num_wires=len(self.wires), dtype=c_dtype)
 
         # TODO: Investigate usefulness of creating numpy random generator
@@ -568,7 +593,8 @@ def execute(
         }
         results = []
         for circuit in circuits:
-            circuit = circuit.map_to_standard_wires()
+            if self._wire_map is not None:
+                [circuit], _ = qml.map_wires(circuit, self._wire_map)
             results.append(simulate(circuit, self._statevector, mcmc=mcmc))
 
         return tuple(results)
@@ -613,8 +639,10 @@ def compute_derivatives(
             Tuple: The jacobian for each trainable parameter
         """
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
+
         return tuple(
-            jacobian(circuit, self._statevector, batch_obs=batch_obs) for circuit in circuits
+            jacobian(circuit, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map)
+            for circuit in circuits
         )
 
     def execute_and_compute_derivatives(
@@ -633,7 +661,10 @@ def execute_and_compute_derivatives(
         """
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
         results = tuple(
-            simulate_and_jacobian(c, self._statevector, batch_obs=batch_obs) for c in circuits
+            simulate_and_jacobian(
+                c, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map
+            )
+            for c in circuits
         )
         return tuple(zip(*results))
 
@@ -686,7 +717,7 @@ def compute_vjp(
         """
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
         return tuple(
-            vjp(circuit, cots, self._statevector, batch_obs=batch_obs)
+            vjp(circuit, cots, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map)
             for circuit, cots in zip(circuits, cotangents)
         )
 
@@ -708,7 +739,9 @@ def execute_and_compute_vjp(
         """
         batch_obs = execution_config.device_options.get("batch_obs", self._batch_obs)
         results = tuple(
-            simulate_and_vjp(circuit, cots, self._statevector, batch_obs=batch_obs)
+            simulate_and_vjp(
+                circuit, cots, self._statevector, batch_obs=batch_obs, wire_map=self._wire_map
+            )
             for circuit, cots in zip(circuits, cotangents)
         )
         return tuple(zip(*results))

tests/new_api/test_device.py

@@ -456,6 +456,29 @@ def test_custom_wires(self, phi, theta, wires):
         assert np.allclose(result[0], np.cos(phi))
         assert np.allclose(result[1], np.cos(phi) * np.cos(theta))
 
+    @pytest.mark.parametrize(
+        "wires, wire_order", [(3, (0, 1, 2)), (("a", "b", "c"), ("a", "b", "c"))]
+    )
+    def test_probs_different_wire_orders(self, wires, wire_order):
+        """Test that measuring probabilities works with custom wires."""
+
+        dev = LightningDevice(wires=wires)
+
+        op = qml.Hadamard(wire_order[1])
+
+        tape = QuantumScript([op], [qml.probs(wires=(wire_order[0], wire_order[1]))])
+
+        res = dev.execute(tape)
+        assert qml.math.allclose(res, np.array([0.5, 0.5, 0.0, 0.0]))
+
+        tape2 = QuantumScript([op], [qml.probs(wires=(wire_order[1], wire_order[2]))])
+        res2 = dev.execute(tape2)
+        assert qml.math.allclose(res2, np.array([0.5, 0.0, 0.5, 0.0]))
+
+        tape3 = QuantumScript([op], [qml.probs(wires=(wire_order[1], wire_order[0]))])
+        res3 = dev.execute(tape3)
+        assert qml.math.allclose(res3, np.array([0.5, 0.0, 0.5, 0.0]))
+
 
 @pytest.mark.parametrize("batch_obs", [True, False])
 class TestDerivatives: