From 0f01d4d018e1ed291acbd80f879923a141d0a104 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Thu, 20 Jul 2023 09:59:14 -0400
Subject: [PATCH 01/19] Add vjp/jvp capabilities

---
 .../devices/experimental/default_qubit_2.py   | 265 ++++++++++++++++++
 pennylane/devices/qubit/adjoint_jacobian.py   |  33 +--
 pennylane/devices/qubit/simulate.py           |  62 ++--
 3 files changed, 317 insertions(+), 43 deletions(-)

diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
index 63a5cf72a8d..ab856033eae 100644
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ b/pennylane/devices/experimental/default_qubit_2.py
@@ -16,6 +16,7 @@
 """
 
 from functools import partial
+from numbers import Number
 from typing import Union, Callable, Tuple, Optional, Sequence
 import concurrent.futures
 import os
@@ -283,6 +284,270 @@ def compute_derivatives(
             f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
         )
 
+    def execute_and_compute_derivatives(
+        self,
+        circuits: QuantumTape_or_Batch,
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+
+        if self.tracker.active:
+            for c in circuits:
+                self.tracker.update(resources=c.specs["resources"])
+            self.tracker.update(batches=1, executions=len(circuits))
+            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.record()
+
+        if execution_config.gradient_method != "adjoint":
+            raise NotImplementedError(
+                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
+            )
+
+        max_workers = self._get_max_workers(execution_config)
+        if max_workers is None:
+            results = tuple(
+                simulate(c, rng=self._rng, debugger=self._debugger, return_final_state=True)
+                for c in circuits
+            )
+            jacs = tuple(adjoint_jacobian(c, state=r[1]) for c, r in zip(circuits, results))
+            results = tuple(r[0] for r in results)
+        else:
+            self._validate_multiprocessing_circuits(circuits)
+
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
+
+            def wrapper(c, rng):
+                res, final_state, _ = simulate(c, rng=rng, debugger=None, return_final_state=True)
+                jac = adjoint_jacobian(c, state=final_state)
+                return res, jac
+
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                results = tuple(executor.map(wrapper, vanilla_circuits, seeds))
+                results, jacs = tuple(zip(*results))
+
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
+
+        return results[0], jacs[0] if is_single_circuit else results, jacs
+
+    def supports_jvp(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[QuantumTape] = None,
+    ) -> bool:
+        """Whether or not this device defines a custom jacobian vector product.
+
+        ``DefaultQubit2`` supports backpropagation derivatives with analytic results, as well as
+        adjoint differentiation.
+
+        Args:
+            execution_config (ExecutionConfig): The configuration of the desired derivative calculation
+            circuit (QuantumTape): An optional circuit to check derivatives support for.
+
+        Returns:
+            bool: Whether or not a derivative can be calculated provided the given information
+        """
+        return self.supports_derivatives(execution_config, circuit)
+
+    def compute_jvp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        tangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            tangents = [tangents]
+
+        if self.tracker.active:
+            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.record()
+
+        if execution_config.gradient_method != "adjoint":
+            raise NotImplementedError(
+                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
+            )
+
+        max_workers = self._get_max_workers(execution_config)
+        if max_workers is None:
+            res = tuple(adjoint_jvp(circuit, tans) for circuit, tans in zip(circuits, tangents))
+        else:
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                res = tuple(executor.map(adjoint_jvp, vanilla_circuits, tangents))
+
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
+
+        return res[0] if is_single_circuit else res
+
+    def execute_and_compute_jvp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        tangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            tangents = [tangents]
+
+        if self.tracker.active:
+            for c in circuits:
+                self.tracker.update(resources=c.specs["resources"])
+            self.tracker.update(batches=1, executions=len(circuits))
+            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.record()
+
+        if execution_config.gradient_method != "adjoint":
+            raise NotImplementedError(
+                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
+            )
+
+        max_workers = self._get_max_workers(execution_config)
+        if max_workers is None:
+            results = tuple(
+                simulate(c, rng=self._rng, debugger=self._debugger, return_final_state=True)
+                for c in circuits
+            )
+            jvps = tuple(
+                adjoint_jvp(c, t, state=r[1]) for c, t, r in zip(circuits, tangents, results)
+            )
+            results = tuple(r[0] for r in results)
+        else:
+            self._validate_multiprocessing_circuits(circuits)
+
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
+
+            def wrapper(c, t, rng):
+                res, final_state, _ = simulate(c, rng=rng, debugger=None, return_final_state=True)
+                jvp = adjoint_jvp(c, t, state=final_state)
+                return res, jvp
+
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                results = tuple(executor.map(wrapper, vanilla_circuits, tangents, seeds))
+                results, jvps = tuple(zip(*results))
+
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
+
+        return results[0], jvps[0] if is_single_circuit else results, jvps
+
+    def supports_vjp(
+        self,
+        execution_config: Optional[ExecutionConfig] = None,
+        circuit: Optional[QuantumTape] = None,
+    ) -> bool:
+        """Whether or not this device defines a custom vector jacobian product.
+
+        ``DefaultQubit2`` supports backpropagation derivatives with analytic results, as well as
+        adjoint differentiation.
+
+        Args:
+            execution_config (ExecutionConfig): A description of the hyperparameters for the desired computation.
+            circuit (None, QuantumTape): A specific circuit to check differentation for.
+
+        Returns:
+            bool: Whether or not a derivative can be calculated provided the given information
+        """
+        return self.supports_derivatives(execution_config, circuit)
+
+    def compute_vjp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            cotangents = [cotangents]
+
+        if self.tracker.active:
+            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.record()
+
+        if execution_config.gradient_method != "adjoint":
+            raise NotImplementedError(
+                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
+            )
+
+        max_workers = self._get_max_workers(execution_config)
+        if max_workers is None:
+            res = tuple(adjoint_vjp(circuit, cots) for circuit, cots in zip(circuits, cotangents))
+        else:
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                res = tuple(executor.map(adjoint_vjp, vanilla_circuits, cotangents))
+
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
+
+        return res[0] if is_single_circuit else res
+
+    def execute_and_compute_jvp(
+        self,
+        circuits: QuantumTape_or_Batch,
+        cotangents: Tuple[Number],
+        execution_config: ExecutionConfig = DefaultExecutionConfig,
+    ):
+        is_single_circuit = False
+        if isinstance(circuits, QuantumScript):
+            is_single_circuit = True
+            circuits = [circuits]
+            cotangents = [cotangents]
+
+        if self.tracker.active:
+            for c in circuits:
+                self.tracker.update(resources=c.specs["resources"])
+            self.tracker.update(batches=1, executions=len(circuits))
+            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.record()
+
+        if execution_config.gradient_method != "adjoint":
+            raise NotImplementedError(
+                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
+            )
+
+        max_workers = self._get_max_workers(execution_config)
+        if max_workers is None:
+            results = tuple(
+                simulate(c, rng=self._rng, debugger=self._debugger, return_final_state=True)
+                for c in circuits
+            )
+            jvps = tuple(
+                adjoint_vjp(c, t, state=r[1]) for c, t, r in zip(circuits, cotangents, results)
+            )
+            results = tuple(r[0] for r in results)
+        else:
+            self._validate_multiprocessing_circuits(circuits)
+
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
+
+            def wrapper(c, t, rng):
+                res, final_state, _ = simulate(c, rng=rng, debugger=None, return_final_state=True)
+                vjp = adjoint_vjp(c, t, state=final_state)
+                return res, vjp
+
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                results = tuple(executor.map(wrapper, vanilla_circuits, cotangents, seeds))
+                results, vjps = tuple(zip(*results))
+
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
+
+        return results[0], vjps[0] if is_single_circuit else results, vjps
+
     # pylint: disable=missing-function-docstring
     def _get_max_workers(self, execution_config=None):
         max_workers = None
diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index bb296bb44e3..4c7de1dc2c5 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -23,6 +23,7 @@
 
 from .apply_operation import apply_operation
 from .initialize_state import create_initial_state
+from .simulate import _final_state
 
 # pylint: disable=protected-access, too-many-branches
 
@@ -34,21 +35,7 @@ def _dot_product_real(bra, ket, num_wires):
     return qml.math.real(qml.math.sum(qml.math.conj(bra) * ket, axis=sum_axes))
 
 
-def _get_output_ket(tape):
-    """Helper function to get the output state of a tape"""
-
-    # Initialization of state
-    prep_operation = tape[0] if isinstance(tape[0], qml.operation.StatePrep) else None
-    ket = create_initial_state(
-        wires=tape.wires, prep_operation=prep_operation
-    )  #  ket(0) if prep_operation is None, else
-    for op in tape.operations[bool(prep_operation) :]:
-        ket = apply_operation(op, ket)
-
-    return ket
-
-
-def adjoint_jacobian(tape: QuantumTape):
+def adjoint_jacobian(tape: QuantumTape, state=None):
     """Implements the adjoint method outlined in
     `Jones and Gacon <https://arxiv.org/abs/2009.02823>`__ to differentiate an input tape.
 
@@ -67,6 +54,8 @@ def adjoint_jacobian(tape: QuantumTape):
 
     Args:
         tape (.QuantumTape): circuit that the function takes the gradient of
+        state (TensorLike): the final state of the circuit; if not provided,
+            the final state will be computed by executing the tape
 
     Returns:
         array or tuple[array]: the derivative of the tape with respect to trainable parameters.
@@ -77,7 +66,7 @@ def adjoint_jacobian(tape: QuantumTape):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = _get_output_ket(tape)
+    ket = state if state is not None else _final_state(tape)[0]
 
     n_obs = len(tape.observables)
     bras = np.empty([n_obs] + [2] * len(tape.wires), dtype=np.complex128)
@@ -119,7 +108,7 @@ def adjoint_jacobian(tape: QuantumTape):
     return tuple(tuple(np.array(j_) for j_ in j) for j in jac)
 
 
-def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number]):
+def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number], state=None):
     """The jacobian vector product used in forward mode calculation of derivatives.
 
     Implements the adjoint method outlined in
@@ -141,6 +130,8 @@ def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number]):
     Args:
         tape (.QuantumTape): circuit that the function takes the gradient of
         tangents (Tuple[Number]): gradient vector for input parameters.
+        state (TensorLike): the final state of the circuit; if not provided,
+            the final state will be computed by executing the tape
 
     Returns:
         Tuple[Number]: gradient vector for output parameters
@@ -150,7 +141,7 @@ def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number]):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = _get_output_ket(tape)
+    ket = state if state is not None else _final_state(tape)[0]
 
     n_obs = len(tape.observables)
     bras = np.empty([n_obs] + [2] * len(tape.wires), dtype=np.complex128)
@@ -191,7 +182,7 @@ def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number]):
     return tuple(np.array(t) for t in tangents_out)
 
 
-def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number]):
+def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
     """The vector jacobian product used in reverse-mode differentiation.
 
     Implements the adjoint method outlined in
@@ -213,6 +204,8 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number]):
     Args:
         tape (.QuantumTape): circuit that the function takes the gradient of
         cotangents (Tuple[Number]): gradient vector for output parameters
+        state (TensorLike): the final state of the circuit; if not provided,
+            the final state will be computed by executing the tape
 
     Returns:
         Tuple[Number]: gradient vector for input parameters
@@ -222,7 +215,7 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number]):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = _get_output_ket(tape)
+    ket = state if state is not None else _final_state(tape)[0]
 
     obs = qml.dot(cotangents, tape.observables)
     bra = apply_operation(obs, ket)
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index bcbc0b060f1..fff2b58cf7f 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -24,7 +24,30 @@
 from .sampling import measure_with_samples
 
 
