Use wire order specified on instantiation

.github/CHANGELOG.md

@@ -105,6 +105,8 @@
 
 ### Bug fixes
 
+* Lightning Qubit once again respects the wire order specified on device instantiation.
+
 * `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/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,19 @@ 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 +117,22 @@ 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 +145,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 +160,7 @@ 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 +173,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 +459,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 +583,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 +629,9 @@ 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 +650,7 @@ 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 +703,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 +725,7 @@ 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))