From 66840b00b9e4e5cb9bd789b105a9e33249d19207 Mon Sep 17 00:00:00 2001
From: Mudit Pandey <mudit.pandey@xanadu.ai>
Date: Mon, 20 May 2024 21:52:33 -0400
Subject: [PATCH] Updated LQ to work with runtime kwargs

---
 .../lightning_qubit/lightning_qubit.py        |  4 ++-
 requirements-dev.txt                          |  2 +-
 tests/new_api/test_device.py                  | 25 +++++++++++++++++++
 3 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/pennylane_lightning/lightning_qubit/lightning_qubit.py b/pennylane_lightning/lightning_qubit/lightning_qubit.py
index fff9d8acaf..e2b0cb6fb4 100644
--- a/pennylane_lightning/lightning_qubit/lightning_qubit.py
+++ b/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,
diff --git a/requirements-dev.txt b/requirements-dev.txt
index b3fb6d76a5..90740a3cbf 100644
--- a/requirements-dev.txt
+++ b/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
diff --git a/tests/new_api/test_device.py b/tests/new_api/test_device.py
index a1a5912160..620e4228a0 100644
--- a/tests/new_api/test_device.py
+++ b/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(