-def simulate(circuit: qml.tape.QuantumScript, rng=None, debugger=None) -> Result:
+def _final_state(circuit, debugger=None):
+    """
+    Get the final state from executing the ops in the circuit
+    """
+    prep = circuit[0] if isinstance(circuit[0], qml.operation.StatePrep) else None
+    state = create_initial_state(circuit.wires, prep)
+
+    # initial state is batched only if the state preparation (if it exists) is batched
+    is_state_batched = False
+    if prep and prep.batch_size is not None:
+        is_state_batched = True
+
+    for op in circuit.operations[bool(prep) :]:
+        state = apply_operation(op, state, is_state_batched=is_state_batched, debugger=debugger)
+
+        # new state is batched if i) the old state is batched, or ii) the new op adds a batch dim
+        is_state_batched = is_state_batched or op.batch_size is not None
+
+    return state, is_state_batched
+
+
+def simulate(
+    circuit: qml.tape.QuantumScript, rng=None, debugger=None, return_final_state=False
+) -> Result:
     """Simulate a single quantum script.
 
     This is an internal function that will be called by the successor to ``default.qubit``.
@@ -35,6 +58,7 @@ def simulate(circuit: qml.tape.QuantumScript, rng=None, debugger=None) -> Result
             seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
             If no value is provided, a default RNG will be used.
         debugger (._Debugger): The debugger to use
+        return_final_state (bool): Whether to return the final state in addition to the results
 
     Returns:
         tuple(TensorLike): The results of the simulation
@@ -55,39 +79,32 @@ def simulate(circuit: qml.tape.QuantumScript, rng=None, debugger=None) -> Result
         wire_map = {w: i for i, w in enumerate(circuit.wires)}
         circuit = qml.map_wires(circuit, wire_map)
 
-    state = create_initial_state(circuit.wires, circuit._prep[0] if circuit._prep else None)
-
-    # initial state is batched only if the state preparation (if it exists) is batched
-    is_state_batched = False
-    if circuit._prep and circuit._prep[0].batch_size is not None:
-        is_state_batched = True
-
-    for op in circuit._ops:
-        state = apply_operation(op, state, is_state_batched=is_state_batched, debugger=debugger)
-
-        # new state is batched if i) the old state is batched, or ii) the new op adds a batch dim
-        is_state_batched = is_state_batched or op.batch_size is not None
+    state, is_state_batched = _final_state(circuit, debugger=debugger)
 
     if not circuit.shots:
         # analytic case
 
         if len(circuit.measurements) == 1:
-            return measure(circuit.measurements[0], state, is_state_batched=is_state_batched)
+            results = measure(circuit.measurements[0], state, is_state_batched=is_state_batched)
 
-        return tuple(
-            measure(mp, state, is_state_batched=is_state_batched) for mp in circuit.measurements
-        )
+        else:
+            results = tuple(
+                measure(mp, state, is_state_batched=is_state_batched) for mp in circuit.measurements
+            )
+
+        return (results, state, is_state_batched) if return_final_state else results
 
     # finite-shot case
 
     if len(circuit.measurements) == 1:
-        return measure_with_samples(
+        results = measure_with_samples(
             circuit.measurements[0],
             state,
             shots=circuit.shots,
             is_state_batched=is_state_batched,
             rng=rng,
         )
+        return (results, state, is_state_batched) if return_final_state else results
 
     rng = default_rng(rng)
     results = tuple(
@@ -97,9 +114,8 @@ def simulate(circuit: qml.tape.QuantumScript, rng=None, debugger=None) -> Result
         for mp in circuit.measurements
     )
 
-    # no shot vector
-    if not circuit.shots.has_partitioned_shots:
-        return results
+    if circuit.shots.has_partitioned_shots:
+        # shot vector case: move the shot vector axis before the measurement axis
+        results = tuple(zip(*results))
 
-    # shot vector case: move the shot vector axis before the measurement axis
-    return tuple(zip(*results))
+    return (results, state, is_state_batched) if return_final_state else results

From 7a975bef408fb01c9995c0c8454aa971e40d75d6 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Thu, 20 Jul 2023 10:16:36 -0400
Subject: [PATCH 02/19] pylint

---
 pennylane/devices/experimental/default_qubit_2.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
index ab856033eae..e4d2a16d941 100644
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ b/pennylane/devices/experimental/default_qubit_2.py
@@ -32,7 +32,7 @@
 from .execution_config import ExecutionConfig, DefaultExecutionConfig
 from ..qubit.simulate import simulate
 from ..qubit.preprocess import preprocess, validate_and_expand_adjoint
-from ..qubit.adjoint_jacobian import adjoint_jacobian
+from ..qubit.adjoint_jacobian import adjoint_jacobian, adjoint_vjp, adjoint_jvp
 
 Result_or_ResultBatch = Union[Result, ResultBatch]
 QuantumTapeBatch = Sequence[QuantumTape]
@@ -494,7 +494,7 @@ def compute_vjp(
 
         return res[0] if is_single_circuit else res
 
-    def execute_and_compute_jvp(
+    def execute_and_compute_vjp(
         self,
         circuits: QuantumTape_or_Batch,
         cotangents: Tuple[Number],
@@ -524,7 +524,7 @@ def execute_and_compute_jvp(
                 simulate(c, rng=self._rng, debugger=self._debugger, return_final_state=True)
                 for c in circuits
             )
-            jvps = tuple(
+            vjps = tuple(
                 adjoint_vjp(c, t, state=r[1]) for c, t, r in zip(circuits, cotangents, results)
             )
             results = tuple(r[0] for r in results)

From e317831c3debe524c63091e703e634295a9acb8e Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Thu, 20 Jul 2023 10:59:26 -0400
Subject: [PATCH 03/19] some tests for simulate

---
 tests/devices/qubit/test_simulate.py | 97 +++++++++++++++++++++++++---
 1 file changed, 87 insertions(+), 10 deletions(-)

diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index 1d8e07cf6e4..1b7c74a926a 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -66,13 +66,24 @@ def test_basis_state(self):
 class TestBasicCircuit:
     """Tests a basic circuit with one rx gate and two simple expectation values."""
 
-    def test_basic_circuit_numpy(self):
+    @pytest.mark.parametrize("return_final_state", [False, True])
+    def test_basic_circuit_numpy(self, return_final_state):
         """Test execution with a basic circuit."""
         phi = np.array(0.397)
         qs = qml.tape.QuantumScript(
             [qml.RX(phi, wires=0)], [qml.expval(qml.PauliY(0)), qml.expval(qml.PauliZ(0))]
         )
-        result = simulate(qs)
+        result = simulate(qs, return_final_state=return_final_state)
+
+        if return_final_state:
+            assert isinstance(result, tuple)
+            assert len(result) == 3
+
+            # the second element is the final state, and the third is whether it is batched
+            assert np.allclose(result[1], np.array([np.cos(phi / 2), -1j * np.sin(phi / 2)]))
+            assert not result[2]
+
+            result = result[0]
 
         assert isinstance(result, tuple)
         assert len(result) == 2
@@ -173,7 +184,8 @@ def test_tf_results_and_backprop(self):
 class TestBroadcasting:
     """Test that simulate works with broadcasted parameters"""
 
-    def test_broadcasted_prep_state(self):
+    @pytest.mark.parametrize("return_final_state", [False, True])
+    def test_broadcasted_prep_state(self, return_final_state):
         """Test that simulate works for state measurements
         when the state prep has broadcasted parameters"""
         x = np.array(1.2)
@@ -183,14 +195,33 @@ def test_broadcasted_prep_state(self):
         prep = [qml.QubitStateVector(np.eye(4), wires=[0, 1])]
 
         qs = qml.tape.QuantumScript(ops, measurements, prep)
-        res = simulate(qs)
+        res = simulate(qs, return_final_state=return_final_state)
+
+        if return_final_state:
+            assert isinstance(res, tuple)
+            assert len(res) == 3
+
+            # the second element is the final state, and the third is whether it is batched
+            expected = np.array(
+                [
+                    [np.cos(x / 2), 0, 0, np.sin(x / 2)],
+                    [0, np.cos(x / 2), np.sin(x / 2), 0],
+                    [-np.sin(x / 2), 0, 0, np.cos(x / 2)],
+                    [0, -np.sin(x / 2), np.cos(x / 2), 0],
+                ]
+            ).reshape((4, 2, 2))
+            assert np.allclose(res[1], expected)
+            assert res[2]
+
+            res = res[0]
 
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert np.allclose(res[0], np.array([np.cos(x), np.cos(x), -np.cos(x), -np.cos(x)]))
         assert np.allclose(res[1], np.array([np.cos(x), -np.cos(x), -np.cos(x), np.cos(x)]))
 
-    def test_broadcasted_op_state(self):
+    @pytest.mark.parametrize("return_final_state", [False, True])
+    def test_broadcasted_op_state(self, return_final_state):
         """Test that simulate works for state measurements
         when an operation has broadcasted parameters"""
         x = np.array([0.8, 1.0, 1.2, 1.4])
@@ -199,14 +230,28 @@ def test_broadcasted_op_state(self):
         measurements = [qml.expval(qml.PauliZ(i)) for i in range(2)]
 
         qs = qml.tape.QuantumScript(ops, measurements)
-        res = simulate(qs)
+        res = simulate(qs, return_final_state=return_final_state)
+
+        if return_final_state:
+            assert isinstance(res, tuple)
+            assert len(res) == 3
+
+            # the second element is the final state, and the third is whether it is batched
+            expected = np.zeros((4, 2, 2))
+            expected[:, 0, 1] = np.cos(x / 2)
+            expected[:, 1, 0] = np.sin(x / 2)
+            assert np.allclose(res[1], expected)
+            assert res[2]
+
+            res = res[0]
 
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert np.allclose(res[0], np.cos(x))
         assert np.allclose(res[1], -np.cos(x))
 
-    def test_broadcasted_prep_sample(self):
+    @pytest.mark.parametrize("return_final_state", [False, True])
+    def test_broadcasted_prep_sample(self, return_final_state):
         """Test that simulate works for sample measurements
         when the state prep has broadcasted parameters"""
         x = np.array(1.2)
@@ -216,7 +261,25 @@ def test_broadcasted_prep_sample(self):
         prep = [qml.QubitStateVector(np.eye(4), wires=[0, 1])]
 
         qs = qml.tape.QuantumScript(ops, measurements, prep, shots=qml.measurements.Shots(10000))
-        res = simulate(qs, rng=123)
+        res = simulate(qs, rng=123, return_final_state=return_final_state)
+
+        if return_final_state:
+            assert isinstance(res, tuple)
+            assert len(res) == 3
+
+            # the second element is the final state, and the third is whether it is batched
+            expected = np.array(
+                [
+                    [np.cos(x / 2), 0, 0, np.sin(x / 2)],
+                    [0, np.cos(x / 2), np.sin(x / 2), 0],
+                    [-np.sin(x / 2), 0, 0, np.cos(x / 2)],
+                    [0, -np.sin(x / 2), np.cos(x / 2), 0],
+                ]
+            ).reshape((4, 2, 2))
+            assert np.allclose(res[1], expected)
+            assert res[2]
+
+            res = res[0]
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -227,7 +290,8 @@ def test_broadcasted_prep_sample(self):
             res[1], np.array([np.cos(x), -np.cos(x), -np.cos(x), np.cos(x)]), atol=0.05
         )
 
-    def test_broadcasted_op_sample(self):
+    @pytest.mark.parametrize("return_final_state", [False, True])
+    def test_broadcasted_op_sample(self, return_final_state):
         """Test that simulate works for sample measurements
         when an operation has broadcasted parameters"""
         x = np.array([0.8, 1.0, 1.2, 1.4])
@@ -236,7 +300,20 @@ def test_broadcasted_op_sample(self):
         measurements = [qml.expval(qml.PauliZ(i)) for i in range(2)]
 
         qs = qml.tape.QuantumScript(ops, measurements, shots=qml.measurements.Shots(10000))
-        res = simulate(qs, rng=123)
+        res = simulate(qs, rng=123, return_final_state=return_final_state)
+
+        if return_final_state:
+            assert isinstance(res, tuple)
+            assert len(res) == 3
+
+            # the second element is the final state, and the third is whether it is batched
+            expected = np.zeros((4, 2, 2))
+            expected[:, 0, 1] = np.cos(x / 2)
+            expected[:, 1, 0] = np.sin(x / 2)
+            assert np.allclose(res[1], expected)
+            assert res[2]
+
+            res = res[0]
 
         assert isinstance(res, tuple)
         assert len(res) == 2

From 7f068834de85b2b894dcab4c223ca1ea923147be Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Thu, 20 Jul 2023 16:14:59 -0400
Subject: [PATCH 04/19] Fix syntax error

---
 pennylane/devices/experimental/default_qubit_2.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
index e4d2a16d941..6a0e26a5926 100644
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ b/pennylane/devices/experimental/default_qubit_2.py
@@ -332,7 +332,7 @@ def wrapper(c, rng):
             # reset _rng to mimic serial behavior
             self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
 
-        return results[0], jacs[0] if is_single_circuit else results, jacs
+        return (results[0], jacs[0]) if is_single_circuit else (results, jacs)
 
     def supports_jvp(
         self,
@@ -439,7 +439,7 @@ def wrapper(c, t, rng):
             # reset _rng to mimic serial behavior
             self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
 
-        return results[0], jvps[0] if is_single_circuit else results, jvps
+        return (results[0], jvps[0]) if is_single_circuit else (results, jvps)
 
     def supports_vjp(
         self,
@@ -546,7 +546,7 @@ def wrapper(c, t, rng):
             # reset _rng to mimic serial behavior
             self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
 
-        return results[0], vjps[0] if is_single_circuit else results, vjps
+        return (results[0], vjps[0]) if is_single_circuit else (results, vjps)
 
     # pylint: disable=missing-function-docstring
     def _get_max_workers(self, execution_config=None):

From d2d507a9e27ac782e9036e1f14b4f0afbbdf3b5b Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Fri, 21 Jul 2023 14:33:25 -0400
Subject: [PATCH 05/19] Split simulate

---
 .../devices/experimental/default_qubit_2.py   |  70 ++++----
 pennylane/devices/qubit/__init__.py           |   2 +-
 pennylane/devices/qubit/adjoint_jacobian.py   |   8 +-
 pennylane/devices/qubit/simulate.py           |  88 +++++-----
 tests/devices/qubit/test_simulate.py          | 153 ++++++++----------
 5 files changed, 149 insertions(+), 172 deletions(-)

diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
index 6a0e26a5926..995d7203cd7 100644
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ b/pennylane/devices/experimental/default_qubit_2.py
@@ -30,7 +30,7 @@
 
 from . import Device
 from .execution_config import ExecutionConfig, DefaultExecutionConfig
-from ..qubit.simulate import simulate
+from ..qubit.simulate import simulate, get_final_state, measure_final_state
 from ..qubit.preprocess import preprocess, validate_and_expand_adjoint
 from ..qubit.adjoint_jacobian import adjoint_jacobian, adjoint_vjp, adjoint_jvp
 
@@ -306,28 +306,26 @@ def execute_and_compute_derivatives(
                 f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
             )
 
+        def wrapper(c, rng=None, debugger=None):
+            state = get_final_state(c, debugger=debugger)
+            jac = adjoint_jacobian(c, state=state)
+            res = measure_final_state(c, state, False, rng=rng)
+            return res, jac
+
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
-            results = tuple(
-                simulate(c, rng=self._rng, debugger=self._debugger, return_final_state=True)
-                for c in circuits
-            )
-            jacs = tuple(adjoint_jacobian(c, state=r[1]) for c, r in zip(circuits, results))
-            results = tuple(r[0] for r in results)
+            results = tuple(wrapper(c, rng=self._rng, debugger=self._debugger) for c in circuits)
+            results, jacs = tuple(zip(*results))
         else:
             self._validate_multiprocessing_circuits(circuits)
 
             vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
             seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
 
-            def wrapper(c, rng):
-                res, final_state, _ = simulate(c, rng=rng, debugger=None, return_final_state=True)
-                jac = adjoint_jacobian(c, state=final_state)
-                return res, jac
-
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
                 results = tuple(executor.map(wrapper, vanilla_circuits, seeds))
-                results, jacs = tuple(zip(*results))
+
+            results, jacs = tuple(zip(*results))
 
             # reset _rng to mimic serial behavior
             self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
@@ -411,30 +409,29 @@ def execute_and_compute_jvp(
                 f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
             )
 
+        def wrapper(c, t, rng=None, debugger=None):
+            state = get_final_state(c, debugger=debugger)
+            jvp = adjoint_jvp(c, t, state=state)
+            res = measure_final_state(c, state, False, rng=rng)
+            return res, jvp
+
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
             results = tuple(
-                simulate(c, rng=self._rng, debugger=self._debugger, return_final_state=True)
-                for c in circuits
-            )
-            jvps = tuple(
-                adjoint_jvp(c, t, state=r[1]) for c, t, r in zip(circuits, tangents, results)
+                wrapper(c, t, rng=self._rng, debugger=self._debugger)
+                for c, t in zip(circuits, tangents)
             )
-            results = tuple(r[0] for r in results)
+            results, jvps = tuple(zip(*results))
         else:
             self._validate_multiprocessing_circuits(circuits)
 
             vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
             seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
 
-            def wrapper(c, t, rng):
-                res, final_state, _ = simulate(c, rng=rng, debugger=None, return_final_state=True)
-                jvp = adjoint_jvp(c, t, state=final_state)
-                return res, jvp
-
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
                 results = tuple(executor.map(wrapper, vanilla_circuits, tangents, seeds))
-                results, jvps = tuple(zip(*results))
+
+            results, jvps = tuple(zip(*results))
 
             # reset _rng to mimic serial behavior
             self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
@@ -518,30 +515,29 @@ def execute_and_compute_vjp(
                 f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
             )
 
+        def wrapper(c, t, rng=None, debugger=None):
+            state = get_final_state(c, debugger=debugger)
+            vjp = adjoint_vjp(c, t, state=state)
+            res = measure_final_state(c, state, False, rng=rng)
+            return res, vjp
+
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
             results = tuple(
-                simulate(c, rng=self._rng, debugger=self._debugger, return_final_state=True)
-                for c in circuits
+                wrapper(c, t, rng=self._rng, debugger=self._debugger)
+                for c, t in zip(circuits, cotangents)
             )
-            vjps = tuple(
-                adjoint_vjp(c, t, state=r[1]) for c, t, r in zip(circuits, cotangents, results)
-            )
-            results = tuple(r[0] for r in results)
+            results, vjps = tuple(zip(*results))
         else:
             self._validate_multiprocessing_circuits(circuits)
 
             vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
             seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
 
-            def wrapper(c, t, rng):
-                res, final_state, _ = simulate(c, rng=rng, debugger=None, return_final_state=True)
-                vjp = adjoint_vjp(c, t, state=final_state)
-                return res, vjp
-
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
                 results = tuple(executor.map(wrapper, vanilla_circuits, cotangents, seeds))
-                results, vjps = tuple(zip(*results))
+
+            results, vjps = tuple(zip(*results))
 
             # reset _rng to mimic serial behavior
             self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
diff --git a/pennylane/devices/qubit/__init__.py b/pennylane/devices/qubit/__init__.py
index 653199c49ac..d69490c2341 100644
--- a/pennylane/devices/qubit/__init__.py
+++ b/pennylane/devices/qubit/__init__.py
@@ -35,4 +35,4 @@
 from .measure import measure
 from .preprocess import preprocess
 from .sampling import sample_state, measure_with_samples
-from .simulate import simulate
+from .simulate import simulate, get_final_state, measure_final_state
diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index 4c7de1dc2c5..dfde18ecb42 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -23,7 +23,7 @@
 
 from .apply_operation import apply_operation
 from .initialize_state import create_initial_state
-from .simulate import _final_state
+from .simulate import get_final_state
 
 # pylint: disable=protected-access, too-many-branches
 
@@ -66,7 +66,7 @@ def adjoint_jacobian(tape: QuantumTape, state=None):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = state if state is not None else _final_state(tape)[0]
+    ket = state if state is not None else get_final_state(tape)[0]
 
     n_obs = len(tape.observables)
     bras = np.empty([n_obs] + [2] * len(tape.wires), dtype=np.complex128)
@@ -141,7 +141,7 @@ def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number], state=None):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = state if state is not None else _final_state(tape)[0]
+    ket = state if state is not None else get_final_state(tape)[0]
 
     n_obs = len(tape.observables)
     bras = np.empty([n_obs] + [2] * len(tape.wires), dtype=np.complex128)
@@ -215,7 +215,7 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = state if state is not None else _final_state(tape)[0]
+    ket = state if state is not None else get_final_state(tape)[0]
 
     obs = qml.dot(cotangents, tape.observables)
     bra = apply_operation(obs, ket)
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index fff2b58cf7f..9dc87d92207 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -24,10 +24,14 @@
 from .sampling import measure_with_samples
 
 
-def _final_state(circuit, debugger=None):
+def get_final_state(circuit, debugger=None):
     """
     Get the final state from executing the ops in the circuit
     """
+    if set(circuit.wires) != set(range(circuit.num_wires)):
+        wire_map = {w: i for i, w in enumerate(circuit.wires)}
+        circuit = qml.map_wires(circuit, wire_map)
+
     prep = circuit[0] if isinstance(circuit[0], qml.operation.StatePrep) else None
     state = create_initial_state(circuit.wires, prep)
 
@@ -45,66 +49,35 @@ def _final_state(circuit, debugger=None):
     return state, is_state_batched
 
 
-def simulate(
-    circuit: qml.tape.QuantumScript, rng=None, debugger=None, return_final_state=False
-) -> Result:
-    """Simulate a single quantum script.
-
-    This is an internal function that will be called by the successor to ``default.qubit``.
-
-    Args:
-        circuit (.QuantumScript): The single circuit to simulate
-        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
-            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
-            If no value is provided, a default RNG will be used.
-        debugger (._Debugger): The debugger to use
-        return_final_state (bool): Whether to return the final state in addition to the results
-
-    Returns:
-        tuple(TensorLike): The results of the simulation
-
-    Note that this function can return measurements for non-commuting observables simultaneously.
-
-    It does currently not support sampling or observables without diagonalizing gates.
-
-    This function assumes that all operations provide matrices.
-
-    >>> qs = qml.tape.QuantumScript([qml.RX(1.2, wires=0)], [qml.expval(qml.PauliZ(0)), qml.probs(wires=(0,1))])
-    >>> simulate(qs)
-    (0.36235775447667357,
-    tensor([0.68117888, 0.        , 0.31882112, 0.        ], requires_grad=True))
-
+def measure_final_state(circuit, state, is_state_batched, rng=None) -> Result:
+    """
+    TODO
     """
     if set(circuit.wires) != set(range(circuit.num_wires)):
         wire_map = {w: i for i, w in enumerate(circuit.wires)}
         circuit = qml.map_wires(circuit, wire_map)
 
