Updated LQ to work with runtime kwargs

pennylane_lightning/lightning_qubit/lightning_qubit.py

@@ -570,7 +570,9 @@ def preprocess(self, execution_config: ExecutionConfig = DefaultExecutionConfig)
         program.add_transform(validate_measurements, name=self.name)
         program.add_transform(validate_observables, accepted_observables, name=self.name)
         program.add_transform(validate_device_wires, self.wires, name=self.name)
-        program.add_transform(mid_circuit_measurements, device=self)
+        program.add_transform(
+            mid_circuit_measurements, device=self, mcm_config=exec_config.mcm_config
+        )
         program.add_transform(
             decompose,
             stopping_condition=stopping_condition,

requirements-dev.txt

@@ -1,5 +1,5 @@
 pip~=22.0
-git+https://github.com/PennyLaneAI/pennylane.git@master
+git+https://github.com/PennyLaneAI/pennylane.git@postselect-choice
 ninja
 flaky
 pybind11

tests/new_api/test_device.py

@@ -350,6 +350,31 @@ def test_preprocess_state_prep_middle_op_decomposition(self, op, decomp_depth):
         )
         assert qml.equal(new_tape, expected_tape)
 
+    @pytest.mark.parametrize("postselect_shots", [True, False])
+    def test_postselect_shots(self, postselect_shots):
+        """Test that results are as expected when a user requests to scale shots with postselection."""
+
+        shots = 10
+        device = LightningDevice(wires=3, shots=shots)
+
+        @qml.qnode(device, postselect_shots=postselect_shots)
+        def func(x):
+            qml.RX(x, 0)
+            qml.measure(0, postselect=1)
+            qml.CNOT([0, 1])
+            qml.sample(qml.PauliZ(1))
+
+        # Choosing small rotation angle ensures that the number of valid shots is less than the
+        # total number of shots. This will allow us to avoid stochastic failures where the number
+        # of valid samples happens to be the same as the number of shots.
+        x = np.pi / 5
+        res = func(x)
+
+        if postselect_shots:
+            assert len(res) < shots
+        else:
+            assert len(res) == shots
+
     @pytest.mark.usefixtures("use_legacy_and_new_opmath")
     @pytest.mark.parametrize("theta, phi", list(zip(THETA, PHI)))
     @pytest.mark.parametrize(