-    state, is_state_batched = _final_state(circuit, debugger=debugger)
-
     if not circuit.shots:
         # analytic case
 
         if len(circuit.measurements) == 1:
-            results = measure(circuit.measurements[0], state, is_state_batched=is_state_batched)
+            return measure(circuit.measurements[0], state, is_state_batched=is_state_batched)
 
-        else:
-            results = tuple(
-                measure(mp, state, is_state_batched=is_state_batched) for mp in circuit.measurements
-            )
-
-        return (results, state, is_state_batched) if return_final_state else results
+        return tuple(
+            measure(mp, state, is_state_batched=is_state_batched) for mp in circuit.measurements
+        )
 
     # finite-shot case
 
     if len(circuit.measurements) == 1:
-        results = measure_with_samples(
+        return measure_with_samples(
             circuit.measurements[0],
             state,
             shots=circuit.shots,
             is_state_batched=is_state_batched,
             rng=rng,
         )
-        return (results, state, is_state_batched) if return_final_state else results
+        return results
 
     rng = default_rng(rng)
     results = tuple(
@@ -118,4 +91,35 @@ def simulate(
         # shot vector case: move the shot vector axis before the measurement axis
         results = tuple(zip(*results))
 
-    return (results, state, is_state_batched) if return_final_state else results
+    return results
+
+
+def simulate(circuit: qml.tape.QuantumScript, rng=None, debugger=None) -> Result:
+    """Simulate a single quantum script.
+
+    This is an internal function that will be called by the successor to ``default.qubit``.
+
+    Args:
+        circuit (.QuantumScript): The single circuit to simulate
+        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+            If no value is provided, a default RNG will be used.
+        debugger (._Debugger): The debugger to use
+
+    Returns:
+        tuple(TensorLike): The results of the simulation
+
+    Note that this function can return measurements for non-commuting observables simultaneously.
+
+    It does currently not support sampling or observables without diagonalizing gates.
+
+    This function assumes that all operations provide matrices.
+
+    >>> qs = qml.tape.QuantumScript([qml.RX(1.2, wires=0)], [qml.expval(qml.PauliZ(0)), qml.probs(wires=(0,1))])
+    >>> simulate(qs)
+    (0.36235775447667357,
+    tensor([0.68117888, 0.        , 0.31882112, 0.        ], requires_grad=True))
+
+    """
+    state, is_state_batched = get_final_state(circuit, debugger=debugger)
+    return measure_final_state(circuit, state, is_state_batched, rng=rng)
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index 1b7c74a926a..509963ad9ab 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -18,7 +18,7 @@
 import numpy as np
 
 import pennylane as qml
-from pennylane.devices.qubit import simulate
+from pennylane.devices.qubit import simulate, get_final_state, measure_final_state
 
 
 class TestCurrentlyUnsupportedCases:
@@ -66,24 +66,13 @@ def test_basis_state(self):
 class TestBasicCircuit:
     """Tests a basic circuit with one rx gate and two simple expectation values."""
 
-    @pytest.mark.parametrize("return_final_state", [False, True])
-    def test_basic_circuit_numpy(self, return_final_state):
+    def test_basic_circuit_numpy(self):
         """Test execution with a basic circuit."""
         phi = np.array(0.397)
         qs = qml.tape.QuantumScript(
             [qml.RX(phi, wires=0)], [qml.expval(qml.PauliY(0)), qml.expval(qml.PauliZ(0))]
         )
-        result = simulate(qs, return_final_state=return_final_state)
-
-        if return_final_state:
-            assert isinstance(result, tuple)
-            assert len(result) == 3
-
-            # the second element is the final state, and the third is whether it is batched
-            assert np.allclose(result[1], np.array([np.cos(phi / 2), -1j * np.sin(phi / 2)]))
-            assert not result[2]
-
-            result = result[0]
+        result = simulate(qs)
 
         assert isinstance(result, tuple)
         assert len(result) == 2
@@ -91,6 +80,12 @@ def test_basic_circuit_numpy(self, return_final_state):
         assert np.allclose(result[0], -np.sin(phi))
         assert np.allclose(result[1], np.cos(phi))
 
+        state, is_state_batched = get_final_state(qs)
+        result = measure_final_state(qs, state, is_state_batched)
+
+        assert np.allclose(state, np.array([np.cos(phi / 2), -1j * np.sin(phi / 2)]))
+        assert not is_state_batched
+
     @pytest.mark.autograd
     def test_autograd_results_and_backprop(self):
         """Tests execution and gradients with autograd"""
@@ -184,8 +179,7 @@ def test_tf_results_and_backprop(self):
 class TestBroadcasting:
     """Test that simulate works with broadcasted parameters"""
 
-    @pytest.mark.parametrize("return_final_state", [False, True])
-    def test_broadcasted_prep_state(self, return_final_state):
+    def test_broadcasted_prep_state(self):
         """Test that simulate works for state measurements
         when the state prep has broadcasted parameters"""
         x = np.array(1.2)
@@ -195,33 +189,28 @@ def test_broadcasted_prep_state(self, return_final_state):
         prep = [qml.QubitStateVector(np.eye(4), wires=[0, 1])]
 
         qs = qml.tape.QuantumScript(ops, measurements, prep)
-        res = simulate(qs, return_final_state=return_final_state)
-
-        if return_final_state:
-            assert isinstance(res, tuple)
-            assert len(res) == 3
-
-            # the second element is the final state, and the third is whether it is batched
-            expected = np.array(
-                [
-                    [np.cos(x / 2), 0, 0, np.sin(x / 2)],
-                    [0, np.cos(x / 2), np.sin(x / 2), 0],
-                    [-np.sin(x / 2), 0, 0, np.cos(x / 2)],
-                    [0, -np.sin(x / 2), np.cos(x / 2), 0],
-                ]
-            ).reshape((4, 2, 2))
-            assert np.allclose(res[1], expected)
-            assert res[2]
-
-            res = res[0]
+        res = simulate(qs)
 
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert np.allclose(res[0], np.array([np.cos(x), np.cos(x), -np.cos(x), -np.cos(x)]))
         assert np.allclose(res[1], np.array([np.cos(x), -np.cos(x), -np.cos(x), np.cos(x)]))
 
-    @pytest.mark.parametrize("return_final_state", [False, True])
-    def test_broadcasted_op_state(self, return_final_state):
+        state, is_state_batched = get_final_state(qs)
+        result = measure_final_state(qs, state, is_state_batched)
+        expected_state = np.array(
+            [
+                [np.cos(x / 2), 0, 0, np.sin(x / 2)],
+                [0, np.cos(x / 2), np.sin(x / 2), 0],
+                [-np.sin(x / 2), 0, 0, np.cos(x / 2)],
+                [0, -np.sin(x / 2), np.cos(x / 2), 0],
+            ]
+        ).reshape((4, 2, 2))
+
+        assert np.allclose(state, expected_state)
+        assert is_state_batched
+
+    def test_broadcasted_op_state(self):
         """Test that simulate works for state measurements
         when an operation has broadcasted parameters"""
         x = np.array([0.8, 1.0, 1.2, 1.4])
@@ -230,28 +219,24 @@ def test_broadcasted_op_state(self, return_final_state):
         measurements = [qml.expval(qml.PauliZ(i)) for i in range(2)]
 
         qs = qml.tape.QuantumScript(ops, measurements)
-        res = simulate(qs, return_final_state=return_final_state)
-
-        if return_final_state:
-            assert isinstance(res, tuple)
-            assert len(res) == 3
-
-            # the second element is the final state, and the third is whether it is batched
-            expected = np.zeros((4, 2, 2))
-            expected[:, 0, 1] = np.cos(x / 2)
-            expected[:, 1, 0] = np.sin(x / 2)
-            assert np.allclose(res[1], expected)
-            assert res[2]
-
-            res = res[0]
+        res = simulate(qs)
 
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert np.allclose(res[0], np.cos(x))
         assert np.allclose(res[1], -np.cos(x))
 
-    @pytest.mark.parametrize("return_final_state", [False, True])
-    def test_broadcasted_prep_sample(self, return_final_state):
+        state, is_state_batched = get_final_state(qs)
+        result = measure_final_state(qs, state, is_state_batched)
+
+        expected_state = np.zeros((4, 2, 2))
+        expected_state[:, 0, 1] = np.cos(x / 2)
+        expected_state[:, 1, 0] = np.sin(x / 2)
+
+        assert np.allclose(state, expected_state)
+        assert is_state_batched
+
+    def test_broadcasted_prep_sample(self):
         """Test that simulate works for sample measurements
         when the state prep has broadcasted parameters"""
         x = np.array(1.2)
@@ -261,25 +246,7 @@ def test_broadcasted_prep_sample(self, return_final_state):
         prep = [qml.QubitStateVector(np.eye(4), wires=[0, 1])]
 
         qs = qml.tape.QuantumScript(ops, measurements, prep, shots=qml.measurements.Shots(10000))
-        res = simulate(qs, rng=123, return_final_state=return_final_state)
-
-        if return_final_state:
-            assert isinstance(res, tuple)
-            assert len(res) == 3
-
-            # the second element is the final state, and the third is whether it is batched
-            expected = np.array(
-                [
-                    [np.cos(x / 2), 0, 0, np.sin(x / 2)],
-                    [0, np.cos(x / 2), np.sin(x / 2), 0],
-                    [-np.sin(x / 2), 0, 0, np.cos(x / 2)],
-                    [0, -np.sin(x / 2), np.cos(x / 2), 0],
-                ]
-            ).reshape((4, 2, 2))
-            assert np.allclose(res[1], expected)
-            assert res[2]
-
-            res = res[0]
+        res = simulate(qs, rng=123)
 
         assert isinstance(res, tuple)
         assert len(res) == 2
@@ -290,8 +257,21 @@ def test_broadcasted_prep_sample(self, return_final_state):
             res[1], np.array([np.cos(x), -np.cos(x), -np.cos(x), np.cos(x)]), atol=0.05
         )
 
-    @pytest.mark.parametrize("return_final_state", [False, True])
-    def test_broadcasted_op_sample(self, return_final_state):
+        state, is_state_batched = get_final_state(qs)
+        result = measure_final_state(qs, state, is_state_batched)
+        expected_state = np.array(
+            [
+                [np.cos(x / 2), 0, 0, np.sin(x / 2)],
+                [0, np.cos(x / 2), np.sin(x / 2), 0],
+                [-np.sin(x / 2), 0, 0, np.cos(x / 2)],
+                [0, -np.sin(x / 2), np.cos(x / 2), 0],
+            ]
+        ).reshape((4, 2, 2))
+
+        assert np.allclose(state, expected_state)
+        assert is_state_batched
+
+    def test_broadcasted_op_sample(self):
         """Test that simulate works for sample measurements
         when an operation has broadcasted parameters"""
         x = np.array([0.8, 1.0, 1.2, 1.4])
@@ -300,26 +280,23 @@ def test_broadcasted_op_sample(self, return_final_state):
         measurements = [qml.expval(qml.PauliZ(i)) for i in range(2)]
 
         qs = qml.tape.QuantumScript(ops, measurements, shots=qml.measurements.Shots(10000))
-        res = simulate(qs, rng=123, return_final_state=return_final_state)
-
-        if return_final_state:
-            assert isinstance(res, tuple)
-            assert len(res) == 3
-
-            # the second element is the final state, and the third is whether it is batched
-            expected = np.zeros((4, 2, 2))
-            expected[:, 0, 1] = np.cos(x / 2)
-            expected[:, 1, 0] = np.sin(x / 2)
-            assert np.allclose(res[1], expected)
-            assert res[2]
-
-            res = res[0]
+        res = simulate(qs, rng=123)
 
         assert isinstance(res, tuple)
         assert len(res) == 2
         assert np.allclose(res[0], np.cos(x), atol=0.05)
         assert np.allclose(res[1], -np.cos(x), atol=0.05)
 
+        state, is_state_batched = get_final_state(qs)
+        result = measure_final_state(qs, state, is_state_batched)
+
+        expected_state = np.zeros((4, 2, 2))
+        expected_state[:, 0, 1] = np.cos(x / 2)
+        expected_state[:, 1, 0] = np.sin(x / 2)
+
+        assert np.allclose(state, expected_state)
+        assert is_state_batched
+
 
 class TestDebugger:
     """Tests that the debugger works for a simple circuit"""

From 2bbfd63ef350bece153b803510108f3d85c602fe Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Fri, 21 Jul 2023 15:41:46 -0400
Subject: [PATCH 06/19] Add some tests

---
 .../devices/experimental/default_qubit_2.py   | 78 ++++++++++++-------
 pennylane/devices/qubit/simulate.py           |  5 +-
 .../experimental/test_default_qubit_2.py      | 61 +++++++++++++--
 3 files changed, 108 insertions(+), 36 deletions(-)

diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
index 995d7203cd7..2998fafc9f7 100644
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ b/pennylane/devices/experimental/default_qubit_2.py
@@ -297,8 +297,11 @@ def execute_and_compute_derivatives(
         if self.tracker.active:
             for c in circuits:
                 self.tracker.update(resources=c.specs["resources"])
-            self.tracker.update(batches=1, executions=len(circuits))
-            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.update(
+                execute_and_derivative_batches=1,
+                executions=len(circuits),
+                derivatives=len(circuits),
+            )
             self.tracker.record()
 
         if execution_config.gradient_method != "adjoint":
@@ -306,15 +309,11 @@ def execute_and_compute_derivatives(
                 f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
             )
 
-        def wrapper(c, rng=None, debugger=None):
-            state = get_final_state(c, debugger=debugger)
-            jac = adjoint_jacobian(c, state=state)
-            res = measure_final_state(c, state, False, rng=rng)
-            return res, jac
-
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
-            results = tuple(wrapper(c, rng=self._rng, debugger=self._debugger) for c in circuits)
+            results = tuple(
+                _adjoint_jac_wrapper(c, rng=self._rng, debugger=self._debugger) for c in circuits
+            )
             results, jacs = tuple(zip(*results))
         else:
             self._validate_multiprocessing_circuits(circuits)
@@ -323,7 +322,7 @@ def wrapper(c, rng=None, debugger=None):
             seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
 
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                results = tuple(executor.map(wrapper, vanilla_circuits, seeds))
+                results = tuple(executor.map(_adjoint_jac_wrapper, vanilla_circuits, seeds))
 
             results, jacs = tuple(zip(*results))
 
@@ -400,8 +399,11 @@ def execute_and_compute_jvp(
         if self.tracker.active:
             for c in circuits:
                 self.tracker.update(resources=c.specs["resources"])
-            self.tracker.update(batches=1, executions=len(circuits))
-            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.update(
+                execute_and_derivative_batches=1,
+                executions=len(circuits),
+                derivatives=len(circuits),
+            )
             self.tracker.record()
 
         if execution_config.gradient_method != "adjoint":
@@ -409,16 +411,10 @@ def execute_and_compute_jvp(
                 f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
             )
 
-        def wrapper(c, t, rng=None, debugger=None):
-            state = get_final_state(c, debugger=debugger)
-            jvp = adjoint_jvp(c, t, state=state)
-            res = measure_final_state(c, state, False, rng=rng)
-            return res, jvp
-
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
             results = tuple(
-                wrapper(c, t, rng=self._rng, debugger=self._debugger)
+                _adjoint_jvp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
                 for c, t in zip(circuits, tangents)
             )
             results, jvps = tuple(zip(*results))
@@ -429,7 +425,9 @@ def wrapper(c, t, rng=None, debugger=None):
             seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
 
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                results = tuple(executor.map(wrapper, vanilla_circuits, tangents, seeds))
+                results = tuple(
+                    executor.map(_adjoint_jvp_wrapper, vanilla_circuits, tangents, seeds)
+                )
 
             results, jvps = tuple(zip(*results))
 
@@ -506,8 +504,11 @@ def execute_and_compute_vjp(
         if self.tracker.active:
             for c in circuits:
                 self.tracker.update(resources=c.specs["resources"])
-            self.tracker.update(batches=1, executions=len(circuits))
-            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.update(
+                execute_and_derivative_batches=1,
+                executions=len(circuits),
+                derivatives=len(circuits),
+            )
             self.tracker.record()
 
         if execution_config.gradient_method != "adjoint":
@@ -515,16 +516,10 @@ def execute_and_compute_vjp(
                 f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
             )
 
-        def wrapper(c, t, rng=None, debugger=None):
-            state = get_final_state(c, debugger=debugger)
-            vjp = adjoint_vjp(c, t, state=state)
-            res = measure_final_state(c, state, False, rng=rng)
-            return res, vjp
-
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
             results = tuple(
-                wrapper(c, t, rng=self._rng, debugger=self._debugger)
+                _adjoint_vjp_wrapper(c, t, rng=self._rng, debugger=self._debugger)
                 for c, t in zip(circuits, cotangents)
             )
             results, vjps = tuple(zip(*results))
@@ -535,7 +530,9 @@ def wrapper(c, t, rng=None, debugger=None):
             seeds = self._rng.integers(2**31 - 1, size=len(vanilla_circuits))
 
             with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                results = tuple(executor.map(wrapper, vanilla_circuits, cotangents, seeds))
+                results = tuple(
+                    executor.map(_adjoint_vjp_wrapper, vanilla_circuits, cotangents, seeds)
+                )
 
             results, vjps = tuple(zip(*results))
 
@@ -613,3 +610,24 @@ def _validate_multiprocessing_workers(max_workers):
             environment variable `{varname}={num_threads_suggest}`.""",
             UserWarning,
         )
+
+
+def _adjoint_jac_wrapper(c, rng=None, debugger=None):
+    state, is_state_batched = get_final_state(c, debugger=debugger)
+    jac = adjoint_jacobian(c, state=state)
+    res = measure_final_state(c, state, is_state_batched, rng=rng)
+    return res, jac
+
+
+def _adjoint_jvp_wrapper(c, t, rng=None, debugger=None):
+    state, is_state_batched = get_final_state(c, debugger=debugger)
+    jvp = adjoint_jvp(c, t, state=state)
+    res = measure_final_state(c, state, is_state_batched, rng=rng)
+    return res, jvp
+
+
+def _adjoint_vjp_wrapper(c, t, rng=None, debugger=None):
+    state, is_state_batched = get_final_state(c, debugger=debugger)
+    vjp = adjoint_vjp(c, t, state=state)
+    res = measure_final_state(c, state, is_state_batched, rng=rng)
+    return res, vjp
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 9dc87d92207..093fa71f90d 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -32,7 +32,10 @@ def get_final_state(circuit, debugger=None):
         wire_map = {w: i for i, w in enumerate(circuit.wires)}
         circuit = qml.map_wires(circuit, wire_map)
 
-    prep = circuit[0] if isinstance(circuit[0], qml.operation.StatePrep) else None
+    prep = None
+    if len(circuit) > 0 and isinstance(circuit[0], qml.operation.StatePrep):
+        prep = circuit[0]
+
     state = create_initial_state(circuit.wires, prep)
 
     # initial state is batched only if the state preparation (if it exists) is batched
diff --git a/tests/devices/experimental/test_default_qubit_2.py b/tests/devices/experimental/test_default_qubit_2.py
index ae5b03fdb35..b411a47638c 100644
--- a/tests/devices/experimental/test_default_qubit_2.py
+++ b/tests/devices/experimental/test_default_qubit_2.py
@@ -154,6 +154,30 @@ def test_tracking_batch(self):
         }
         assert tracker.latest == {"batches": 1, "executions": 2}
 
+    def test_tracking_execute_and_derivatives(self):
+        """Test that the execute_and_compute_* calls are being tracked for the
+        experimental default qubit device"""
+
+        qs = qml.tape.QuantumScript([], [qml.expval(qml.PauliZ(0))])
+        dev = DefaultQubit2()
+        config = ExecutionConfig(gradient_method="adjoint")
+
+        with qml.Tracker(dev) as tracker:
+            dev.compute_derivatives(qs, config)
+            dev.execute_and_compute_derivatives([qs] * 2, config)
+            dev.compute_jvp([qs] * 3, [(0,)] * 3, config)
+            dev.execute_and_compute_jvp([qs] * 4, [(0,)] * 4, config)
+            dev.compute_vjp([qs] * 5, [(0,)] * 5, config)
+            dev.execute_and_compute_vjp([qs] * 6, [(0,)] * 6, config)
+
+        assert tracker.history == {
+            "executions": [2, 4, 6],
+            "derivatives": [1, 2, 3, 4, 5, 6],
+            "derivative_batches": [1, 1, 1],
+            "execute_and_derivative_batches": [1, 1, 1],
+            "resources": [Resources(num_wires=1)] * 12,
+        }
+
     def test_tracking_resources(self):
         """Test that resources are tracked for the experimental default qubit device."""
         qs = qml.tape.QuantumScript(
@@ -234,24 +258,36 @@ def test_supports_backprop(self):
         """Test that DefaultQubit2 says that it supports backpropagation."""
         dev = DefaultQubit2()
         assert dev.supports_derivatives() is True
+        assert dev.supports_jvp() is True
+        assert dev.supports_vjp() is True
 
         config = ExecutionConfig(gradient_method="backprop")
         assert dev.supports_derivatives(config) is True
+        assert dev.supports_jvp(config) is True
+        assert dev.supports_vjp(config) is True
 
         qs = qml.tape.QuantumScript([], [qml.state()])
-        assert dev.supports_derivatives(config, qs)
+        assert dev.supports_derivatives(config, qs) is True
+        assert dev.supports_jvp(config, qs) is True
+        assert dev.supports_vjp(config, qs) is True
 
         config = ExecutionConfig(gradient_method="backprop", device_options={"max_workers": 1})
         assert dev.supports_derivatives(config) is False
+        assert dev.supports_jvp(config) is False
+        assert dev.supports_vjp(config) is False
 
     def test_supports_adjoint(self):
         """Test that DefaultQubit2 says that it supports adjoint differentiation."""
         dev = DefaultQubit2()
         config = ExecutionConfig(gradient_method="adjoint")
         assert dev.supports_derivatives(config) is True
+        assert dev.supports_jvp(config) is True
+        assert dev.supports_vjp(config) is True
 
         qs = qml.tape.QuantumScript([], [qml.expval(qml.PauliZ(0))])
         assert dev.supports_derivatives(config, qs) is True
+        assert dev.supports_jvp(config, qs) is True
+        assert dev.supports_vjp(config, qs) is True
 
     def test_doesnt_support_adjoint_with_invalid_tape(self):
         """Tests that DefaultQubit2 does not support adjoint differentiation with invalid circuits."""
@@ -259,13 +295,17 @@ def test_doesnt_support_adjoint_with_invalid_tape(self):
         config = ExecutionConfig(gradient_method="adjoint")
         circuit = qml.tape.QuantumScript([], [qml.probs()])
         assert dev.supports_derivatives(config, circuit=circuit) is False
+        assert dev.supports_jvp(config, circuit=circuit) is False
+        assert dev.supports_vjp(config, circuit=circuit) is False
 
     @pytest.mark.parametrize("gradient_method", ["parameter-shift", "finite-diff", "device"])
     def test_doesnt_support_other_gradient_methods(self, gradient_method):
         """Test that DefaultQubit2 currently does not support other gradient methods natively."""
         dev = DefaultQubit2()
         config = ExecutionConfig(gradient_method=gradient_method)
-        assert not dev.supports_derivatives(config)
+        assert dev.supports_derivatives(config) is False
+        assert dev.supports_jvp(config) is False
+        assert dev.supports_vjp(config) is False
 
 
 class TestBasicCircuit:
@@ -997,20 +1037,31 @@ class TestAdjointDifferentiation:
     ec = ExecutionConfig(gradient_method="adjoint")
 
     @pytest.mark.parametrize("max_workers", [None, 1, 2])
-    def test_single_circuit(self, max_workers):
-        """Tests a basic example with a single circuit."""
+    def test_compute_derivatives_single_circuit(self, max_workers):
+        """Tests compute derivatives with a single circuit."""
         dev = DefaultQubit2(max_workers=max_workers)
         x = np.array(np.pi / 7)
         qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
         qs = validate_and_expand_adjoint(qs)
         expected_grad = -qml.math.sin(x)
         actual_grad = dev.compute_derivatives(qs, self.ec)
+
         assert isinstance(actual_grad, np.ndarray)
         assert actual_grad.shape == ()  # pylint: disable=no-member
         assert np.isclose(actual_grad, expected_grad)
 
-        expected_val = qml.math.cos(x)
+    @pytest.mark.parametrize("max_workers", [None, 1, 2])
+    def test_execute_and_compute_derivatives_single_circuit(self, max_workers):
+        """Test execute and compute derivatives with a single circuit."""
+        dev = DefaultQubit2(max_workers=max_workers)
+        x = np.array(np.pi / 7)
+        qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+        qs = validate_and_expand_adjoint(qs)
+
         actual_val, actual_grad = dev.execute_and_compute_derivatives(qs, self.ec)
+        expected_grad = -qml.math.sin(x)
+        expected_val = qml.math.cos(x)
+
         assert np.isclose(actual_val, expected_val)
         assert np.isclose(actual_grad, expected_grad)
 

From d5644a8d0c3e748b01ff8f35b0fbf64d88e53a23 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Fri, 21 Jul 2023 15:47:17 -0400
Subject: [PATCH 07/19] pylint

---
 pennylane/devices/qubit/adjoint_jacobian.py | 1 -
 pennylane/devices/qubit/simulate.py         | 1 -
 2 files changed, 2 deletions(-)

diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index dfde18ecb42..6bca30831b8 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -22,7 +22,6 @@
 from pennylane.tape import QuantumTape
 
 from .apply_operation import apply_operation
-from .initialize_state import create_initial_state
 from .simulate import get_final_state
 
 # pylint: disable=protected-access, too-many-branches
diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index 093fa71f90d..f209f6c0a99 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -80,7 +80,6 @@ def measure_final_state(circuit, state, is_state_batched, rng=None) -> Result:
             is_state_batched=is_state_batched,
             rng=rng,
         )
-        return results
 
     rng = default_rng(rng)
     results = tuple(

From 4999f40771cfc9821add15853b3fd87bae698771 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Mon, 24 Jul 2023 10:28:44 -0400
Subject: [PATCH 08/19] more device tests

---
 .../experimental/test_default_qubit_2.py      | 212 ++++++++++++++++--
 1 file changed, 191 insertions(+), 21 deletions(-)

diff --git a/tests/devices/experimental/test_default_qubit_2.py b/tests/devices/experimental/test_default_qubit_2.py
index b411a47638c..4f804b22fbb 100644
--- a/tests/devices/experimental/test_default_qubit_2.py
+++ b/tests/devices/experimental/test_default_qubit_2.py
@@ -1031,42 +1031,30 @@ def test_tf_backprop(self, convert_to_hamiltonian):
         assert qml.math.allclose(g1, g2)
 
 
+@pytest.mark.parametrize("max_workers", [None, 1, 2])
 class TestAdjointDifferentiation:
     """Tests adjoint differentiation integration with DefaultQubit2."""
 
     ec = ExecutionConfig(gradient_method="adjoint")
 
-    @pytest.mark.parametrize("max_workers", [None, 1, 2])
-    def test_compute_derivatives_single_circuit(self, max_workers):
-        """Tests compute derivatives with a single circuit."""
+    def test_derivatives_single_circuit(self, max_workers):
+        """Tests derivatives with a single circuit."""
         dev = DefaultQubit2(max_workers=max_workers)
         x = np.array(np.pi / 7)
         qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
         qs = validate_and_expand_adjoint(qs)
         expected_grad = -qml.math.sin(x)
         actual_grad = dev.compute_derivatives(qs, self.ec)
-
         assert isinstance(actual_grad, np.ndarray)
         assert actual_grad.shape == ()  # pylint: disable=no-member
         assert np.isclose(actual_grad, expected_grad)
 
-    @pytest.mark.parametrize("max_workers", [None, 1, 2])
-    def test_execute_and_compute_derivatives_single_circuit(self, max_workers):
-        """Test execute and compute derivatives with a single circuit."""
-        dev = DefaultQubit2(max_workers=max_workers)
-        x = np.array(np.pi / 7)
-        qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
-        qs = validate_and_expand_adjoint(qs)
-
-        actual_val, actual_grad = dev.execute_and_compute_derivatives(qs, self.ec)
-        expected_grad = -qml.math.sin(x)
         expected_val = qml.math.cos(x)
-
+        actual_val, actual_grad = dev.execute_and_compute_derivatives(qs, self.ec)
         assert np.isclose(actual_val, expected_val)
         assert np.isclose(actual_grad, expected_grad)
 
-    @pytest.mark.parametrize("max_workers", [None, 1, 2])
-    def test_list_with_single_circuit(self, max_workers):
+    def test_derivatives_list_with_single_circuit(self, max_workers):
         """Tests a basic example with a batch containing a single circuit."""
         dev = DefaultQubit2(max_workers=max_workers)
         x = np.array(np.pi / 7)
@@ -1083,8 +1071,7 @@ def test_list_with_single_circuit(self, max_workers):
         assert np.isclose(expected_val, actual_val[0])
         assert np.isclose(expected_grad, actual_grad[0])
 
-    @pytest.mark.parametrize("max_workers", [None, 1, 2])
-    def test_many_tapes_many_results(self, max_workers):
+    def test_derivatives_many_tapes_many_results(self, max_workers):
         """Tests a basic example with a batch of circuits of varying return shapes."""
         dev = DefaultQubit2(max_workers=max_workers)
         x = np.array(np.pi / 7)
@@ -1098,8 +1085,7 @@ def test_many_tapes_many_results(self, max_workers):
         assert isinstance(actual_grad[1], tuple)
         assert qml.math.allclose(actual_grad[1], expected_grad[1])
 
-    @pytest.mark.parametrize("max_workers", [None, 1, 2])
-    def test_integration(self, max_workers):
+    def test_derivatives_integration(self, max_workers):
         """Tests the expected workflow done by a calling method."""
         dev = DefaultQubit2(max_workers=max_workers)
         x = np.array(np.pi / 7)
@@ -1119,6 +1105,190 @@ def test_integration(self, max_workers):
         assert isinstance(actual_grad[1], tuple)
         assert qml.math.allclose(actual_grad[1], expected_grad[1])
 
+    def test_jvps_single_circuit(self, max_workers):
+        """Tests jvps with a single circuit."""
+        dev = DefaultQubit2(max_workers=max_workers)
+        x = np.array(np.pi / 7)
+        tangent = (0.456,)
+
+        qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+        qs = validate_and_expand_adjoint(qs)
+
+        expected_grad = -qml.math.sin(x) * tangent[0]
+        actual_grad = dev.compute_jvp(qs, tangent, self.ec)
+        assert isinstance(actual_grad, np.ndarray)
+        assert actual_grad.shape == ()  # pylint: disable=no-member
+        assert np.isclose(actual_grad, expected_grad)
+
+        expected_val = qml.math.cos(x)
+        actual_val, actual_grad = dev.execute_and_compute_jvp(qs, tangent, self.ec)
+        assert np.isclose(actual_val, expected_val)
+        assert np.isclose(actual_grad, expected_grad)
+
+    def test_jvps_list_with_single_circuit(self, max_workers):
+        """Tests a basic example with a batch containing a single circuit."""
+        dev = DefaultQubit2(max_workers=max_workers)
+        x = np.array(np.pi / 7)
+        tangent = (0.456,)
+
+        qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+        qs = validate_and_expand_adjoint(qs)
+
+        expected_grad = -qml.math.sin(x) * tangent[0]
+        actual_grad = dev.compute_jvp([qs], [tangent], self.ec)
+        assert isinstance(actual_grad, tuple)
+        assert isinstance(actual_grad[0], np.ndarray)
+        assert np.isclose(actual_grad[0], expected_grad)
+
+        expected_val = qml.math.cos(x)
+        actual_val, actual_grad = dev.execute_and_compute_jvp([qs], [tangent], self.ec)
+        assert np.isclose(expected_val, actual_val[0])
+        assert np.isclose(expected_grad, actual_grad[0])
+
+    def test_jvps_many_tapes_many_results(self, max_workers):
+        """Tests a basic example with a batch of circuits of varying return shapes."""
+        dev = DefaultQubit2(max_workers=max_workers)
+        x = np.array(np.pi / 7)
+        single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+        multi_meas = qml.tape.QuantumScript(
+            [qml.RY(x, 0)], [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(0))]
+        )
+        tangents = [(0.456,), (0.789,)]
+
+        expected_grad = (
+            -qml.math.sin(x) * tangents[0][0],
+            (qml.math.cos(x) * tangents[1][0], -qml.math.sin(x) * tangents[1][0]),
+        )
+        actual_grad = dev.compute_jvp([single_meas, multi_meas], tangents, self.ec)
+        assert np.isclose(actual_grad[0], expected_grad[0])
+        assert isinstance(actual_grad[1], tuple)
+        assert qml.math.allclose(actual_grad[1], expected_grad[1])
+
+        expected_val = (qml.math.cos(x), (qml.math.sin(x), qml.math.cos(x)))
+        actual_val, actual_grad = dev.execute_and_compute_jvp(
+            [single_meas, multi_meas], tangents, self.ec
+        )
+        assert np.isclose(actual_val[0], expected_val[0])
+        assert qml.math.allclose(actual_val[1], expected_val[1])
+        assert np.isclose(actual_grad[0], expected_grad[0])
+        assert qml.math.allclose(actual_grad[1], expected_grad[1])
+
+    def test_jvps_integration(self, max_workers):
+        """Tests the expected workflow done by a calling method."""
+        dev = DefaultQubit2(max_workers=max_workers)
+        x = np.array(np.pi / 7)
+
+        single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+        multi_meas = qml.tape.QuantumScript(
+            [qml.RY(x, 0)], [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(0))]
+        )
+        tangents = [(0.456,), (0.789,)]
+
+        circuits, _, new_ec = dev.preprocess([single_meas, multi_meas], self.ec)
+        actual_grad = dev.compute_jvp(circuits, tangents, self.ec)
+        expected_grad = (
+            -qml.math.sin(x) * tangents[0][0],
+            (qml.math.cos(x) * tangents[1][0], -qml.math.sin(x) * tangents[1][0]),
+        )
+
+        assert new_ec.use_device_gradient
+        assert new_ec.grad_on_execution
+
+        assert np.isclose(actual_grad[0], expected_grad[0])
+        assert isinstance(actual_grad[1], tuple)
+        assert qml.math.allclose(actual_grad[1], expected_grad[1])
+
+    def test_vjps_single_circuit(self, max_workers):
+        """Tests vjps with a single circuit."""
+        dev = DefaultQubit2(max_workers=max_workers)
+        x = np.array(np.pi / 7)
+        cotangent = (0.456,)
+
+        qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+        qs = validate_and_expand_adjoint(qs)
+
+        expected_grad = -qml.math.sin(x) * cotangent[0]
+        actual_grad = dev.compute_vjp(qs, cotangent, self.ec)
+        assert isinstance(actual_grad, np.ndarray)
+        assert actual_grad.shape == ()  # pylint: disable=no-member
+        assert np.isclose(actual_grad, expected_grad)
+
+        expected_val = qml.math.cos(x)
+        actual_val, actual_grad = dev.execute_and_compute_vjp(qs, cotangent, self.ec)
+        assert np.isclose(actual_val, expected_val)
+        assert np.isclose(actual_grad, expected_grad)
+
+    def test_vjps_list_with_single_circuit(self, max_workers):
+        """Tests a basic example with a batch containing a single circuit."""
+        dev = DefaultQubit2(max_workers=max_workers)
+        x = np.array(np.pi / 7)
+        cotangent = (0.456,)
+
+        qs = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+        qs = validate_and_expand_adjoint(qs)
+
+        expected_grad = -qml.math.sin(x) * cotangent[0]
+        actual_grad = dev.compute_vjp([qs], [cotangent], self.ec)
+        assert isinstance(actual_grad, tuple)
+        assert isinstance(actual_grad[0], np.ndarray)
+        assert np.isclose(actual_grad[0], expected_grad)
+
+        expected_val = qml.math.cos(x)
+        actual_val, actual_grad = dev.execute_and_compute_vjp([qs], [cotangent], self.ec)
+        assert np.isclose(expected_val, actual_val[0])
+        assert np.isclose(expected_grad, actual_grad[0])
+
+    def test_vjps_many_tapes_many_results(self, max_workers):
+        """Tests a basic example with a batch of circuits of varying return shapes."""
+        dev = DefaultQubit2(max_workers=max_workers)
+        x = np.array(np.pi / 7)
+        single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+        multi_meas = qml.tape.QuantumScript(
+            [qml.RY(x, 0)], [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(0))]
+        )
+        cotangents = [(0.456,), (0.789, 0.123)]
+
+        expected_grad = (
+            -qml.math.sin(x) * cotangents[0][0],
+            qml.math.cos(x) * cotangents[1][0] - qml.math.sin(x) * cotangents[1][1],
+        )
+        actual_grad = dev.compute_vjp([single_meas, multi_meas], cotangents, self.ec)
+        assert np.isclose(actual_grad[0], expected_grad[0])
+        assert np.isclose(actual_grad[1], expected_grad[1])
+
+        expected_val = (qml.math.cos(x), (qml.math.sin(x), qml.math.cos(x)))
+        actual_val, actual_grad = dev.execute_and_compute_vjp(
+            [single_meas, multi_meas], cotangents, self.ec
+        )
+        assert np.isclose(actual_val[0], expected_val[0])
+        assert qml.math.allclose(actual_val[1], expected_val[1])
+        assert np.isclose(actual_grad[0], expected_grad[0])
+        assert np.isclose(actual_grad[1], expected_grad[1])
+
+    def test_vjps_integration(self, max_workers):
+        """Tests the expected workflow done by a calling method."""
+        dev = DefaultQubit2(max_workers=max_workers)
+        x = np.array(np.pi / 7)
+
+        single_meas = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.PauliZ(0))])
+        multi_meas = qml.tape.QuantumScript(
+            [qml.RY(x, 0)], [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(0))]
+        )
+        cotangents = [(0.456,), (0.789, 0.123)]
+
+        circuits, _, new_ec = dev.preprocess([single_meas, multi_meas], self.ec)
+        actual_grad = dev.compute_vjp(circuits, cotangents, self.ec)
+        expected_grad = (
+            -qml.math.sin(x) * cotangents[0][0],
+            qml.math.cos(x) * cotangents[1][0] - qml.math.sin(x) * cotangents[1][1],
+        )
+
+        assert new_ec.use_device_gradient
+        assert new_ec.grad_on_execution
+
+        assert np.isclose(actual_grad[0], expected_grad[0])
+        assert np.isclose(actual_grad[1], expected_grad[1])
+
 
 class TestPreprocessingIntegration:
     """Test preprocess produces output that can be executed by the device."""

From 64c298866c0e9b4cc95eb00340cdd223107a7dc7 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Mon, 24 Jul 2023 11:15:30 -0400
Subject: [PATCH 09/19] pylint and fix tests

---
 tests/devices/qubit/test_simulate.py          | 33 ++++++++++++++++---
 .../test_autograd_default_qubit_2.py          |  5 ++-
 .../test_jax_default_qubit_2.py               |  5 ++-
 .../test_tensorflow_default_qubit_2.py        |  5 ++-
 .../test_torch_default_qubit_2.py             |  5 ++-
 5 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index 509963ad9ab..e6ef6e0c269 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -86,6 +86,11 @@ def test_basic_circuit_numpy(self):
         assert np.allclose(state, np.array([np.cos(phi / 2), -1j * np.sin(phi / 2)]))
         assert not is_state_batched
 
+        assert isinstance(result, tuple)
+        assert len(result) == 2
+        assert np.allclose(result[0], -np.sin(phi))
+        assert np.allclose(result[1], np.cos(phi))
+
     @pytest.mark.autograd
     def test_autograd_results_and_backprop(self):
         """Tests execution and gradients with autograd"""
@@ -197,7 +202,7 @@ def test_broadcasted_prep_state(self):
         assert np.allclose(res[1], np.array([np.cos(x), -np.cos(x), -np.cos(x), np.cos(x)]))
 
         state, is_state_batched = get_final_state(qs)
-        result = measure_final_state(qs, state, is_state_batched)
+        res = measure_final_state(qs, state, is_state_batched)
         expected_state = np.array(
             [
                 [np.cos(x / 2), 0, 0, np.sin(x / 2)],
@@ -209,6 +214,10 @@ def test_broadcasted_prep_state(self):
 
         assert np.allclose(state, expected_state)
         assert is_state_batched
+        assert isinstance(res, tuple)
+        assert len(res) == 2
+        assert np.allclose(res[0], np.array([np.cos(x), np.cos(x), -np.cos(x), -np.cos(x)]))
+        assert np.allclose(res[1], np.array([np.cos(x), -np.cos(x), -np.cos(x), np.cos(x)]))
 
     def test_broadcasted_op_state(self):
         """Test that simulate works for state measurements
@@ -227,7 +236,7 @@ def test_broadcasted_op_state(self):
         assert np.allclose(res[1], -np.cos(x))
 
         state, is_state_batched = get_final_state(qs)
-        result = measure_final_state(qs, state, is_state_batched)
+        res = measure_final_state(qs, state, is_state_batched)
 
         expected_state = np.zeros((4, 2, 2))
         expected_state[:, 0, 1] = np.cos(x / 2)
@@ -235,6 +244,10 @@ def test_broadcasted_op_state(self):
 
         assert np.allclose(state, expected_state)
         assert is_state_batched
+        assert isinstance(res, tuple)
+        assert len(res) == 2
+        assert np.allclose(res[0], np.cos(x))
+        assert np.allclose(res[1], -np.cos(x))
 
     def test_broadcasted_prep_sample(self):
         """Test that simulate works for sample measurements
@@ -258,7 +271,7 @@ def test_broadcasted_prep_sample(self):
         )
 
         state, is_state_batched = get_final_state(qs)
-        result = measure_final_state(qs, state, is_state_batched)
+        res = measure_final_state(qs, state, is_state_batched, rng=123)
         expected_state = np.array(
             [
                 [np.cos(x / 2), 0, 0, np.sin(x / 2)],
@@ -270,6 +283,14 @@ def test_broadcasted_prep_sample(self):
 
         assert np.allclose(state, expected_state)
         assert is_state_batched
+        assert isinstance(res, tuple)
+        assert len(res) == 2
+        assert np.allclose(
+            res[0], np.array([np.cos(x), np.cos(x), -np.cos(x), -np.cos(x)]), atol=0.05
+        )
+        assert np.allclose(
+            res[1], np.array([np.cos(x), -np.cos(x), -np.cos(x), np.cos(x)]), atol=0.05
+        )
 
     def test_broadcasted_op_sample(self):
         """Test that simulate works for sample measurements
@@ -288,7 +309,7 @@ def test_broadcasted_op_sample(self):
         assert np.allclose(res[1], -np.cos(x), atol=0.05)
 
         state, is_state_batched = get_final_state(qs)
-        result = measure_final_state(qs, state, is_state_batched)
+        result = measure_final_state(qs, state, is_state_batched, rng=123)
 
         expected_state = np.zeros((4, 2, 2))
         expected_state[:, 0, 1] = np.cos(x / 2)
@@ -296,6 +317,10 @@ def test_broadcasted_op_sample(self):
 
         assert np.allclose(state, expected_state)
         assert is_state_batched
+        assert isinstance(res, tuple)
+        assert len(res) == 2
+        assert np.allclose(res[0], np.cos(x), atol=0.05)
+        assert np.allclose(res[1], -np.cos(x), atol=0.05)
 
 
 class TestDebugger:
diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
index 296d8c2c5f1..7c0fcfd8cb6 100644
--- a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py
@@ -158,7 +158,10 @@ def cost(a, b):
         with device.tracker:
             res = cost(a, b)
 
-        assert device.tracker.totals["batches"] == 1
+        if execute_kwargs.get("grad_on_execution", False):
+            assert device.tracker.totals["execute_and_derivative_batches"] == 1
+        else:
+            assert device.tracker.totals["batches"] == 1
         assert device.tracker.totals["executions"] == 2  # different wires so different hashes
 
         assert len(res) == 2
diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
index 4fdf6b6d2e0..d008cd33e14 100644
--- a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py
@@ -141,7 +141,10 @@ def cost(a, b):
         with device.tracker:
             res = cost(a, b)
 
-        assert device.tracker.totals["batches"] == 1
+        if execute_kwargs.get("gradient_fn", None) == "adjoint":
+            assert device.tracker.totals["execute_and_derivative_batches"] == 1
+        else:
+            assert device.tracker.totals["batches"] == 1
         assert device.tracker.totals["executions"] == 2  # different wires so different hashes
 
         assert len(res) == 2
diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
index 9ec4b133288..0b48df29678 100644
--- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py
@@ -147,7 +147,10 @@ def cost(a, b):
         with device.tracker:
             res = cost(a, b)
 
-        assert device.tracker.totals["batches"] == 1
+        if execute_kwargs.get("gradient_fn", None) == "adjoint":
+            assert device.tracker.totals["execute_and_derivative_batches"] == 1
+        else:
+            assert device.tracker.totals["batches"] == 1
         assert device.tracker.totals["executions"] == 2  # different wires so different hashes
 
         assert len(res) == 2
diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
index a09ec5abeeb..02228bf7853 100644
--- a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
+++ b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py
@@ -164,7 +164,10 @@ def cost(a, b):
         with device.tracker:
             res = cost(a, b)
 
-        assert device.tracker.totals["batches"] == 1
+        if execute_kwargs.get("grad_on_execution", False):
+            assert device.tracker.totals["execute_and_derivative_batches"] == 1
+        else:
+            assert device.tracker.totals["batches"] == 1
         assert device.tracker.totals["executions"] == 2  # different wires so different hashes
 
         assert len(res) == 2

From 1162c898a40c4d9bf166f207fa7b2843b406b582 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Mon, 24 Jul 2023 13:06:31 -0400
Subject: [PATCH 10/19] change tracker keys

---
 pennylane/devices/experimental/default_qubit_2.py  | 12 ++++--------
 tests/devices/experimental/test_default_qubit_2.py | 12 +++++++++---
 2 files changed, 13 insertions(+), 11 deletions(-)

diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
index 2998fafc9f7..dc72edfd1c4 100644
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ b/pennylane/devices/experimental/default_qubit_2.py
@@ -363,7 +363,7 @@ def compute_jvp(
             tangents = [tangents]
 
         if self.tracker.active:
-            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.update(jvp_batches=1, jvps=len(circuits))
             self.tracker.record()
 
         if execution_config.gradient_method != "adjoint":
@@ -400,9 +400,7 @@ def execute_and_compute_jvp(
             for c in circuits:
                 self.tracker.update(resources=c.specs["resources"])
             self.tracker.update(
-                execute_and_derivative_batches=1,
-                executions=len(circuits),
-                derivatives=len(circuits),
+                execute_and_jvp_batches=1, executions=len(circuits), jvps=len(circuits)
             )
             self.tracker.record()
 
@@ -468,7 +466,7 @@ def compute_vjp(
             cotangents = [cotangents]
 
         if self.tracker.active:
-            self.tracker.update(derivative_batches=1, derivatives=len(circuits))
+            self.tracker.update(vjp_batches=1, vjps=len(circuits))
             self.tracker.record()
 
         if execution_config.gradient_method != "adjoint":
@@ -505,9 +503,7 @@ def execute_and_compute_vjp(
             for c in circuits:
                 self.tracker.update(resources=c.specs["resources"])
             self.tracker.update(
-                execute_and_derivative_batches=1,
-                executions=len(circuits),
-                derivatives=len(circuits),
+                execute_and_vjp_batches=1, executions=len(circuits), vjps=len(circuits)
             )
             self.tracker.record()
 
diff --git a/tests/devices/experimental/test_default_qubit_2.py b/tests/devices/experimental/test_default_qubit_2.py
index 4f804b22fbb..9c171902446 100644
--- a/tests/devices/experimental/test_default_qubit_2.py
+++ b/tests/devices/experimental/test_default_qubit_2.py
@@ -172,9 +172,15 @@ def test_tracking_execute_and_derivatives(self):
 
         assert tracker.history == {
             "executions": [2, 4, 6],
-            "derivatives": [1, 2, 3, 4, 5, 6],
-            "derivative_batches": [1, 1, 1],
-            "execute_and_derivative_batches": [1, 1, 1],
+            "derivatives": [1, 2],
+            "derivative_batches": [1],
+            "execute_and_derivative_batches": [1],
+            "jvps": [3, 4],
+            "jvp_batches": [1],
+            "execute_and_jvp_batches": [1],
+            "vjps": [5, 6],
+            "vjp_batches": [1],
+            "execute_and_vjp_batches": [1],
             "resources": [Resources(num_wires=1)] * 12,
         }
 

From d29df4f5ba6b971152c2949aeb9a23ede0b61972 Mon Sep 17 00:00:00 2001
From: Edward Jiang <34989448+eddddddy@users.noreply.github.com>
Date: Mon, 24 Jul 2023 13:06:51 -0400
Subject: [PATCH 11/19] Update pennylane/devices/qubit/adjoint_jacobian.py

Co-authored-by: Christina Lee <christina@xanadu.ai>
---
 pennylane/devices/qubit/adjoint_jacobian.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index 6bca30831b8..e556d38ff45 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -65,7 +65,7 @@ def adjoint_jacobian(tape: QuantumTape, state=None):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = state if state is not None else get_final_state(tape)[0]
+    ket = state or get_final_state(tape)[0]
 
     n_obs = len(tape.observables)
     bras = np.empty([n_obs] + [2] * len(tape.wires), dtype=np.complex128)

From 32525c5c336942093a6dc779108a81df5bb602b3 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Mon, 24 Jul 2023 13:09:15 -0400
Subject: [PATCH 12/19] other places

---
 pennylane/devices/qubit/adjoint_jacobian.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index e556d38ff45..65dc1582d46 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -140,7 +140,7 @@ def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number], state=None):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = state if state is not None else get_final_state(tape)[0]
+    ket = state or get_final_state(tape)[0]
 
     n_obs = len(tape.observables)
     bras = np.empty([n_obs] + [2] * len(tape.wires), dtype=np.complex128)
@@ -214,7 +214,7 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = state if state is not None else get_final_state(tape)[0]
+    ket = state or get_final_state(tape)[0]
 
     obs = qml.dot(cotangents, tape.observables)
     bra = apply_operation(obs, ket)

From ee86b6f614a9acd62df2192f1a5f5345d7b5813d Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Mon, 24 Jul 2023 13:18:25 -0400
Subject: [PATCH 13/19] Add docs for simulate funcs

---
 pennylane/devices/qubit/simulate.py | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py
index f209f6c0a99..c9a9adbf4ca 100644
--- a/pennylane/devices/qubit/simulate.py
+++ b/pennylane/devices/qubit/simulate.py
@@ -26,7 +26,18 @@
 
 def get_final_state(circuit, debugger=None):
     """
-    Get the final state from executing the ops in the circuit
+    Get the final state that results from executing the given quantum script.
+
+    This is an internal function that will be called by the successor to ``default.qubit``.
+
+    Args:
+        circuit (.QuantumScript): The single circuit to simulate
+        debugger (._Debugger): The debugger to use
+
+    Returns:
+        Tuple[TensorLike, bool]: A tuple containing the final state of the quantum script and
+            whether the state has a batch dimension.
+
     """
     if set(circuit.wires) != set(range(circuit.num_wires)):
         wire_map = {w: i for i, w in enumerate(circuit.wires)}
@@ -54,7 +65,20 @@ def get_final_state(circuit, debugger=None):
 
 def measure_final_state(circuit, state, is_state_batched, rng=None) -> Result:
     """
-    TODO
+    Perform the measurements required by the circuit on the provided state.
+
+    This is an internal function that will be called by the successor to ``default.qubit``.
+
+    Args:
+        circuit (.QuantumScript): The single circuit to simulate
+        state (TensorLike): The state to perform measurement on
+        is_state_batched (bool): Whether the state has a batch dimension or not.
+        rng (Union[None, int, array_like[int], SeedSequence, BitGenerator, Generator]): A
+            seed-like parameter matching that of ``seed`` for ``numpy.random.default_rng``.
+            If no value is provided, a default RNG will be used.
+
+    Returns:
+        Tuple[TensorLike]: The measurement results
     """
     if set(circuit.wires) != set(range(circuit.num_wires)):
         wire_map = {w: i for i, w in enumerate(circuit.wires)}

From 1fdc1e5d7ad7b5739047cdc2ec78b1f6b1475e75 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Mon, 24 Jul 2023 13:32:04 -0400
Subject: [PATCH 14/19] revert suggestion

---
 pennylane/devices/qubit/adjoint_jacobian.py | 6 +++---
 tests/devices/qubit/test_simulate.py        | 2 +-
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index 65dc1582d46..6bca30831b8 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -65,7 +65,7 @@ def adjoint_jacobian(tape: QuantumTape, state=None):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = state or get_final_state(tape)[0]
+    ket = state if state is not None else get_final_state(tape)[0]
 
     n_obs = len(tape.observables)
     bras = np.empty([n_obs] + [2] * len(tape.wires), dtype=np.complex128)
@@ -140,7 +140,7 @@ def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number], state=None):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = state or get_final_state(tape)[0]
+    ket = state if state is not None else get_final_state(tape)[0]
 
     n_obs = len(tape.observables)
     bras = np.empty([n_obs] + [2] * len(tape.wires), dtype=np.complex128)
@@ -214,7 +214,7 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
         wire_map = {w: i for i, w in enumerate(tape.wires)}
         tape = qml.map_wires(tape, wire_map)
 
-    ket = state or get_final_state(tape)[0]
+    ket = state if state is not None else get_final_state(tape)[0]
 
     obs = qml.dot(cotangents, tape.observables)
     bra = apply_operation(obs, ket)
diff --git a/tests/devices/qubit/test_simulate.py b/tests/devices/qubit/test_simulate.py
index e6ef6e0c269..f70cd99ee8c 100644
--- a/tests/devices/qubit/test_simulate.py
+++ b/tests/devices/qubit/test_simulate.py
@@ -309,7 +309,7 @@ def test_broadcasted_op_sample(self):
         assert np.allclose(res[1], -np.cos(x), atol=0.05)
 
         state, is_state_batched = get_final_state(qs)
-        result = measure_final_state(qs, state, is_state_batched, rng=123)
+        res = measure_final_state(qs, state, is_state_batched, rng=123)
 
         expected_state = np.zeros((4, 2, 2))
         expected_state[:, 0, 1] = np.cos(x / 2)

From c65381c7e0cbc2125d9202131e6c9831ad047d3e Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Tue, 25 Jul 2023 10:25:39 -0400
Subject: [PATCH 15/19] Remove validation

---
 .../devices/experimental/default_qubit_2.py   | 56 +++++--------------
 .../experimental/test_default_qubit_2.py      | 37 ------------
 2 files changed, 13 insertions(+), 80 deletions(-)

diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
index dc72edfd1c4..7a7d20ef99a 100644
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ b/pennylane/devices/experimental/default_qubit_2.py
@@ -265,24 +265,19 @@ def compute_derivatives(
             self.tracker.update(derivative_batches=1, derivatives=len(circuits))
             self.tracker.record()
 
-        if execution_config.gradient_method == "adjoint":
-            max_workers = self._get_max_workers(execution_config)
-            if max_workers is None:
-                res = tuple(adjoint_jacobian(circuit) for circuit in circuits)
-            else:
-                vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
-                with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
-                    exec_map = executor.map(adjoint_jacobian, vanilla_circuits)
-                    res = tuple(circuit for circuit in exec_map)
-
-                # reset _rng to mimic serial behavior
-                self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
-
-            return res[0] if is_single_circuit else res
-
-        raise NotImplementedError(
-            f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
-        )
+        max_workers = self._get_max_workers(execution_config)
+        if max_workers is None:
+            res = tuple(adjoint_jacobian(circuit) for circuit in circuits)
+        else:
+            vanilla_circuits = [convert_to_numpy_parameters(c) for c in circuits]
+            with concurrent.futures.ProcessPoolExecutor(max_workers=max_workers) as executor:
+                exec_map = executor.map(adjoint_jacobian, vanilla_circuits)
+                res = tuple(circuit for circuit in exec_map)
+
+            # reset _rng to mimic serial behavior
+            self._rng = np.random.default_rng(self._rng.integers(2**31 - 1))
+
+        return res[0] if is_single_circuit else res
 
     def execute_and_compute_derivatives(
         self,
@@ -304,11 +299,6 @@ def execute_and_compute_derivatives(
             )
             self.tracker.record()
 
-        if execution_config.gradient_method != "adjoint":
-            raise NotImplementedError(
-                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
-            )
-
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
             results = tuple(
@@ -366,11 +356,6 @@ def compute_jvp(
             self.tracker.update(jvp_batches=1, jvps=len(circuits))
             self.tracker.record()
 
-        if execution_config.gradient_method != "adjoint":
-            raise NotImplementedError(
-                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
-            )
-
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
             res = tuple(adjoint_jvp(circuit, tans) for circuit, tans in zip(circuits, tangents))
@@ -404,11 +389,6 @@ def execute_and_compute_jvp(
             )
             self.tracker.record()
 
-        if execution_config.gradient_method != "adjoint":
-            raise NotImplementedError(
-                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
-            )
-
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
             results = tuple(
@@ -469,11 +449,6 @@ def compute_vjp(
             self.tracker.update(vjp_batches=1, vjps=len(circuits))
             self.tracker.record()
 
-        if execution_config.gradient_method != "adjoint":
-            raise NotImplementedError(
-                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
-            )
-
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
             res = tuple(adjoint_vjp(circuit, cots) for circuit, cots in zip(circuits, cotangents))
@@ -507,11 +482,6 @@ def execute_and_compute_vjp(
             )
             self.tracker.record()
 
-        if execution_config.gradient_method != "adjoint":
-            raise NotImplementedError(
-                f"{self.name} cannot compute derivatives via {execution_config.gradient_method}"
-            )
-
         max_workers = self._get_max_workers(execution_config)
         if max_workers is None:
             results = tuple(
diff --git a/tests/devices/experimental/test_default_qubit_2.py b/tests/devices/experimental/test_default_qubit_2.py
index 9c171902446..22d589bf6b1 100644
--- a/tests/devices/experimental/test_default_qubit_2.py
+++ b/tests/devices/experimental/test_default_qubit_2.py
@@ -29,43 +29,6 @@ def test_name():
     assert DefaultQubit2().name == "default.qubit.2"
 
 
-def test_no_jvp_functionality():
-    """Test that jvp is not supported on DefaultQubit2."""
-    dev = DefaultQubit2()
-
-    assert not dev.supports_jvp(ExecutionConfig())
-
-    with pytest.raises(NotImplementedError):
-        dev.compute_jvp(qml.tape.QuantumScript(), (10, 10))
-
-    with pytest.raises(NotImplementedError):
-        dev.execute_and_compute_jvp(qml.tape.QuantumScript(), (10, 10))
-
-
-def test_no_vjp_functionality():
-    """Test that vjp is not supported on DefaultQubit2."""
-    dev = DefaultQubit2()
-
-    assert not dev.supports_vjp(ExecutionConfig())
-
-    with pytest.raises(NotImplementedError):
-        dev.compute_vjp(qml.tape.QuantumScript(), (10.0, 10.0))
-
-    with pytest.raises(NotImplementedError):
-        dev.execute_and_compute_vjp(qml.tape.QuantumScript(), (10.0, 10.0))
-
-
-def test_no_device_derivatives():
-    """Test that DefaultQubit2 currently doesn't support device derivatives."""
-    dev = DefaultQubit2()
-
-    with pytest.raises(NotImplementedError):
-        dev.compute_derivatives(qml.tape.QuantumScript())
-
-    with pytest.raises(NotImplementedError):
-        dev.execute_and_compute_derivatives(qml.tape.QuantumScript())
-
-
 def test_debugger_attribute():
     """Test that DefaultQubit2 has a debugger attribute and that it is `None`"""
     # pylint: disable=protected-access

From 6906c71f54eb8e8e424d8e470eb0722ded63acc6 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Tue, 25 Jul 2023 12:20:36 -0400
Subject: [PATCH 16/19] Add validation for use_device_gradient

---
 .../devices/experimental/default_qubit_2.py    |  6 +++++-
 .../experimental/test_default_qubit_2.py       | 18 ++++++++++++++++++
 2 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
index 7a7d20ef99a..dd73f255464 100644
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ b/pennylane/devices/experimental/default_qubit_2.py
@@ -170,9 +170,13 @@ def supports_derivatives(
         ):
             return True
 
-        if execution_config.gradient_method == "adjoint":
+        if (
+            execution_config.gradient_method == "adjoint"
+            and execution_config.use_device_gradient in [None, True]
+        ):
             if circuit is None:
                 return True
+
             return isinstance(validate_and_expand_adjoint(circuit), QuantumScript)
 
         return False
diff --git a/tests/devices/experimental/test_default_qubit_2.py b/tests/devices/experimental/test_default_qubit_2.py
index 22d589bf6b1..c10ee5688cb 100644
--- a/tests/devices/experimental/test_default_qubit_2.py
+++ b/tests/devices/experimental/test_default_qubit_2.py
@@ -258,6 +258,24 @@ def test_supports_adjoint(self):
         assert dev.supports_jvp(config, qs) is True
         assert dev.supports_vjp(config, qs) is True
 
+        config = ExecutionConfig(gradient_method="adjoint", use_device_gradient=True)
+        assert dev.supports_derivatives(config) is True
+        assert dev.supports_jvp(config) is True
+        assert dev.supports_vjp(config) is True
+
+        assert dev.supports_derivatives(config, qs) is True
+        assert dev.supports_jvp(config, qs) is True
+        assert dev.supports_vjp(config, qs) is True
+
+        config = ExecutionConfig(gradient_method="adjoint", use_device_gradient=False)
+        assert dev.supports_derivatives(config) is False
+        assert dev.supports_jvp(config) is False
+        assert dev.supports_vjp(config) is False
+
+        assert dev.supports_derivatives(config, qs) is False
+        assert dev.supports_jvp(config, qs) is False
+        assert dev.supports_vjp(config, qs) is False
+
     def test_doesnt_support_adjoint_with_invalid_tape(self):
         """Tests that DefaultQubit2 does not support adjoint differentiation with invalid circuits."""
         dev = DefaultQubit2()

From 45375d1d3e32001b890943d8dd47a86ea738b504 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Tue, 25 Jul 2023 14:27:30 -0400
Subject: [PATCH 17/19] remove None case

---
 pennylane/devices/experimental/default_qubit_2.py  |  5 +----
 tests/devices/experimental/test_default_qubit_2.py | 11 +----------
 2 files changed, 2 insertions(+), 14 deletions(-)

diff --git a/pennylane/devices/experimental/default_qubit_2.py b/pennylane/devices/experimental/default_qubit_2.py
index dd73f255464..8bd3fb7dd1a 100644
--- a/pennylane/devices/experimental/default_qubit_2.py
+++ b/pennylane/devices/experimental/default_qubit_2.py
@@ -170,10 +170,7 @@ def supports_derivatives(
         ):
             return True
 
-        if (
-            execution_config.gradient_method == "adjoint"
-            and execution_config.use_device_gradient in [None, True]
-        ):
+        if execution_config.gradient_method == "adjoint" and execution_config.use_device_gradient:
             if circuit is None:
                 return True
 
diff --git a/tests/devices/experimental/test_default_qubit_2.py b/tests/devices/experimental/test_default_qubit_2.py
index c10ee5688cb..8de95f5d994 100644
--- a/tests/devices/experimental/test_default_qubit_2.py
+++ b/tests/devices/experimental/test_default_qubit_2.py
@@ -248,21 +248,12 @@ def test_supports_backprop(self):
     def test_supports_adjoint(self):
         """Test that DefaultQubit2 says that it supports adjoint differentiation."""
         dev = DefaultQubit2()
-        config = ExecutionConfig(gradient_method="adjoint")
-        assert dev.supports_derivatives(config) is True
-        assert dev.supports_jvp(config) is True
-        assert dev.supports_vjp(config) is True
-
-        qs = qml.tape.QuantumScript([], [qml.expval(qml.PauliZ(0))])
-        assert dev.supports_derivatives(config, qs) is True
-        assert dev.supports_jvp(config, qs) is True
-        assert dev.supports_vjp(config, qs) is True
-
         config = ExecutionConfig(gradient_method="adjoint", use_device_gradient=True)
         assert dev.supports_derivatives(config) is True
         assert dev.supports_jvp(config) is True
         assert dev.supports_vjp(config) is True
 
+        qs = qml.tape.QuantumScript([], [qml.expval(qml.PauliZ(0))])
         assert dev.supports_derivatives(config, qs) is True
         assert dev.supports_jvp(config, qs) is True
         assert dev.supports_vjp(config, qs) is True

From 0e21a13d80054ba9eee964e5775dc68148948c4e Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Wed, 26 Jul 2023 10:33:57 -0400
Subject: [PATCH 18/19] use device gradient in validation

---
 pennylane/qnode.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/pennylane/qnode.py b/pennylane/qnode.py
index d5bc695ae56..3fc7218e754 100644
--- a/pennylane/qnode.py
+++ b/pennylane/qnode.py
@@ -762,7 +762,9 @@ def _validate_adjoint_method(device):
         # cannot be done here since we don't yet know the composition of the circuit.
 
         if isinstance(device, qml.devices.experimental.Device):
-            config = qml.devices.experimental.ExecutionConfig(gradient_method="adjoint")
+            config = qml.devices.experimental.ExecutionConfig(
+                gradient_method="adjoint", use_device_gradient=True
+            )
             if device.supports_derivatives(config):
                 return "adjoint", {}, device
             raise ValueError(f"The {device} device does not support adjoint differentiation.")

From 3b98dac5289570529c091d248d2dd2b07cd38544 Mon Sep 17 00:00:00 2001
From: Edward Jiang <eddy3.14159@gmail.com>
Date: Fri, 28 Jul 2023 10:28:48 -0400
Subject: [PATCH 19/19] changelog

---
 doc/releases/changelog-dev.md | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 5f2db9d7826..07afc84ebfb 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -88,6 +88,9 @@
 * When given a callable, `qml.ctrl` now does its custom pre-processing on all queued operators from the callable.
   [(#4370)](https://github.com/PennyLaneAI/pennylane/pull/4370)
 
+* The experimental `DefaultQubit2` device now supports computing VJPs and JVPs using the adjoint method.
+  [(#4374)](https://github.com/PennyLaneAI/pennylane/pull/4374)
+
 <h3>Breaking changes 💔</h3>
 
 * `Operator.expand` now uses the output of `Operator.decomposition` instead of what it queues.