From a5ad41c1fc76fbf1ecdcc8a3989c1f76741e2be2 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 31 Jul 2024 11:47:16 +0300
Subject: [PATCH 01/33] Adding code for batch circuit submit.

---
 pennylane_ionq/device.py | 322 +++++++++++++++++++++++++++++++++++++--
 1 file changed, 309 insertions(+), 13 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 4f46216..6517861 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -14,6 +14,9 @@
 """
 This module contains the device class for constructing IonQ devices for PennyLane.
 """
+import inspect
+import logging
+from typing import Union
 import warnings
 from time import sleep
 
@@ -21,9 +24,32 @@
 
 from pennylane import QubitDevice, DeviceError
 
+from pennylane.measurements import (
+    ClassicalShadowMP,
+    CountsMP,
+    ExpectationMP,
+    MeasurementTransform,
+    MeasurementValue,
+    MutualInfoMP,
+    ProbabilityMP,
+    SampleMeasurement,
+    SampleMP,
+    ShadowExpvalMP,
+    Shots,
+    StateMeasurement,
+    StateMP,
+    VarianceMP,
+    VnEntropyMP,
+)
+from pennylane.resource import Resources
+from pennylane.tape import QuantumTape
+
 from .api_client import Job, JobExecutionError
 from ._version import __version__
 
+logger = logging.getLogger(__name__)
+logger.addHandler(logging.NullHandler())
+
 _qis_operation_map = {
     # native PennyLane operations also native to IonQ
     "PauliX": "x",
@@ -118,26 +144,39 @@ def __init__(
             raise ValueError("The ionq device does not support analytic expectation values.")
 
         super().__init__(wires=wires, shots=shots)
+        self._current_circuit_index = None
         self.target = target
         self.api_key = api_key
         self.gateset = gateset
         self.error_mitigation = error_mitigation
         self.sharpen = sharpen
         self._operation_map = _GATESET_OPS[gateset]
+        self.histogram = None
+        self.histograms = None
         self.reset()
 
-    def reset(self):
+    def reset(self, circuits_array_length=0):
         """Reset the device"""
-        self._prob_array = None
+        self._current_circuit_index = None
         self.histogram = None
-        self.circuit = {
-            "format": "ionq.circuit.v0",
-            "qubits": self.num_wires,
-            "circuit": [],
-            "gateset": self.gateset,
-        }
+        self.histograms = None
+        self._prob_array = None
+        if circuits_array_length == 0:
+            self.input = {
+                "format": "ionq.circuit.v0",
+                "qubits": self.num_wires,
+                "circuit": [],
+                "gateset": self.gateset,
+            }
+        else:
+            self.input = {
+                "format": "ionq.circuit.v0",
+                "qubits": self.num_wires,
+                "circuits": [{"circuit": []} for _ in range(circuits_array_length)],
+                "gateset": self.gateset,
+            }
         self.job = {
-            "input": self.circuit,
+            "input": self.input,
             "target": self.target,
             "shots": self.shots,
         }
@@ -151,6 +190,116 @@ def reset(self):
                 stacklevel=2,
             )
 
+    def batch_execute(self, circuits):
+        """Execute a batch of quantum circuits on the device.
+
+        The circuits are represented by tapes, and they are executed one-by-one using the
+        device's ``execute`` method. The results are collected in a list.
+
+        Args:
+            circuits (list[~.tape.QuantumTape]): circuits to execute on the device
+
+        Returns:
+            list[array[float]]: list of measured value(s)
+        """
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(
+                """Entry with args=(circuits=%s) called by=%s""",
+                circuits,
+                "::L".join(
+                    str(i)
+                    for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]
+                ),
+            )
+
+        self.reset(circuits_array_length=len(circuits))
+
+        for circuit_index, circuit in enumerate(circuits):
+            self.check_validity(circuit.operations, circuit.observables)
+            self.batch_apply(
+                circuit.operations,
+                rotations=self._get_diagonalizing_gates(circuit),
+                circuit_index=circuit_index,
+            )
+
+        self._submit_job()
+
+        results = []
+        for circuit_index, circuit in enumerate(circuits):
+            self._current_circuit_index = circuit_index
+            sample_type = (SampleMP, CountsMP, ClassicalShadowMP, ShadowExpvalMP)
+            if self.shots is not None or any(
+                isinstance(m, sample_type) for m in circuit.measurements
+            ):
+                self._samples = self.generate_samples()
+
+            # compute the required statistics
+            if self._shot_vector is not None:
+                result = self.shot_vec_statistics(circuit)
+            else:
+                result = self.statistics(circuit)
+                single_measurement = len(circuit.measurements) == 1
+
+                result = result[0] if single_measurement else tuple(result)
+
+            self._current_circuit_index = None
+            self._samples = None
+            results.append(result)
+
+        # increment counter for number of executions of qubit device
+        self._num_executions += 1
+
+        if self.tracker.active:
+            for circuit in circuits:
+                shots_from_dev = (
+                    self._shots if not self.shot_vector else self._raw_shot_sequence
+                )
+                tape_resources = circuit.specs["resources"]
+
+                resources = Resources(  # temporary until shots get updated on tape !
+                    tape_resources.num_wires,
+                    tape_resources.num_gates,
+                    tape_resources.gate_types,
+                    tape_resources.gate_sizes,
+                    tape_resources.depth,
+                    Shots(shots_from_dev),
+                )
+                self.tracker.update(
+                    executions=1,
+                    shots=self._shots,
+                    results=results,
+                    resources=resources,
+                )
+
+            self.tracker.update(batches=1, batch_len=len(circuits))
+            self.tracker.record()
+
+        return results
+
+    def batch_apply(self, operations, circuit_index, **kwargs):
+
+        rotations = kwargs.pop("rotations", [])
+
+        if len(operations) == 0 and len(rotations) == 0:
+            warnings.warn(
+                "Circuit is empty. Empty circuits return failures. Submitting anyway."
+            )
+
+        for i, operation in enumerate(operations):
+            if i > 0 and operation.name in {
+                "BasisState",
+                "QubitStateVector",
+                "StatePrep",
+            }:
+                raise DeviceError(
+                    f"The operation {operation.name} is only supported at the beginning of a circuit."
+                )
+            self._apply_operation(operation, circuit_index)
+
+        # diagonalize observables
+        for operation in rotations:
+            self._apply_operation(operation, circuit_index)
+
     @property
     def operations(self):
         """Get the supported set of operations.
@@ -184,7 +333,7 @@ def apply(self, operations, **kwargs):
 
         self._submit_job()
 
-    def _apply_operation(self, operation):
+    def _apply_operation(self, operation, circuit_index=None):
         name = operation.name
         wires = self.map_wires(operation.wires).tolist()
         gate = {"gate": self._operation_map[name]}
@@ -202,13 +351,18 @@ def _apply_operation(self, operation):
 
         if self.gateset == "native":
             if len(par) > 1:
-                gate["phases"] = [float(v) for v in par]
+                gate["phases"] = [float(v) for v in par[:2]]
+                if len(par) > 2:
+                    gate["angle"] = float(par[2])
             else:
                 gate["phase"] = float(par[0])
         elif par:
             gate["rotation"] = float(par[0])
 
-        self.circuit["circuit"].append(gate)
+        if "circuit" in self.input:
+            self.input["circuit"].append(gate)
+        else:
+            self.input["circuits"][circuit_index]["circuit"].append(gate)
 
     def _submit_job(self):
         job = Job(api_key=self.api_key)
@@ -232,13 +386,24 @@ def _submit_job(self):
         # state (as a base-10 integer string) to the probability
         # as a floating point value between 0 and 1.
         # e.g., {"0": 0.413, "9": 0.111, "17": 0.476}
-        self.histogram = job.data.value
+        some_inner_value = next(iter(job.data.value.values()))
+        if isinstance(some_inner_value, dict):
+            self.histograms = []
+            for key in job.data.value.keys():
+                self.histograms.append(job.data.value[key])
+        else:
+            self.histograms = []
+            self.histograms.append(job.data.value)
+            self.histogram = job.data.value
 
     @property
     def prob(self):
         """None or array[float]: Array of computational basis state probabilities. If
         no job has been submitted, returns ``None``.
         """
+        if self._current_circuit_index is not None:
+            self.histogram = self.histograms[self._current_circuit_index]
+
         if self.histogram is None:
             return None
 
@@ -259,6 +424,9 @@ def prob(self):
             norm = sum(histogram_values)
             self._prob_array[idx] = np.fromiter(histogram_values, float) / norm
 
+        if self._current_circuit_index is not None:
+            self.histogram = None
+
         return self._prob_array
 
     def probability(self, wires=None, shot_range=None, bin_size=None):
@@ -269,6 +437,134 @@ def probability(self, wires=None, shot_range=None, bin_size=None):
 
         return self.estimate_probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
 
+    def expval(self, observable, shot_range=None, bin_size=None):
+        """
+        Returns the expectation value of a Hamiltonian observable.
+        Overwrites the method in QubitDevice class to accomodate batch submission of circuits.
+        When invoking this method after submitting circuits using any other method than 
+        QubitDevice.execute() method, self_current_circuit_index and self._samples should be 
+        initialized using the circuit which is targeted out of  one or potentially several circuits submitted 
+        in a batch. This is likely not very good design but the alternative would be to create
+        huge amounts of code duplication with respect to the Pennylane base code.
+
+        Args:
+            observable (~.Observable): a PennyLane observable
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+
+        """
+        try:
+            return super().expval(observable, shot_range, bin_size)
+        except TypeError as e:
+            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index and self._samples should be initialized using the circuit which is targeted out of  one or potentially several circuits submitted in a batch.", *e.args)
+            raise
+
+    def estimate_probability(
+        self, wires=None, shot_range=None, bin_size=None
+    ):
+        """Return the estimated probability of each computational basis state
+        using the generated samples. Overwrites the method in QubitDevice class
+        to accomodate batch submission of circuits. When invoking this method after 
+        submitting circuits using any other method than QubitDevice.execute() 
+        method, self_current_circuit_index and self._samples should be initialized to the 
+        circuit which is targeted out of  one or potentially several circuits submitted in a batch.
+        This is likely not very good design but the alternative would be to create
+        huge amounts of code duplication with respect to the Pennylane base code.
+
+        Args:
+            wires (Iterable[Number, str], Number, str, Wires): wires to calculate
+                marginal probabilities for. Wires not provided are traced out of  one or the system.
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+
+        Returns:
+            array[float]: list of the probabilities
+        """
+        try:
+            return super().estimate_probability(wires, shot_range, bin_size)
+        except TypeError as e:
+            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index and self._samples should be initialized using the circuit which is targeted out of  one or potentially several circuits submitted in a batch.", *e.args)
+            raise
+
+    def sample(
+        self,
+        observable,
+        shot_range=None,
+        bin_size=None,
+        counts=False,
+    ):
+        """Return samples of an observable. Overwrites the method in QubitDevice class
+        to accomodate batch submission of circuits. When invoking this method after submitting 
+        circuits using any other method than QubitDevice.execute() method, 
+        self_current_circuit_index and self._samples should be initialized using the circuit which 
+        is targeted out of  one or potentially several circuits submitted in a batch.
+        This is likely not very good design but the alternative would be to create
+        huge amounts of code duplication with respect to the Pennylane base code.
+
+        Args:
+            observable (Observable): the observable to sample
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+            counts (bool): whether counts (``True``) or raw samples (``False``)
+                should be returned
+            circuit_index (int): index of circuit in case of batch circuit submission
+
+        Raises:
+            EigvalsUndefinedError: if no information is available about the
+                eigenvalues of the observable
+
+        Returns:
+            Union[array[float], dict, list[dict]]: samples in an array of
+            dimension ``(shots,)`` or counts
+        """
+        try:
+            return super().sample(observable, shot_range, bin_size, counts)
+        except TypeError as e:
+            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method self_current_circuit_index and self._samples should be initialized using the circuit which is targeted out of  one or potentially several circuits submitted in a batch.", *e.args)
+            raise
+
+    def _measure(
+        self,
+        measurement: Union[SampleMeasurement, StateMeasurement],
+        shot_range=None,
+        bin_size=None,
+    ):
+        """Compute the corresponding measurement process depending on ``shots`` and the measurement
+        type. Overwrites the method in QubitDevice class to accomodate batch submission of circuits.
+        When invoking this method after submitting circuits using any other method than QubitDevice.execute() 
+        method, self_current_circuit_index and self._samples should be initialized using the circuit which 
+        is targeted out of  one or potentially several circuits submitted in a batch. This is likely not very good 
+        design but the alternative would be to create huge amounts of code duplication with respect to the 
+        Pennylane base code.
+
+        Args:
+            measurement (Union[SampleMeasurement, StateMeasurement]): measurement process
+            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
+                to use. If not specified, all samples are used.
+            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
+                returns the measurement statistic separately over each bin. If not
+                provided, the entire shot range is treated as a single bin.
+
+        Raises:
+            ValueError: if the measurement cannot be computed
+
+        Returns:
+            Union[float, dict, list[float]]: result of the measurement
+        """
+        try:
+            return super()._measure(measurement, shot_range, bin_size)
+        except TypeError as e:
+            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index and self._samples should be initialized using the circuit which is targeted out of  one or potentially several circuits submitted in a batch. ", *e.args)
+            raise
 
 class SimulatorDevice(IonQDevice):
     r"""Simulator device for IonQ.

From 311c09fbfbee30feaabbb2efad6078fa8bd3cacd Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 31 Jul 2024 11:58:14 +0300
Subject: [PATCH 02/33] Remove unused imports.

---
 pennylane_ionq/device.py | 9 ---------
 1 file changed, 9 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 6517861..cdf93e4 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -27,22 +27,13 @@
 from pennylane.measurements import (
     ClassicalShadowMP,
     CountsMP,
-    ExpectationMP,
-    MeasurementTransform,
-    MeasurementValue,
-    MutualInfoMP,
-    ProbabilityMP,
     SampleMeasurement,
     SampleMP,
     ShadowExpvalMP,
     Shots,
     StateMeasurement,
-    StateMP,
-    VarianceMP,
-    VnEntropyMP,
 )
 from pennylane.resource import Resources
-from pennylane.tape import QuantumTape
 
 from .api_client import Job, JobExecutionError
 from ._version import __version__

From bad9168282bed26ecc599af6444fa11d33224d62 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 31 Jul 2024 11:59:18 +0300
Subject: [PATCH 03/33] Update version number.

---
 pennylane_ionq/_version.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane_ionq/_version.py b/pennylane_ionq/_version.py
index c32bfdf..54057c1 100644
--- a/pennylane_ionq/_version.py
+++ b/pennylane_ionq/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.36.0"
+__version__ = "0.37.0"

From b22e88fd78acf7dcdef154e344c275dba24e99ac Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 31 Jul 2024 12:24:31 +0300
Subject: [PATCH 04/33] Initialize self._samples in overloaded methods.

---
 pennylane_ionq/device.py | 56 ++++++++++++++++++++++++++--------------
 1 file changed, 36 insertions(+), 20 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index cdf93e4..0f49199 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -433,9 +433,9 @@ def expval(self, observable, shot_range=None, bin_size=None):
         Returns the expectation value of a Hamiltonian observable.
         Overwrites the method in QubitDevice class to accomodate batch submission of circuits.
         When invoking this method after submitting circuits using any other method than 
-        QubitDevice.execute() method, self_current_circuit_index and self._samples should be 
-        initialized using the circuit which is targeted out of  one or potentially several circuits submitted 
-        in a batch. This is likely not very good design but the alternative would be to create
+        QubitDevice.execute() method, self_current_circuit_index should be 
+        initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted 
+        in a batch. This is not very good design but the alternative would be to create
         huge amounts of code duplication with respect to the Pennylane base code.
 
         Args:
@@ -448,9 +448,13 @@ def expval(self, observable, shot_range=None, bin_size=None):
 
         """
         try:
-            return super().expval(observable, shot_range, bin_size)
+            self._samples = self.generate_samples()
+            result = super().expval(observable, shot_range, bin_size)
+            self._samples = None
+            return result
         except TypeError as e:
-            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index and self._samples should be initialized using the circuit which is targeted out of  one or potentially several circuits submitted in a batch.", *e.args)
+            self._samples = None
+            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index should be initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
             raise
 
     def estimate_probability(
@@ -460,14 +464,14 @@ def estimate_probability(
         using the generated samples. Overwrites the method in QubitDevice class
         to accomodate batch submission of circuits. When invoking this method after 
         submitting circuits using any other method than QubitDevice.execute() 
-        method, self_current_circuit_index and self._samples should be initialized to the 
-        circuit which is targeted out of  one or potentially several circuits submitted in a batch.
-        This is likely not very good design but the alternative would be to create
+        method, self_current_circuit_index should be initialized to the 
+        circuit which is targeted out of one or potentially several circuits submitted in a batch.
+        This is not very good design but the alternative would be to create
         huge amounts of code duplication with respect to the Pennylane base code.
 
         Args:
             wires (Iterable[Number, str], Number, str, Wires): wires to calculate
-                marginal probabilities for. Wires not provided are traced out of  one or the system.
+                marginal probabilities for. Wires not provided are traced out of one or the system.
             shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
                 to use. If not specified, all samples are used.
             bin_size (int): Divides the shot range into bins of size ``bin_size``, and
@@ -478,9 +482,13 @@ def estimate_probability(
             array[float]: list of the probabilities
         """
         try:
-            return super().estimate_probability(wires, shot_range, bin_size)
+            self._samples = self.generate_samples()
+            result = super().estimate_probability(wires, shot_range, bin_size)
+            self._samples = None
+            return result
         except TypeError as e:
-            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index and self._samples should be initialized using the circuit which is targeted out of  one or potentially several circuits submitted in a batch.", *e.args)
+            self._samples = None
+            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index should be initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
             raise
 
     def sample(
@@ -493,9 +501,9 @@ def sample(
         """Return samples of an observable. Overwrites the method in QubitDevice class
         to accomodate batch submission of circuits. When invoking this method after submitting 
         circuits using any other method than QubitDevice.execute() method, 
-        self_current_circuit_index and self._samples should be initialized using the circuit which 
-        is targeted out of  one or potentially several circuits submitted in a batch.
-        This is likely not very good design but the alternative would be to create
+        self_current_circuit_index should be initialized pointing to the circuit which 
+        is targeted out of one or potentially several circuits submitted in a batch.
+        This is not very good design but the alternative would be to create
         huge amounts of code duplication with respect to the Pennylane base code.
 
         Args:
@@ -518,9 +526,13 @@ def sample(
             dimension ``(shots,)`` or counts
         """
         try:
-            return super().sample(observable, shot_range, bin_size, counts)
+            self._samples = self.generate_samples()
+            result = super().sample(observable, shot_range, bin_size, counts)
+            self._samples = None
+            return result
         except TypeError as e:
-            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method self_current_circuit_index and self._samples should be initialized using the circuit which is targeted out of  one or potentially several circuits submitted in a batch.", *e.args)
+            self._samples = None
+            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method self_current_circuit_index should be initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
             raise
 
     def _measure(
@@ -532,8 +544,8 @@ def _measure(
         """Compute the corresponding measurement process depending on ``shots`` and the measurement
         type. Overwrites the method in QubitDevice class to accomodate batch submission of circuits.
         When invoking this method after submitting circuits using any other method than QubitDevice.execute() 
-        method, self_current_circuit_index and self._samples should be initialized using the circuit which 
-        is targeted out of  one or potentially several circuits submitted in a batch. This is likely not very good 
+        method, self_current_circuit_index should be initialized pointing to the circuit which 
+        is targeted out of one or potentially several circuits submitted in a batch. This is not very good 
         design but the alternative would be to create huge amounts of code duplication with respect to the 
         Pennylane base code.
 
@@ -552,9 +564,13 @@ def _measure(
             Union[float, dict, list[float]]: result of the measurement
         """
         try:
-            return super()._measure(measurement, shot_range, bin_size)
+            self._samples = self.generate_samples()
+            result = super()._measure(measurement, shot_range, bin_size)
+            self._samples = None
+            return result
         except TypeError as e:
-            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index and self._samples should be initialized using the circuit which is targeted out of  one or potentially several circuits submitted in a batch. ", *e.args)
+            self._samples = None
+            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index should be initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted in a batch. ", *e.args)
             raise
 
 class SimulatorDevice(IonQDevice):

From d7e16a9acb2d3a6093dd5835555c3a3b2bf09746 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 31 Jul 2024 16:31:34 +0300
Subject: [PATCH 05/33] Various fixes.

---
 pennylane_ionq/device.py | 78 ++++++++++++----------------------------
 1 file changed, 23 insertions(+), 55 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 0f49199..46b025d 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -181,6 +181,9 @@ def reset(self, circuits_array_length=0):
                 stacklevel=2,
             )
 
+    def set_current_circuit_index(self, circuit_index):
+        self._current_circuit_index = circuit_index
+
     def batch_execute(self, circuits):
         """Execute a batch of quantum circuits on the device.
 
@@ -217,7 +220,7 @@ def batch_execute(self, circuits):
 
         results = []
         for circuit_index, circuit in enumerate(circuits):
-            self._current_circuit_index = circuit_index
+            self.set_current_circuit_index(circuit_index)
             sample_type = (SampleMP, CountsMP, ClassicalShadowMP, ShadowExpvalMP)
             if self.shots is not None or any(
                 isinstance(m, sample_type) for m in circuit.measurements
@@ -233,7 +236,7 @@ def batch_execute(self, circuits):
 
                 result = result[0] if single_measurement else tuple(result)
 
-            self._current_circuit_index = None
+            self.set_current_circuit_index(None)
             self._samples = None
             results.append(result)
 
@@ -420,43 +423,6 @@ def prob(self):
 
         return self._prob_array
 
-    def probability(self, wires=None, shot_range=None, bin_size=None):
-        wires = wires or self.wires
-
-        if shot_range is None and bin_size is None:
-            return self.marginal_prob(self.prob, wires)
-
-        return self.estimate_probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
-
-    def expval(self, observable, shot_range=None, bin_size=None):
-        """
-        Returns the expectation value of a Hamiltonian observable.
-        Overwrites the method in QubitDevice class to accomodate batch submission of circuits.
-        When invoking this method after submitting circuits using any other method than 
-        QubitDevice.execute() method, self_current_circuit_index should be 
-        initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted 
-        in a batch. This is not very good design but the alternative would be to create
-        huge amounts of code duplication with respect to the Pennylane base code.
-
-        Args:
-            observable (~.Observable): a PennyLane observable
-            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
-                to use. If not specified, all samples are used.
-            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
-                returns the measurement statistic separately over each bin. If not
-                provided, the entire shot range is treated as a single bin.
-
-        """
-        try:
-            self._samples = self.generate_samples()
-            result = super().expval(observable, shot_range, bin_size)
-            self._samples = None
-            return result
-        except TypeError as e:
-            self._samples = None
-            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index should be initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
-            raise
-
     def estimate_probability(
         self, wires=None, shot_range=None, bin_size=None
     ):
@@ -464,10 +430,11 @@ def estimate_probability(
         using the generated samples. Overwrites the method in QubitDevice class
         to accomodate batch submission of circuits. When invoking this method after 
         submitting circuits using any other method than QubitDevice.execute() 
-        method, self_current_circuit_index should be initialized to the 
-        circuit which is targeted out of one or potentially several circuits submitted in a batch.
-        This is not very good design but the alternative would be to create
-        huge amounts of code duplication with respect to the Pennylane base code.
+        method, current circuit index should be set via set_current_circuit_index() 
+        method to the circuit which is targeted out of one or potentially several 
+        circuits submitted in a batch. This is not very good design but the 
+        alternative would be to create huge amounts of code duplication with respect
+        to the Pennylane base code.
 
         Args:
             wires (Iterable[Number, str], Number, str, Wires): wires to calculate
@@ -488,7 +455,7 @@ def estimate_probability(
             return result
         except TypeError as e:
             self._samples = None
-            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index should be initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
+            e.args = ("When invoking QubitDevice estimate_probability() method after submitting circuits using any other method than QubitDevice.execute() method, current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
             raise
 
     def sample(
@@ -500,11 +467,11 @@ def sample(
     ):
         """Return samples of an observable. Overwrites the method in QubitDevice class
         to accomodate batch submission of circuits. When invoking this method after submitting 
-        circuits using any other method than QubitDevice.execute() method, 
-        self_current_circuit_index should be initialized pointing to the circuit which 
-        is targeted out of one or potentially several circuits submitted in a batch.
-        This is not very good design but the alternative would be to create
-        huge amounts of code duplication with respect to the Pennylane base code.
+        circuits using any other method than QubitDevice.execute() method, current circuit 
+        index should be set via set_current_circuit_index() method pointing to index of the 
+        circuit which is targeted out of one or potentially several circuits submitted in a batch.
+        This is not very good design but the alternative would be to create huge amounts of 
+        code duplication with respect to the Pennylane base code.
 
         Args:
             observable (Observable): the observable to sample
@@ -532,7 +499,7 @@ def sample(
             return result
         except TypeError as e:
             self._samples = None
-            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method self_current_circuit_index should be initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
+            e.args = ("When invoking QubitDevice sample() method after submitting circuits using any other method than QubitDevice.execute() method current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
             raise
 
     def _measure(
@@ -544,10 +511,10 @@ def _measure(
         """Compute the corresponding measurement process depending on ``shots`` and the measurement
         type. Overwrites the method in QubitDevice class to accomodate batch submission of circuits.
         When invoking this method after submitting circuits using any other method than QubitDevice.execute() 
-        method, self_current_circuit_index should be initialized pointing to the circuit which 
-        is targeted out of one or potentially several circuits submitted in a batch. This is not very good 
-        design but the alternative would be to create huge amounts of code duplication with respect to the 
-        Pennylane base code.
+        method, current circuit index should be set via set_current_circuit_index() method pointing to index 
+        of the circuit which is targeted out of one or potentially several circuits submitted in a batch. 
+        This is not very good design but the alternative would be to create huge amounts of code duplication 
+        with respect to the Pennylane base code.
 
         Args:
             measurement (Union[SampleMeasurement, StateMeasurement]): measurement process
@@ -570,9 +537,10 @@ def _measure(
             return result
         except TypeError as e:
             self._samples = None
-            e.args = ("When invoking this method after submitting circuits using any other method than QubitDevice.execute() method, self_current_circuit_index should be initialized pointing to the circuit which is targeted out of one or potentially several circuits submitted in a batch. ", *e.args)
+            e.args = ("When invoking QubitDevice _measure() method after submitting circuits using any other method than QubitDevice.execute() method, current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch. ", *e.args)
             raise
 
+
 class SimulatorDevice(IonQDevice):
     r"""Simulator device for IonQ.
 

From dede94e75d132fd356d17f761499430d5976ec8c Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Tue, 6 Aug 2024 14:42:17 +0300
Subject: [PATCH 06/33] Running existing unit tests from previous
 implementation and fixing bugs.

---
 pennylane_ionq/device.py |  38 +++---
 tests/test_device.py     | 255 ++++++++++++++++++++++++++++++++++++---
 2 files changed, 263 insertions(+), 30 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 46b025d..379a6a3 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -151,7 +151,6 @@ def reset(self, circuits_array_length=0):
         self._current_circuit_index = None
         self.histogram = None
         self.histograms = None
-        self._prob_array = None
         if circuits_array_length == 0:
             self.input = {
                 "format": "ionq.circuit.v0",
@@ -401,27 +400,26 @@ def prob(self):
         if self.histogram is None:
             return None
 
-        if self._prob_array is None:
-            # The IonQ API returns basis states using little-endian ordering.
-            # Here, we rearrange the states to match the big-endian ordering
-            # expected by PennyLane.
-            basis_states = (
-                int(bin(int(k))[2:].rjust(self.num_wires, "0")[::-1], 2) for k in self.histogram
-            )
-            idx = np.fromiter(basis_states, dtype=int)
+        # The IonQ API returns basis states using little-endian ordering.
+        # Here, we rearrange the states to match the big-endian ordering
+        # expected by PennyLane.
+        basis_states = (
+            int(bin(int(k))[2:].rjust(self.num_wires, "0")[::-1], 2) for k in self.histogram
+        )
+        idx = np.fromiter(basis_states, dtype=int)
 
-            # convert the sparse probs into a probability array
-            self._prob_array = np.zeros([2**self.num_wires])
+        # convert the sparse probs into a probability array
+        prob_array = np.zeros([2**self.num_wires])
 
-            # histogram values don't always perfectly sum to exactly one
-            histogram_values = self.histogram.values()
-            norm = sum(histogram_values)
-            self._prob_array[idx] = np.fromiter(histogram_values, float) / norm
+        # histogram values don't always perfectly sum to exactly one
+        histogram_values = self.histogram.values()
+        norm = sum(histogram_values)
+        prob_array[idx] = np.fromiter(histogram_values, float) / norm
 
         if self._current_circuit_index is not None:
             self.histogram = None
 
-        return self._prob_array
+        return prob_array
 
     def estimate_probability(
         self, wires=None, shot_range=None, bin_size=None
@@ -540,6 +538,14 @@ def _measure(
             e.args = ("When invoking QubitDevice _measure() method after submitting circuits using any other method than QubitDevice.execute() method, current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch. ", *e.args)
             raise
 
+    def probability(self, wires=None, shot_range=None, bin_size=None):
+        wires = wires or self.wires
+
+        if shot_range is None and bin_size is None:
+            return self.marginal_prob(self.prob, wires)
+
+        return self.estimate_probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
+
 
 class SimulatorDevice(IonQDevice):
     r"""Simulator device for IonQ.
diff --git a/tests/test_device.py b/tests/test_device.py
index c8842fd..95e2247 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -17,12 +17,12 @@
 import pennylane as qml
 import pytest
 import requests
-from unittest.mock import PropertyMock, patch
 
 from conftest import shortnames
-from pennylane_ionq.api_client import JobExecutionError, ResourceManager, Job, Field
-from pennylane_ionq.device import QPUDevice, IonQDevice
+from pennylane_ionq.api_client import JobExecutionError, ResourceManager, Job
+from pennylane_ionq.device import QPUDevice, IonQDevice, SimulatorDevice
 from pennylane_ionq.ops import GPI, GPI2, MS
+from pennylane.measurements import SampleMeasurement
 
 FAKE_API_KEY = "ABC123"
 
@@ -230,18 +230,18 @@ def test_prob_no_results(self, d):
         dev = qml.device(d, wires=1, shots=1)
         assert dev.prob is None
 
-    def test_probability(self):
-        """Test that device.probability works."""
-        dev = qml.device("ionq.simulator", wires=2)
-        dev._samples = np.array([[1, 1], [1, 1], [0, 0], [0, 0]])
-        assert np.array_equal(dev.probability(shot_range=(0, 2)), [0, 0, 0, 1])
+    # def test_probability(self):
+    #     """Test that device.probability works."""
+    #     dev = qml.device("ionq.simulator", wires=2)
+    #     dev._samples = np.array([[1, 1], [1, 1], [0, 0], [0, 0]])
+    #     assert np.array_equal(dev.probability(shot_range=(0, 2)), [0, 0, 0, 1])
 
-        uniform_prob = [0.25] * 4
-        with patch(
-            "pennylane_ionq.device.SimulatorDevice.prob", new_callable=PropertyMock
-        ) as mock_prob:
-            mock_prob.return_value = uniform_prob
-            assert np.array_equal(dev.probability(), uniform_prob)
+    #     uniform_prob = [0.25] * 4
+    #     with patch(
+    #         "pennylane_ionq.device.SimulatorDevice.prob"
+    #     ) as mock_prob:
+    #         mock_prob.return_value = uniform_prob
+    #         assert np.array_equal(dev.probability(), uniform_prob)
 
     @pytest.mark.parametrize(
         "backend", ["harmony", "aria-1", "aria-2", "forte-1", None]
@@ -256,6 +256,127 @@ def test_backend_initialization(self, backend):
         )
         assert dev.backend == backend
 
+    def test_batch_execute_probabilities(self, requires_api):
+        """Test batch_execute method when computing circuit probabilities."""
+        dev = SimulatorDevice(wires=(0, 1, 2), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0.5, wires=[0])
+            GPI2(0, wires=[1])
+            MS(0, 0.5, wires=[1, 2])
+            qml.probs(wires=[0, 1, 2])
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            GPI(0.5, wires=[1])
+            MS(0, 0.5, wires=[1, 2])
+            qml.probs(wires=[0, 1, 2])
+        results = dev.batch_execute([tape1, tape2])
+        assert np.array_equal(results[0], [0.0, 0.0, 0.0, 0.0, 0.25, 0.25, 0.25, 0.25])
+        assert np.array_equal(results[1], [0.0, 0.25, 0.25, 0.0, 0.0, 0.25, 0.25, 0.0])
+        dev.set_current_circuit_index(0)
+        assert np.array_equal(
+            dev.probability(),
+            [0.0, 0.0, 0.0, 0.0, 0.25, 0.25, 0.25, 0.25],
+        )
+        dev.set_current_circuit_index(1)
+        assert np.array_equal(
+            dev.probability(),
+            [0.0, 0.25, 0.25, 0.0, 0.0, 0.25, 0.25, 0.0],
+        )
+
+    def test_batch_execute_variance(self, requires_api):
+        """Test batch_execute method when computing variance of an observable."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.var(qml.PauliZ(0))
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            qml.var(qml.PauliZ(0))
+        results = dev.batch_execute([tape1, tape2])
+        assert results[0] == pytest.approx(0, abs=0.01)
+        assert results[1] == pytest.approx(1, abs=0.01)
+
+    def test_batch_execute_expectation_value(self, requires_api):
+        """Test batch_execute method when computing expectation value of an observable."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.expval(qml.PauliZ(0))
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            qml.expval(qml.PauliZ(0))
+        results = dev.batch_execute([tape1, tape2])
+        assert results[0] == pytest.approx(-1, abs=0.1)
+        assert results[1] == pytest.approx(0, abs=0.1)
+
+    def test_batch_execute_sample(self, requires_api):
+        """Test batch_execute method when sampling."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.sample(qml.PauliZ(0))
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            qml.sample(qml.PauliZ(0))
+        results = dev.batch_execute([tape1, tape2])
+        dev.set_current_circuit_index(0)
+        results0 = dev.sample(qml.PauliZ(0))
+        dev.set_current_circuit_index(1)
+        results1 = dev.sample(qml.PauliZ(0))
+        assert set(results[0]) == set(results0) == {-1}
+        assert set(results[1]) == set(results1) == {-1, 1}
+
+    def test_batch_execute_counts(self, requires_api):
+        """Test batch_execute method when computing counts."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.counts(qml.PauliZ(0))
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            qml.counts(qml.PauliZ(0))
+        results = dev.batch_execute([tape1, tape2])
+        assert results[0][-1] == 1024
+        assert results[1][-1] == pytest.approx(512, abs=100)
+
+    def test_sample_measurements(self, requires_api):
+        """Test branch of code activated by using SampleMeasurement."""
+
+        class CountState(SampleMeasurement):
+            def __init__(self, state: str):
+                self.state = state  # string identifying the state e.g. "0101"
+                wires = list(range(len(state)))
+                super().__init__(wires=wires)
+
+            def process_samples(
+                self, samples, wire_order, shot_range=None, bin_size=None
+            ):
+                counts_mp = qml.counts(wires=self._wires)
+                counts = counts_mp.process_samples(
+                    samples, wire_order, shot_range, bin_size
+                )
+                return float(counts.get(self.state, 0))
+
+            def process_counts(self, counts, wire_order):
+                return float(counts.get(self.state, 0))
+
+            def __copy__(self):
+                return CountState(state=self.state)
+
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            CountState(state="1")
+
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            CountState(state="1")
+
+        results = dev.batch_execute([tape1, tape2])
+        assert results[0] == 1024
+        assert results[1] == pytest.approx(512, abs=100)
+
 
 class TestJobAttribute:
     """Tests job creation with mocked submission."""
@@ -282,6 +403,32 @@ def mock_submit_job(*args):
         assert len(dev.job["input"]["circuit"]) == 1
         assert dev.job["input"]["circuit"][0] == {"gate": "x", "target": 0}
 
+    def test_nonparametrized_tape_batch_submit(self, mocker):
+        """Tests job attribute after single paulix tape, on batch submit."""
+
+        def mock_submit_job(*args):
+            pass
+
+        mocker.patch("pennylane_ionq.device.IonQDevice._submit_job", mock_submit_job)
+        dev = IonQDevice(wires=(0,), target="foo")
+
+        with qml.tape.QuantumTape() as tape:
+            qml.PauliX(0)
+
+        dev.reset(circuits_array_length=1)
+        dev.batch_apply(tape.operations, circuit_index=0)
+
+        assert dev.job["input"]["format"] == "ionq.circuit.v0"
+        assert dev.job["input"]["gateset"] == "qis"
+        assert dev.job["target"] == "foo"
+        assert dev.job["input"]["qubits"] == 1
+
+        assert len(dev.job["input"]["circuits"]) == 1
+        assert dev.job["input"]["circuits"][0]["circuit"][0] == {
+            "gate": "x",
+            "target": 0,
+        }
+
     def test_parameterized_op(self, mocker):
         """Tests job attribute several parameterized operations."""
 
@@ -313,6 +460,38 @@ def mock_submit_job(*args):
             "rotation": 2.3456,
         }
 
+    def test_parameterized_op_batch_submit(self, mocker):
+        """Tests job attribute several parameterized operations."""
+
+        def mock_submit_job(*args):
+            pass
+
+        mocker.patch("pennylane_ionq.device.IonQDevice._submit_job", mock_submit_job)
+        dev = IonQDevice(wires=(0,))
+
+        with qml.tape.QuantumTape() as tape:
+            qml.RX(1.2345, wires=0)
+            qml.RY(qml.numpy.array(2.3456), wires=0)
+
+        dev.reset(circuits_array_length=1)
+        dev.batch_apply(tape.operations, circuit_index=0)
+
+        assert dev.job["input"]["format"] == "ionq.circuit.v0"
+        assert dev.job["input"]["gateset"] == "qis"
+        assert dev.job["input"]["qubits"] == 1
+
+        assert len(dev.job["input"]["circuits"][0]["circuit"]) == 2
+        assert dev.job["input"]["circuits"][0]["circuit"][0] == {
+            "gate": "rx",
+            "target": 0,
+            "rotation": 1.2345,
+        }
+        assert dev.job["input"]["circuits"][0]["circuit"][1] == {
+            "gate": "ry",
+            "target": 0,
+            "rotation": 2.3456,
+        }
+
     def test_parameterized_native_op(self, mocker):
         """Tests job attribute several parameterized native operations."""
 
@@ -358,6 +537,54 @@ def mock_submit_job(*args):
             "angle": 0.1,
         }
 
+    def test_parameterized_native_op_batch_submit(self, mocker):
+        """Tests job attribute for several parameterized native operations with batch_execute."""
+
+        class StopExecute(Exception):
+            pass
+
+        def mock_submit_job(*args):
+            raise StopExecute()
+
+        mocker.patch(
+            "pennylane_ionq.device.SimulatorDevice._submit_job", mock_submit_job
+        )
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0.7, wires=[0])
+            GPI2(0.8, wires=[0])
+            qml.expval(qml.PauliZ(0))
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0.9, wires=[0])
+            qml.expval(qml.PauliZ(0))
+
+        try:
+            dev.batch_execute([tape1, tape2])
+        except StopExecute:
+            pass
+
+        assert dev.job["input"]["format"] == "ionq.circuit.v0"
+        assert dev.job["input"]["gateset"] == "native"
+        assert dev.job["target"] == "simulator"
+        assert dev.job["input"]["qubits"] == 1
+        assert len(dev.job["input"]["circuits"]) == 2
+        assert dev.job["input"]["circuits"][0]["circuit"][0] == {
+            "gate": "gpi",
+            "target": 0,
+            "phase": 0.7,
+        }
+        assert dev.job["input"]["circuits"][0]["circuit"][1] == {
+            "gate": "gpi2",
+            "target": 0,
+            "phase": 0.8,
+        }
+        assert dev.job["input"]["circuits"][1]["circuit"][0] == {
+            "gate": "gpi2",
+            "target": 0,
+            "phase": 0.9,
+        }
+
     @pytest.mark.parametrize(
         "phi0, phi1, theta",
         [

From 07557877d90066a414a827c2b8a5677ea4f8c03d Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 11 Sep 2024 11:09:26 +0300
Subject: [PATCH 07/33] Correct unit test after updating pennylane baseline
 code to latest version.

---
 tests/test_device.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_device.py b/tests/test_device.py
index ccfed8a..033e7ef 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -67,7 +67,7 @@ def test_load_device(self, d):
         """Test that the device loads correctly"""
         dev = qml.device(d, wires=2, shots=1024)
         assert dev.num_wires == 2
-        assert dev.shots == 1024
+        assert dev.shots.total_shots == 1024
         assert dev.short_name == d
 
     @pytest.mark.parametrize("d", shortnames)

From 34be4271bfaee555fededb13065e364b18c543ce Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 19 Sep 2024 10:37:06 +0300
Subject: [PATCH 08/33] Fix codefactor issues.

---
 pennylane_ionq/device.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 3c67cff..0fbca31 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -145,6 +145,7 @@ def __init__(
         self._operation_map = _GATESET_OPS[gateset]
         self.histogram = None
         self.histograms = None
+        self._samples = None
         self.reset()
 
     def reset(self, circuits_array_length=0):
@@ -272,6 +273,8 @@ def batch_execute(self, circuits):
 
     def batch_apply(self, operations, circuit_index, **kwargs):
 
+        "Apply circuit operations when submitting for execution a batch of circuits."
+
         rotations = kwargs.pop("rotations", [])
 
         if len(operations) == 0 and len(rotations) == 0:

From 2832900c84636d85b05f8bd6e30831bfa9ec9c35 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 19 Sep 2024 10:41:43 +0300
Subject: [PATCH 09/33] Run black code formatter.

---
 pennylane_ionq/device.py | 60 +++++++++++++++++++++++++---------------
 1 file changed, 37 insertions(+), 23 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 0fbca31..f6e2e2d 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -133,7 +133,9 @@ def __init__(
         sharpen=False,
     ):
         if shots is None:
-            raise ValueError("The ionq device does not support analytic expectation values.")
+            raise ValueError(
+                "The ionq device does not support analytic expectation values."
+            )
 
         super().__init__(wires=wires, shots=shots)
         self._current_circuit_index = None
@@ -311,7 +313,9 @@ def apply(self, operations, **kwargs):
         rotations = kwargs.pop("rotations", [])
 
         if len(operations) == 0 and len(rotations) == 0:
-            warnings.warn("Circuit is empty. Empty circuits return failures. Submitting anyway.")
+            warnings.warn(
+                "Circuit is empty. Empty circuits return failures. Submitting anyway."
+            )
 
         for i, operation in enumerate(operations):
             if i > 0 and operation.name in {
@@ -408,7 +412,8 @@ def prob(self):
         # Here, we rearrange the states to match the big-endian ordering
         # expected by PennyLane.
         basis_states = (
-            int(bin(int(k))[2:].rjust(self.num_wires, "0")[::-1], 2) for k in self.histogram
+            int(bin(int(k))[2:].rjust(self.num_wires, "0")[::-1], 2)
+            for k in self.histogram
         )
         idx = np.fromiter(basis_states, dtype=int)
 
@@ -425,16 +430,14 @@ def prob(self):
 
         return prob_array
 
-    def estimate_probability(
-        self, wires=None, shot_range=None, bin_size=None
-    ):
+    def estimate_probability(self, wires=None, shot_range=None, bin_size=None):
         """Return the estimated probability of each computational basis state
         using the generated samples. Overwrites the method in QubitDevice class
-        to accomodate batch submission of circuits. When invoking this method after 
-        submitting circuits using any other method than QubitDevice.execute() 
-        method, current circuit index should be set via set_current_circuit_index() 
-        method to the circuit which is targeted out of one or potentially several 
-        circuits submitted in a batch. This is not very good design but the 
+        to accomodate batch submission of circuits. When invoking this method after
+        submitting circuits using any other method than QubitDevice.execute()
+        method, current circuit index should be set via set_current_circuit_index()
+        method to the circuit which is targeted out of one or potentially several
+        circuits submitted in a batch. This is not very good design but the
         alternative would be to create huge amounts of code duplication with respect
         to the Pennylane base code.
 
@@ -457,7 +460,10 @@ def estimate_probability(
             return result
         except TypeError as e:
             self._samples = None
-            e.args = ("When invoking QubitDevice estimate_probability() method after submitting circuits using any other method than QubitDevice.execute() method, current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
+            e.args = (
+                "When invoking QubitDevice estimate_probability() method after submitting circuits using any other method than QubitDevice.execute() method, current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch.",
+                *e.args,
+            )
             raise
 
     def sample(
@@ -468,11 +474,11 @@ def sample(
         counts=False,
     ):
         """Return samples of an observable. Overwrites the method in QubitDevice class
-        to accomodate batch submission of circuits. When invoking this method after submitting 
-        circuits using any other method than QubitDevice.execute() method, current circuit 
-        index should be set via set_current_circuit_index() method pointing to index of the 
+        to accomodate batch submission of circuits. When invoking this method after submitting
+        circuits using any other method than QubitDevice.execute() method, current circuit
+        index should be set via set_current_circuit_index() method pointing to index of the
         circuit which is targeted out of one or potentially several circuits submitted in a batch.
-        This is not very good design but the alternative would be to create huge amounts of 
+        This is not very good design but the alternative would be to create huge amounts of
         code duplication with respect to the Pennylane base code.
 
         Args:
@@ -501,7 +507,10 @@ def sample(
             return result
         except TypeError as e:
             self._samples = None
-            e.args = ("When invoking QubitDevice sample() method after submitting circuits using any other method than QubitDevice.execute() method current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch.", *e.args)
+            e.args = (
+                "When invoking QubitDevice sample() method after submitting circuits using any other method than QubitDevice.execute() method current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch.",
+                *e.args,
+            )
             raise
 
     def _measure(
@@ -512,10 +521,10 @@ def _measure(
     ):
         """Compute the corresponding measurement process depending on ``shots`` and the measurement
         type. Overwrites the method in QubitDevice class to accomodate batch submission of circuits.
-        When invoking this method after submitting circuits using any other method than QubitDevice.execute() 
-        method, current circuit index should be set via set_current_circuit_index() method pointing to index 
-        of the circuit which is targeted out of one or potentially several circuits submitted in a batch. 
-        This is not very good design but the alternative would be to create huge amounts of code duplication 
+        When invoking this method after submitting circuits using any other method than QubitDevice.execute()
+        method, current circuit index should be set via set_current_circuit_index() method pointing to index
+        of the circuit which is targeted out of one or potentially several circuits submitted in a batch.
+        This is not very good design but the alternative would be to create huge amounts of code duplication
         with respect to the Pennylane base code.
 
         Args:
@@ -539,7 +548,10 @@ def _measure(
             return result
         except TypeError as e:
             self._samples = None
-            e.args = ("When invoking QubitDevice _measure() method after submitting circuits using any other method than QubitDevice.execute() method, current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch. ", *e.args)
+            e.args = (
+                "When invoking QubitDevice _measure() method after submitting circuits using any other method than QubitDevice.execute() method, current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch. ",
+                *e.args,
+            )
             raise
 
     def probability(self, wires=None, shot_range=None, bin_size=None):
@@ -548,7 +560,9 @@ def probability(self, wires=None, shot_range=None, bin_size=None):
         if shot_range is None and bin_size is None:
             return self.marginal_prob(self.prob, wires)
 
-        return self.estimate_probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
+        return self.estimate_probability(
+            wires=wires, shot_range=shot_range, bin_size=bin_size
+        )
 
 
 class SimulatorDevice(IonQDevice):

From c19af9a03b1d72cd6f4f459c9a293756fcd8a7ba Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 19 Sep 2024 12:04:10 +0300
Subject: [PATCH 10/33] Shots cannot be none in an IonQDevice. Remove check on
 shots.

---
 pennylane_ionq/device.py | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index f6e2e2d..c84cc6b 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -224,11 +224,7 @@ def batch_execute(self, circuits):
         results = []
         for circuit_index, circuit in enumerate(circuits):
             self.set_current_circuit_index(circuit_index)
-            sample_type = (SampleMP, CountsMP, ClassicalShadowMP, ShadowExpvalMP)
-            if self.shots is not None or any(
-                isinstance(m, sample_type) for m in circuit.measurements
-            ):
-                self._samples = self.generate_samples()
+            self._samples = self.generate_samples()
 
             # compute the required statistics
             if self._shot_vector is not None:

From db1f366a3d34cb22e918b8ec718245f366be00bb Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 19 Sep 2024 12:15:26 +0300
Subject: [PATCH 11/33] Correct docstring comment.

---
 pennylane_ionq/device.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index c84cc6b..4d23e8a 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -486,7 +486,6 @@ def sample(
                 provided, the entire shot range is treated as a single bin.
             counts (bool): whether counts (``True``) or raw samples (``False``)
                 should be returned
-            circuit_index (int): index of circuit in case of batch circuit submission
 
         Raises:
             EigvalsUndefinedError: if no information is available about the

From 4755b9dc11cd42c29753a0ab4c45a58f06cbba67 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Tue, 24 Sep 2024 14:29:43 +0300
Subject: [PATCH 12/33] Remove exeception handling code.

---
 pennylane_ionq/device.py | 50 +++++++++++-----------------------------
 1 file changed, 14 insertions(+), 36 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 4d23e8a..c70813e 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -449,18 +449,10 @@ def estimate_probability(self, wires=None, shot_range=None, bin_size=None):
         Returns:
             array[float]: list of the probabilities
         """
-        try:
-            self._samples = self.generate_samples()
-            result = super().estimate_probability(wires, shot_range, bin_size)
-            self._samples = None
-            return result
-        except TypeError as e:
-            self._samples = None
-            e.args = (
-                "When invoking QubitDevice estimate_probability() method after submitting circuits using any other method than QubitDevice.execute() method, current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch.",
-                *e.args,
-            )
-            raise
+        self._samples = self.generate_samples()
+        result = super().estimate_probability(wires, shot_range, bin_size)
+        self._samples = None
+        return result
 
     def sample(
         self,
@@ -495,18 +487,11 @@ def sample(
             Union[array[float], dict, list[dict]]: samples in an array of
             dimension ``(shots,)`` or counts
         """
-        try:
-            self._samples = self.generate_samples()
-            result = super().sample(observable, shot_range, bin_size, counts)
-            self._samples = None
-            return result
-        except TypeError as e:
-            self._samples = None
-            e.args = (
-                "When invoking QubitDevice sample() method after submitting circuits using any other method than QubitDevice.execute() method current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch.",
-                *e.args,
-            )
-            raise
+        self._samples = self.generate_samples()
+        result = super().sample(observable, shot_range, bin_size, counts)
+        self._samples = None
+        return result
+
 
     def _measure(
         self,
@@ -536,18 +521,11 @@ def _measure(
         Returns:
             Union[float, dict, list[float]]: result of the measurement
         """
-        try:
-            self._samples = self.generate_samples()
-            result = super()._measure(measurement, shot_range, bin_size)
-            self._samples = None
-            return result
-        except TypeError as e:
-            self._samples = None
-            e.args = (
-                "When invoking QubitDevice _measure() method after submitting circuits using any other method than QubitDevice.execute() method, current circuit index should be set via set_current_circuit_index() method pointing to index of the circuit which is targeted out of one or potentially several circuits submitted in a batch. ",
-                *e.args,
-            )
-            raise
+        self._samples = self.generate_samples()
+        result = super()._measure(measurement, shot_range, bin_size)
+        self._samples = None
+        return result
+
 
     def probability(self, wires=None, shot_range=None, bin_size=None):
         wires = wires or self.wires

From 287c552563d506e66bae2c0703140e4395a4a7c5 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Tue, 24 Sep 2024 14:54:32 +0300
Subject: [PATCH 13/33] Remove self.histogram, replace with self.histograms.

---
 pennylane_ionq/device.py | 24 ++++++++++--------------
 tests/test_device.py     |  4 ++--
 2 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index c70813e..896e837 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -145,16 +145,14 @@ def __init__(
         self.error_mitigation = error_mitigation
         self.sharpen = sharpen
         self._operation_map = _GATESET_OPS[gateset]
-        self.histogram = None
-        self.histograms = None
+        self.histograms = []
         self._samples = None
         self.reset()
 
     def reset(self, circuits_array_length=0):
         """Reset the device"""
         self._current_circuit_index = None
-        self.histogram = None
-        self.histograms = None
+        self.histograms = []
         if circuits_array_length == 0:
             self.input = {
                 "format": "ionq.circuit.v0",
@@ -391,7 +389,6 @@ def _submit_job(self):
         else:
             self.histograms = []
             self.histograms.append(job.data.value)
-            self.histogram = job.data.value
 
     @property
     def prob(self):
@@ -399,17 +396,19 @@ def prob(self):
         no job has been submitted, returns ``None``.
         """
         if self._current_circuit_index is not None:
-            self.histogram = self.histograms[self._current_circuit_index]
-
-        if self.histogram is None:
-            return None
+            histogram = self.histograms[self._current_circuit_index]
+        else:
+            try:
+                histogram = self.histograms[0]
+            except IndexError:
+                return None            
 
         # The IonQ API returns basis states using little-endian ordering.
         # Here, we rearrange the states to match the big-endian ordering
         # expected by PennyLane.
         basis_states = (
             int(bin(int(k))[2:].rjust(self.num_wires, "0")[::-1], 2)
-            for k in self.histogram
+            for k in histogram
         )
         idx = np.fromiter(basis_states, dtype=int)
 
@@ -417,13 +416,10 @@ def prob(self):
         prob_array = np.zeros([2**self.num_wires])
 
         # histogram values don't always perfectly sum to exactly one
-        histogram_values = self.histogram.values()
+        histogram_values = histogram.values()
         norm = sum(histogram_values)
         prob_array[idx] = np.fromiter(histogram_values, float) / norm
 
-        if self._current_circuit_index is not None:
-            self.histogram = None
-
         return prob_array
 
     def estimate_probability(self, wires=None, shot_range=None, bin_size=None):
diff --git a/tests/test_device.py b/tests/test_device.py
index 033e7ef..8e1dd52 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -43,10 +43,10 @@ def test_generate_samples_qpu_device(self, wires, histogram):
         """Test that the generate_samples method for QPUDevices shuffles the samples between calls."""
 
         dev = QPUDevice(wires, shots=1024, api_key=FAKE_API_KEY)
-        dev.histogram = histogram
+        dev.histograms = [histogram]
 
         sample1 = dev.generate_samples()
-        assert dev.histogram == histogram  # make sure histogram is still the same
+        assert dev.histograms[0] == histogram  # make sure histogram is still the same
         sample2 = dev.generate_samples()
         assert not np.all(sample1 == sample2)  # some rows are different
 

From d9220163ef53e2f17112c498c31a45c3db0b9b2e Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Tue, 24 Sep 2024 16:16:11 +0300
Subject: [PATCH 14/33] Uniformize treatment of one vs multiple circuits.

---
 pennylane_ionq/device.py | 29 +++++++++--------------------
 tests/test_device.py     | 20 ++++++++++----------
 2 files changed, 19 insertions(+), 30 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 896e837..85fe52a 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -149,24 +149,16 @@ def __init__(
         self._samples = None
         self.reset()
 
-    def reset(self, circuits_array_length=0):
+    def reset(self, circuits_array_length=1):
         """Reset the device"""
         self._current_circuit_index = None
         self.histograms = []
-        if circuits_array_length == 0:
-            self.input = {
-                "format": "ionq.circuit.v0",
-                "qubits": self.num_wires,
-                "circuit": [],
-                "gateset": self.gateset,
-            }
-        else:
-            self.input = {
-                "format": "ionq.circuit.v0",
-                "qubits": self.num_wires,
-                "circuits": [{"circuit": []} for _ in range(circuits_array_length)],
-                "gateset": self.gateset,
-            }
+        self.input = {
+            "format": "ionq.circuit.v0",
+            "qubits": self.num_wires,
+            "circuits": [{"circuit": []} for _ in range(circuits_array_length)],
+            "gateset": self.gateset,
+        }
         self.job = {
             "input": self.input,
             "target": self.target,
@@ -328,7 +320,7 @@ def apply(self, operations, **kwargs):
 
         self._submit_job()
 
-    def _apply_operation(self, operation, circuit_index=None):
+    def _apply_operation(self, operation, circuit_index=0):
         name = operation.name
         wires = self.map_wires(operation.wires).tolist()
         gate = {"gate": self._operation_map[name]}
@@ -354,10 +346,7 @@ def _apply_operation(self, operation, circuit_index=None):
         elif par:
             gate["rotation"] = float(par[0])
 
-        if "circuit" in self.input:
-            self.input["circuit"].append(gate)
-        else:
-            self.input["circuits"][circuit_index]["circuit"].append(gate)
+        self.input["circuits"][circuit_index]["circuit"].append(gate)
 
     def _submit_job(self):
         job = Job(api_key=self.api_key)
diff --git a/tests/test_device.py b/tests/test_device.py
index 8e1dd52..0c63e10 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -366,8 +366,8 @@ def mock_submit_job(*args):
         assert dev.job["target"] == "foo"
         assert dev.job["input"]["qubits"] == 1
 
-        assert len(dev.job["input"]["circuit"]) == 1
-        assert dev.job["input"]["circuit"][0] == {"gate": "x", "target": 0}
+        assert len(dev.job["input"]["circuits"][0]) == 1
+        assert dev.job["input"]["circuits"][0]["circuit"][0] == {"gate": "x", "target": 0}
 
     def test_nonparametrized_tape_batch_submit(self, mocker):
         """Tests job attribute after single paulix tape, on batch submit."""
@@ -414,13 +414,13 @@ def mock_submit_job(*args):
         assert dev.job["input"]["gateset"] == "qis"
         assert dev.job["input"]["qubits"] == 1
 
-        assert len(dev.job["input"]["circuit"]) == 2
-        assert dev.job["input"]["circuit"][0] == {
+        assert len(dev.job["input"]["circuits"][0]["circuit"]) == 2
+        assert dev.job["input"]["circuits"][0]["circuit"][0] == {
             "gate": "rx",
             "target": 0,
             "rotation": 1.2345,
         }
-        assert dev.job["input"]["circuit"][1] == {
+        assert dev.job["input"]["circuits"][0]["circuit"][1] == {
             "gate": "ry",
             "target": 0,
             "rotation": 2.3456,
@@ -479,24 +479,24 @@ def mock_submit_job(*args):
         assert dev.job["input"]["gateset"] == "native"
         assert dev.job["input"]["qubits"] == 3
 
-        assert len(dev.job["input"]["circuit"]) == 4
-        assert dev.job["input"]["circuit"][0] == {
+        assert len(dev.job["input"]["circuits"][0]["circuit"]) == 4
+        assert dev.job["input"]["circuits"][0]["circuit"][0] == {
             "gate": "gpi",
             "target": 0,
             "phase": 0.1,
         }
-        assert dev.job["input"]["circuit"][1] == {
+        assert dev.job["input"]["circuits"][0]["circuit"][1] == {
             "gate": "gpi2",
             "target": 1,
             "phase": 0.2,
         }
-        assert dev.job["input"]["circuit"][2] == {
+        assert dev.job["input"]["circuits"][0]["circuit"][2] == {
             "gate": "ms",
             "targets": [1, 2],
             "phases": [0.2, 0.3],
             "angle": 0.25,
         }
-        assert dev.job["input"]["circuit"][3] == {
+        assert dev.job["input"]["circuits"][0]["circuit"][3] == {
             "gate": "ms",
             "targets": [1, 2],
             "phases": [0.4, 0.5],

From ec02077122d107f4a220262e9f42acc0b784ac27 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 25 Sep 2024 15:09:14 +0300
Subject: [PATCH 15/33] Improve current_circuit_index handling with raising
 exceptions, add tests.

---
 pennylane_ionq/device.py |  13 +++++
 tests/test_device.py     | 107 ++++++++++++++++++++++++++++++++++++++-
 2 files changed, 119 insertions(+), 1 deletion(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 85fe52a..681bcdb 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -76,6 +76,16 @@
     "qis": _qis_operation_map,
 }
 
+class CircuitIndexNotSetException(Exception):
+    """Raised when after submitting multiple circuits circuit index is not set
+       before the user want to access implementation methods of IonQDevice
+       like probability(), estimate_probability(), sample() or the prob property.
+    """
+    
+    def __init__(self):
+        self.message = "Because multiple circuits have been submitted in this job, the index of the circuit \
+you want to acces must be first set via the set_current_circuit_index device method."
+        super().__init__(self.message)
 
 class IonQDevice(QubitDevice):
     r"""IonQ device for PennyLane.
@@ -384,6 +394,9 @@ def prob(self):
         """None or array[float]: Array of computational basis state probabilities. If
         no job has been submitted, returns ``None``.
         """
+        if self._current_circuit_index is None and len(self.histograms) > 1:
+            raise CircuitIndexNotSetException()
+
         if self._current_circuit_index is not None:
             histogram = self.histograms[self._current_circuit_index]
         else:
diff --git a/tests/test_device.py b/tests/test_device.py
index 0c63e10..f2c07de 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -20,7 +20,7 @@
 
 from conftest import shortnames
 from pennylane_ionq.api_client import JobExecutionError, ResourceManager, Job
-from pennylane_ionq.device import QPUDevice, IonQDevice, SimulatorDevice
+from pennylane_ionq.device import QPUDevice, IonQDevice, SimulatorDevice, CircuitIndexNotSetException
 from pennylane_ionq.ops import GPI, GPI2, MS
 from pennylane.measurements import SampleMeasurement
 
@@ -222,6 +222,22 @@ def test_backend_initialization(self, backend):
         )
         assert dev.backend == backend
 
+    def test_batch_execute_probabilities_raises (self, requires_api):
+        """Test invoking probability() method raises exception if circuit index not
+            previously set when multiple circuits are submitted in one job.
+        """
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0.5, wires=[0])
+            qml.probs(wires=[0])
+        dev.batch_execute([tape1, tape1])
+        with pytest.raises(
+            CircuitIndexNotSetException,
+            match="Because multiple circuits have been submitted in this job, the index of the circuit \
+you want to acces must be first set via the set_current_circuit_index device method."
+            ):
+            dev.probability() 
+
     def test_batch_execute_probabilities(self, requires_api):
         """Test batch_execute method when computing circuit probabilities."""
         dev = SimulatorDevice(wires=(0, 1, 2), gateset="native", shots=1024)
@@ -275,6 +291,22 @@ def test_batch_execute_expectation_value(self, requires_api):
         assert results[0] == pytest.approx(-1, abs=0.1)
         assert results[1] == pytest.approx(0, abs=0.1)
 
+    def test_batch_execute_sample_raises (self, requires_api):
+        """Test invoking sample() method raises exception if circuit index not
+            previously set when multiple circuits are submitted in one job.
+        """
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0.5, wires=[0])
+            qml.probs(wires=[0])
+        dev.batch_execute([tape1, tape1])
+        with pytest.raises(
+            CircuitIndexNotSetException,
+            match="Because multiple circuits have been submitted in this job, the index of the circuit \
+you want to acces must be first set via the set_current_circuit_index device method."
+            ):
+            dev.sample(qml.PauliZ(0))
+
     def test_batch_execute_sample(self, requires_api):
         """Test batch_execute method when sampling."""
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
@@ -292,6 +324,79 @@ def test_batch_execute_sample(self, requires_api):
         assert set(results[0]) == set(results0) == {-1}
         assert set(results[1]) == set(results1) == {-1, 1}
 
+    def test_batch_execute_estimate_probability_raises (self, requires_api):
+        """Test invoking estimate_probability() method raises exception if circuit index not
+            previously set when multiple circuits are submitted in one job.
+        """
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.probs(wires=[0])
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            qml.probs(wires=[0])
+        dev.batch_execute([tape1, tape2])
+        with pytest.raises(
+            CircuitIndexNotSetException,
+            match="Because multiple circuits have been submitted in this job, the index of the circuit \
+you want to acces must be first set via the set_current_circuit_index device method."
+            ):
+            dev.estimate_probability()
+
+    def test_batch_execute_estimate_probability(self, requires_api):
+        """Test batch_execute method with invoking estimate_probability method."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.sample(qml.PauliZ(0))
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            qml.sample(qml.PauliZ(0))
+        dev.batch_execute([tape1, tape2])
+        dev.set_current_circuit_index(0)
+        prob0 = dev.estimate_probability()
+        dev.set_current_circuit_index(1)
+        prob1 = dev.estimate_probability()
+        np.testing.assert_array_almost_equal(prob0, [0., 1.], decimal = 1)
+        np.testing.assert_array_almost_equal(prob1, [0.5, 0.5], decimal = 1)
+
+
+    def test_batch_execute_invoking_prob_property_raises (self, requires_api):
+        """Test invoking prob device property raises exception if circuit index not
+            previously set when multiple circuits are submitted in one job.
+        """
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.probs(wires=[0])
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            qml.probs(wires=[0])
+        dev.batch_execute([tape1, tape2])
+        with pytest.raises(
+            CircuitIndexNotSetException,
+            match="Because multiple circuits have been submitted in this job, the index of the circuit \
+you want to acces must be first set via the set_current_circuit_index device method."
+            ):
+            dev.prob
+
+    def test_batch_execute_prob_property(self, requires_api):
+        """Test batch_execute method with invoking invoking prob device property."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.sample(qml.PauliZ(0))
+        with qml.tape.QuantumTape() as tape2:
+            GPI2(0, wires=[0])
+            qml.sample(qml.PauliZ(0))
+        dev.batch_execute([tape1, tape2])
+        dev.set_current_circuit_index(0)
+        prob0 = dev.prob
+        dev.set_current_circuit_index(1)
+        prob1 = dev.prob
+        np.testing.assert_array_almost_equal(prob0, [0., 1.], decimal = 1)
+        np.testing.assert_array_almost_equal(prob1, [0.5, 0.5], decimal = 1)
+
     def test_batch_execute_counts(self, requires_api):
         """Test batch_execute method when computing counts."""
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)

From 8db4dbcddf5011e248c5db99fcfcdafa4523dfab Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 25 Sep 2024 16:18:02 +0300
Subject: [PATCH 16/33] Reset samples in reset function.

---
 pennylane_ionq/device.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 681bcdb..e4a8aec 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -162,6 +162,7 @@ def __init__(
     def reset(self, circuits_array_length=1):
         """Reset the device"""
         self._current_circuit_index = None
+        self._samples = None
         self.histograms = []
         self.input = {
             "format": "ionq.circuit.v0",

From 12d4020d31f858afdf0ca00705abc34c7ab51be1 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 25 Sep 2024 17:00:03 +0300
Subject: [PATCH 17/33] Run black code formatter.

---
 pennylane_ionq/device.py | 15 +++----
 tests/test_device.py     | 96 ++++++++++++++++++++++++++--------------
 2 files changed, 69 insertions(+), 42 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index e4a8aec..53d2fab 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -76,17 +76,19 @@
     "qis": _qis_operation_map,
 }
 
+
 class CircuitIndexNotSetException(Exception):
     """Raised when after submitting multiple circuits circuit index is not set
-       before the user want to access implementation methods of IonQDevice
-       like probability(), estimate_probability(), sample() or the prob property.
+    before the user want to access implementation methods of IonQDevice
+    like probability(), estimate_probability(), sample() or the prob property.
     """
-    
+
     def __init__(self):
         self.message = "Because multiple circuits have been submitted in this job, the index of the circuit \
 you want to acces must be first set via the set_current_circuit_index device method."
         super().__init__(self.message)
 
+
 class IonQDevice(QubitDevice):
     r"""IonQ device for PennyLane.
 
@@ -404,14 +406,13 @@ def prob(self):
             try:
                 histogram = self.histograms[0]
             except IndexError:
-                return None            
+                return None
 
         # The IonQ API returns basis states using little-endian ordering.
         # Here, we rearrange the states to match the big-endian ordering
         # expected by PennyLane.
         basis_states = (
-            int(bin(int(k))[2:].rjust(self.num_wires, "0")[::-1], 2)
-            for k in histogram
+            int(bin(int(k))[2:].rjust(self.num_wires, "0")[::-1], 2) for k in histogram
         )
         idx = np.fromiter(basis_states, dtype=int)
 
@@ -491,7 +492,6 @@ def sample(
         self._samples = None
         return result
 
-
     def _measure(
         self,
         measurement: Union[SampleMeasurement, StateMeasurement],
@@ -525,7 +525,6 @@ def _measure(
         self._samples = None
         return result
 
-
     def probability(self, wires=None, shot_range=None, bin_size=None):
         wires = wires or self.wires
 
diff --git a/tests/test_device.py b/tests/test_device.py
index f2c07de..582a6ec 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -20,7 +20,12 @@
 
 from conftest import shortnames
 from pennylane_ionq.api_client import JobExecutionError, ResourceManager, Job
-from pennylane_ionq.device import QPUDevice, IonQDevice, SimulatorDevice, CircuitIndexNotSetException
+from pennylane_ionq.device import (
+    QPUDevice,
+    IonQDevice,
+    SimulatorDevice,
+    CircuitIndexNotSetException,
+)
 from pennylane_ionq.ops import GPI, GPI2, MS
 from pennylane.measurements import SampleMeasurement
 
@@ -52,11 +57,15 @@ def test_generate_samples_qpu_device(self, wires, histogram):
 
         unique_outcomes1 = np.unique(sample1, axis=0)
         unique_outcomes2 = np.unique(sample2, axis=0)
-        assert np.all(unique_outcomes1 == unique_outcomes2)  # possible outcomes are the same
+        assert np.all(
+            unique_outcomes1 == unique_outcomes2
+        )  # possible outcomes are the same
 
         sorted_outcomes1 = np.sort(sample1, axis=0)
         sorted_outcomes2 = np.sort(sample2, axis=0)
-        assert np.all(sorted_outcomes1 == sorted_outcomes2)  # set of outcomes is the same
+        assert np.all(
+            sorted_outcomes1 == sorted_outcomes2
+        )  # set of outcomes is the same
 
 
 class TestDeviceIntegration:
@@ -96,7 +105,9 @@ def test_failedcircuit(self, monkeypatch):
         monkeypatch.setattr(
             requests, "post", lambda url, timeout, data, headers: (url, data, headers)
         )
-        monkeypatch.setattr(ResourceManager, "handle_response", lambda self, response: None)
+        monkeypatch.setattr(
+            ResourceManager, "handle_response", lambda self, response: None
+        )
         monkeypatch.setattr(Job, "is_complete", False)
         monkeypatch.setattr(Job, "is_failed", True)
 
@@ -111,13 +122,17 @@ def test_shots(self, shots, monkeypatch, mocker, tol):
         monkeypatch.setattr(
             requests, "post", lambda url, timeout, data, headers: (url, data, headers)
         )
-        monkeypatch.setattr(ResourceManager, "handle_response", lambda self, response: None)
+        monkeypatch.setattr(
+            ResourceManager, "handle_response", lambda self, response: None
+        )
         monkeypatch.setattr(Job, "is_complete", True)
 
         def fake_response(self, resource_id=None, params=None):
             """Return fake response data"""
             fake_json = {"0": 1}
-            setattr(self.resource, "data", type("data", tuple(), {"value": fake_json})())
+            setattr(
+                self.resource, "data", type("data", tuple(), {"value": fake_json})()
+            )
 
         monkeypatch.setattr(ResourceManager, "get", fake_response)
 
@@ -133,20 +148,26 @@ def circuit():
         circuit()
         assert json.loads(spy.call_args[1]["data"])["shots"] == shots
 
-    @pytest.mark.parametrize("error_mitigation", [None, {"debias": True}, {"debias": False}])
+    @pytest.mark.parametrize(
+        "error_mitigation", [None, {"debias": True}, {"debias": False}]
+    )
     def test_error_mitigation(self, error_mitigation, monkeypatch, mocker):
         """Test that shots are correctly specified when submitting a job to the API."""
 
         monkeypatch.setattr(
             requests, "post", lambda url, timeout, data, headers: (url, data, headers)
         )
-        monkeypatch.setattr(ResourceManager, "handle_response", lambda self, response: None)
+        monkeypatch.setattr(
+            ResourceManager, "handle_response", lambda self, response: None
+        )
         monkeypatch.setattr(Job, "is_complete", True)
 
         def fake_response(self, resource_id=None, params=None):
             """Return fake response data"""
             fake_json = {"0": 1}
-            setattr(self.resource, "data", type("data", tuple(), {"value": fake_json})())
+            setattr(
+                self.resource, "data", type("data", tuple(), {"value": fake_json})()
+            )
 
         monkeypatch.setattr(ResourceManager, "get", fake_response)
 
@@ -167,7 +188,10 @@ def circuit():
         spy = mocker.spy(requests, "post")
         circuit()
         if error_mitigation is not None:
-            assert json.loads(spy.call_args[1]["data"])["error_mitigation"] == error_mitigation
+            assert (
+                json.loads(spy.call_args[1]["data"])["error_mitigation"]
+                == error_mitigation
+            )
         else:
             with pytest.raises(KeyError, match="error_mitigation"):
                 json.loads(spy.call_args[1]["data"])["error_mitigation"]
@@ -211,7 +235,9 @@ def test_prob_no_results(self, d):
         dev = qml.device(d, wires=1, shots=1)
         assert dev.prob is None
 
-    @pytest.mark.parametrize("backend", ["harmony", "aria-1", "aria-2", "forte-1", None])
+    @pytest.mark.parametrize(
+        "backend", ["harmony", "aria-1", "aria-2", "forte-1", None]
+    )
     def test_backend_initialization(self, backend):
         """Test that the device initializes with the correct backend."""
         dev = qml.device(
@@ -222,9 +248,9 @@ def test_backend_initialization(self, backend):
         )
         assert dev.backend == backend
 
-    def test_batch_execute_probabilities_raises (self, requires_api):
+    def test_batch_execute_probabilities_raises(self, requires_api):
         """Test invoking probability() method raises exception if circuit index not
-            previously set when multiple circuits are submitted in one job.
+        previously set when multiple circuits are submitted in one job.
         """
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
         with qml.tape.QuantumTape() as tape1:
@@ -234,9 +260,9 @@ def test_batch_execute_probabilities_raises (self, requires_api):
         with pytest.raises(
             CircuitIndexNotSetException,
             match="Because multiple circuits have been submitted in this job, the index of the circuit \
-you want to acces must be first set via the set_current_circuit_index device method."
-            ):
-            dev.probability() 
+you want to acces must be first set via the set_current_circuit_index device method.",
+        ):
+            dev.probability()
 
     def test_batch_execute_probabilities(self, requires_api):
         """Test batch_execute method when computing circuit probabilities."""
@@ -291,9 +317,9 @@ def test_batch_execute_expectation_value(self, requires_api):
         assert results[0] == pytest.approx(-1, abs=0.1)
         assert results[1] == pytest.approx(0, abs=0.1)
 
-    def test_batch_execute_sample_raises (self, requires_api):
+    def test_batch_execute_sample_raises(self, requires_api):
         """Test invoking sample() method raises exception if circuit index not
-            previously set when multiple circuits are submitted in one job.
+        previously set when multiple circuits are submitted in one job.
         """
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
         with qml.tape.QuantumTape() as tape1:
@@ -303,8 +329,8 @@ def test_batch_execute_sample_raises (self, requires_api):
         with pytest.raises(
             CircuitIndexNotSetException,
             match="Because multiple circuits have been submitted in this job, the index of the circuit \
-you want to acces must be first set via the set_current_circuit_index device method."
-            ):
+you want to acces must be first set via the set_current_circuit_index device method.",
+        ):
             dev.sample(qml.PauliZ(0))
 
     def test_batch_execute_sample(self, requires_api):
@@ -324,9 +350,9 @@ def test_batch_execute_sample(self, requires_api):
         assert set(results[0]) == set(results0) == {-1}
         assert set(results[1]) == set(results1) == {-1, 1}
 
-    def test_batch_execute_estimate_probability_raises (self, requires_api):
+    def test_batch_execute_estimate_probability_raises(self, requires_api):
         """Test invoking estimate_probability() method raises exception if circuit index not
-            previously set when multiple circuits are submitted in one job.
+        previously set when multiple circuits are submitted in one job.
         """
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
         with qml.tape.QuantumTape() as tape1:
@@ -339,8 +365,8 @@ def test_batch_execute_estimate_probability_raises (self, requires_api):
         with pytest.raises(
             CircuitIndexNotSetException,
             match="Because multiple circuits have been submitted in this job, the index of the circuit \
-you want to acces must be first set via the set_current_circuit_index device method."
-            ):
+you want to acces must be first set via the set_current_circuit_index device method.",
+        ):
             dev.estimate_probability()
 
     def test_batch_execute_estimate_probability(self, requires_api):
@@ -357,13 +383,12 @@ def test_batch_execute_estimate_probability(self, requires_api):
         prob0 = dev.estimate_probability()
         dev.set_current_circuit_index(1)
         prob1 = dev.estimate_probability()
-        np.testing.assert_array_almost_equal(prob0, [0., 1.], decimal = 1)
-        np.testing.assert_array_almost_equal(prob1, [0.5, 0.5], decimal = 1)
-
+        np.testing.assert_array_almost_equal(prob0, [0.0, 1.0], decimal=1)
+        np.testing.assert_array_almost_equal(prob1, [0.5, 0.5], decimal=1)
 
-    def test_batch_execute_invoking_prob_property_raises (self, requires_api):
+    def test_batch_execute_invoking_prob_property_raises(self, requires_api):
         """Test invoking prob device property raises exception if circuit index not
-            previously set when multiple circuits are submitted in one job.
+        previously set when multiple circuits are submitted in one job.
         """
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
         with qml.tape.QuantumTape() as tape1:
@@ -376,8 +401,8 @@ def test_batch_execute_invoking_prob_property_raises (self, requires_api):
         with pytest.raises(
             CircuitIndexNotSetException,
             match="Because multiple circuits have been submitted in this job, the index of the circuit \
-you want to acces must be first set via the set_current_circuit_index device method."
-            ):
+you want to acces must be first set via the set_current_circuit_index device method.",
+        ):
             dev.prob
 
     def test_batch_execute_prob_property(self, requires_api):
@@ -394,8 +419,8 @@ def test_batch_execute_prob_property(self, requires_api):
         prob0 = dev.prob
         dev.set_current_circuit_index(1)
         prob1 = dev.prob
-        np.testing.assert_array_almost_equal(prob0, [0., 1.], decimal = 1)
-        np.testing.assert_array_almost_equal(prob1, [0.5, 0.5], decimal = 1)
+        np.testing.assert_array_almost_equal(prob0, [0.0, 1.0], decimal=1)
+        np.testing.assert_array_almost_equal(prob1, [0.5, 0.5], decimal=1)
 
     def test_batch_execute_counts(self, requires_api):
         """Test batch_execute method when computing counts."""
@@ -472,7 +497,10 @@ def mock_submit_job(*args):
         assert dev.job["input"]["qubits"] == 1
 
         assert len(dev.job["input"]["circuits"][0]) == 1
-        assert dev.job["input"]["circuits"][0]["circuit"][0] == {"gate": "x", "target": 0}
+        assert dev.job["input"]["circuits"][0]["circuit"][0] == {
+            "gate": "x",
+            "target": 0,
+        }
 
     def test_nonparametrized_tape_batch_submit(self, mocker):
         """Tests job attribute after single paulix tape, on batch submit."""

From ea651295cabdeddb44b0a203163142fc885513c6 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 25 Sep 2024 17:02:13 +0300
Subject: [PATCH 18/33] Remove unused includes.

---
 pennylane_ionq/device.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 53d2fab..9b75785 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -26,11 +26,7 @@
 from pennylane.devices import QubitDevice
 
 from pennylane.measurements import (
-    ClassicalShadowMP,
-    CountsMP,
     SampleMeasurement,
-    SampleMP,
-    ShadowExpvalMP,
     Shots,
     StateMeasurement,
 )

From 80f221e1fcb8d83dea82eb17d7ceb1e549ec600f Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 25 Sep 2024 17:43:03 +0300
Subject: [PATCH 19/33] Fix codefactor reported issues.

---
 pennylane_ionq/device.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 9b75785..5ece13d 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -184,6 +184,10 @@ def reset(self, circuits_array_length=1):
             )
 
     def set_current_circuit_index(self, circuit_index):
+        """Sets the index of the current circuit for which operations are applied upon.
+        In case of multiple circuits being submitted via batch_execute method
+        self._current_circuit_index tracks the index of the current circuit.
+        """
         self._current_circuit_index = circuit_index
 
     def batch_execute(self, circuits):
@@ -330,6 +334,12 @@ def apply(self, operations, **kwargs):
         self._submit_job()
 
     def _apply_operation(self, operation, circuit_index=0):
+        """Applies operations to the internal device state.
+
+        Args:
+            operation (.Operation): operation to apply on the device
+            circuit_index: index of the circuit to apply operation to
+        """
         name = operation.name
         wires = self.map_wires(operation.wires).tolist()
         gate = {"gate": self._operation_map[name]}

From 9a7ac798131983d49ea6c538424e556092317f09 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 25 Sep 2024 17:52:01 +0300
Subject: [PATCH 20/33] Adding doc string to method.

---
 pennylane_ionq/device.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 5ece13d..bf04712 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -308,6 +308,9 @@ def operations(self):
         return set(self._operation_map.keys())
 
     def apply(self, operations, **kwargs):
+        
+        "Implementation of abstract method apply method."
+
         self.reset()
         rotations = kwargs.pop("rotations", [])
 

From c7102f23415f366329e627795d6f37b0d4f1ab35 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 25 Sep 2024 17:56:02 +0300
Subject: [PATCH 21/33] Fix docstring.

---
 pennylane_ionq/device.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index bf04712..aadcdfe 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -309,7 +309,7 @@ def operations(self):
 
     def apply(self, operations, **kwargs):
         
-        "Implementation of abstract method apply method."
+        "Implementation of QubitDevice abstract method apply."
 
         self.reset()
         rotations = kwargs.pop("rotations", [])
@@ -371,6 +371,7 @@ def _apply_operation(self, operation, circuit_index=0):
         self.input["circuits"][circuit_index]["circuit"].append(gate)
 
     def _submit_job(self):
+
         job = Job(api_key=self.api_key)
 
         # send job for exection

From e8cd6e33e69fcf52f7f047f7b81936b3e8a17b8b Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 26 Sep 2024 10:38:34 +0300
Subject: [PATCH 22/33] Add test with shot vector.

---
 tests/test_device.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/tests/test_device.py b/tests/test_device.py
index 582a6ec..596d63c 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -27,7 +27,7 @@
     CircuitIndexNotSetException,
 )
 from pennylane_ionq.ops import GPI, GPI2, MS
-from pennylane.measurements import SampleMeasurement
+from pennylane.measurements import SampleMeasurement, ShotCopies
 
 FAKE_API_KEY = "ABC123"
 
@@ -291,6 +291,16 @@ def test_batch_execute_probabilities(self, requires_api):
             [0.0, 0.25, 0.25, 0.0, 0.0, 0.25, 0.25, 0.0],
         )
 
+    def test_batch_execute_probabilities_with_shot_vector(self, requires_api):
+        """Test batch_execute method with shot vector."""
+        dev = SimulatorDevice(wires=(0, 1, 2), gateset="native", shots=1024)
+        dev._shot_vector = (ShotCopies(1, 3),)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.probs(wires=[0])
+        results = dev.batch_execute([tape1])
+        assert(len(results[0]) == 3)
+
     def test_batch_execute_variance(self, requires_api):
         """Test batch_execute method when computing variance of an observable."""
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)

From 1c20c3e72dc2b24bd2363739160c4cace6bde7b7 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 26 Sep 2024 13:07:37 +0300
Subject: [PATCH 23/33] Add test with an observable that requires rotations for
 diagonalization.

---
 tests/test_device.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/tests/test_device.py b/tests/test_device.py
index 596d63c..41f777c 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -299,7 +299,7 @@ def test_batch_execute_probabilities_with_shot_vector(self, requires_api):
             GPI(0, wires=[0])
             qml.probs(wires=[0])
         results = dev.batch_execute([tape1])
-        assert(len(results[0]) == 3)
+        assert len(results[0]) == 3
 
     def test_batch_execute_variance(self, requires_api):
         """Test batch_execute method when computing variance of an observable."""
@@ -327,6 +327,21 @@ def test_batch_execute_expectation_value(self, requires_api):
         assert results[0] == pytest.approx(-1, abs=0.1)
         assert results[1] == pytest.approx(0, abs=0.1)
 
+    def test_batch_execute_expectation_value_with_diagonalization_rotations(
+        self, requires_api
+    ):
+        """Test batch_execute method when computing expectation value of an
+        observable that requires rotations for diagonalization."""
+        dev = SimulatorDevice(wires=(0,), gateset="qis", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            qml.Hadamard(0)
+            qml.expval(qml.PauliX(0))
+        with qml.tape.QuantumTape() as tape2:
+            qml.expval(qml.PauliX(0))
+        results = dev.batch_execute([tape1, tape2])
+        assert results[0] == pytest.approx(1, abs=0.1)
+        assert results[1] == pytest.approx(0, abs=0.1)
+
     def test_batch_execute_sample_raises(self, requires_api):
         """Test invoking sample() method raises exception if circuit index not
         previously set when multiple circuits are submitted in one job.

From 54809ad3935f338c9574638b865e3dcc55d21efd Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 26 Sep 2024 14:18:11 +0300
Subject: [PATCH 24/33] Adding unit tests for using pennylane tracker in
 batch_execute method. Adding unit tests for user warnings.

---
 tests/test_device.py | 46 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 46 insertions(+)

diff --git a/tests/test_device.py b/tests/test_device.py
index 41f777c..f2b00c5 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -248,6 +248,52 @@ def test_backend_initialization(self, backend):
         )
         assert dev.backend == backend
 
+    def test_recording_when_pennylane_tracker_active(self):
+        """Test recording device execution history via pennnylane tracker class."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        dev.tracker = qml.Tracker()
+        dev.tracker.active = True
+        dev.tracker.reset()
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.probs(wires=[0])
+        dev.batch_execute([tape1, tape1])
+        assert dev.tracker.history["executions"] == [1, 1]
+        assert dev.tracker.history["shots"] == [1024, 1024]
+        assert dev.tracker.history["batches"] == [1]
+        assert dev.tracker.history["batch_len"] == [2]
+        assert len(dev.tracker.history["resources"]) == 2
+        assert dev.tracker.history["resources"][0].num_wires == 1
+        assert dev.tracker.history["resources"][0].num_gates == 1
+        assert dev.tracker.history["resources"][0].depth == 1
+        assert dev.tracker.history["resources"][0].gate_types == {"GPI": 1}
+        assert dev.tracker.history["resources"][0].gate_sizes == {1: 1}
+        assert dev.tracker.history["resources"][0].shots.total_shots == 1024
+        assert len(dev.tracker.history["results"]) == 2
+
+    def test_not_recording_when_pennylane_tracker_not_active(self):
+        """Test recording device not executed when tracker is inactive."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        dev.tracker = qml.Tracker()
+        dev.tracker.active = False
+        dev.tracker.reset()
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.probs(wires=[0])
+        dev.batch_execute([tape1])
+        assert dev.tracker.history == {}
+
+    def test_warning_on_empty_circuit(self):
+        """Test warning are shown when circuit is empty."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            qml.probs(wires=[0])
+        with pytest.warns(
+            UserWarning,
+            match="Circuit is empty. Empty circuits return failures. Submitting anyway.",
+        ):
+            dev.batch_execute([tape1])
+
     def test_batch_execute_probabilities_raises(self, requires_api):
         """Test invoking probability() method raises exception if circuit index not
         previously set when multiple circuits are submitted in one job.

From 17f6f5221339f3a74b806a115271f9d882924587 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 26 Sep 2024 16:00:30 +0300
Subject: [PATCH 25/33] Add tests for logging in batch_execute.

---
 tests/test_device.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/tests/test_device.py b/tests/test_device.py
index f2b00c5..5e1eb1c 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 """Tests that plugin devices are accessible and integrate with PennyLane"""
 import json
+import logging
 import numpy as np
 import pennylane as qml
 import pytest
@@ -28,6 +29,7 @@
 )
 from pennylane_ionq.ops import GPI, GPI2, MS
 from pennylane.measurements import SampleMeasurement, ShotCopies
+from unittest import mock
 
 FAKE_API_KEY = "ABC123"
 
@@ -294,6 +296,23 @@ def test_warning_on_empty_circuit(self):
         ):
             dev.batch_execute([tape1])
 
+    @mock.patch("logging.Logger.isEnabledFor", return_value=True)
+    @mock.patch("logging.Logger.debug")
+    def test_batch_execute_logging_when_enabled(
+        self,
+        mock_logging_debug_method,
+        mock_logging_is_enabled_for_method,
+    ):
+        """Test logging invoked in batch_execute method."""
+        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
+        with qml.tape.QuantumTape() as tape1:
+            GPI(0, wires=[0])
+            qml.probs(wires=[0])
+        dev.batch_execute([tape1])
+        assert mock_logging_is_enabled_for_method.called
+        assert mock_logging_is_enabled_for_method.call_args[0][0] == logging.DEBUG
+        mock_logging_debug_method.assert_called()
+
     def test_batch_execute_probabilities_raises(self, requires_api):
         """Test invoking probability() method raises exception if circuit index not
         previously set when multiple circuits are submitted in one job.

From 0c30e00f51d460cf4a198c19b056a482cdf94695 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 26 Sep 2024 16:01:47 +0300
Subject: [PATCH 26/33] Run black.

---
 pennylane_ionq/device.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index aadcdfe..b3c59f8 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -308,7 +308,7 @@ def operations(self):
         return set(self._operation_map.keys())
 
     def apply(self, operations, **kwargs):
-        
+
         "Implementation of QubitDevice abstract method apply."
 
         self.reset()

From 87516abfd2dac3418e944a584acb3c82e5ba1aed Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 2 Oct 2024 16:02:00 +0300
Subject: [PATCH 27/33] Remove method override.

---
 pennylane_ionq/device.py | 33 ---------------------------------
 1 file changed, 33 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index b3c59f8..451a88c 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -502,39 +502,6 @@ def sample(
         self._samples = None
         return result
 
-    def _measure(
-        self,
-        measurement: Union[SampleMeasurement, StateMeasurement],
-        shot_range=None,
-        bin_size=None,
-    ):
-        """Compute the corresponding measurement process depending on ``shots`` and the measurement
-        type. Overwrites the method in QubitDevice class to accomodate batch submission of circuits.
-        When invoking this method after submitting circuits using any other method than QubitDevice.execute()
-        method, current circuit index should be set via set_current_circuit_index() method pointing to index
-        of the circuit which is targeted out of one or potentially several circuits submitted in a batch.
-        This is not very good design but the alternative would be to create huge amounts of code duplication
-        with respect to the Pennylane base code.
-
-        Args:
-            measurement (Union[SampleMeasurement, StateMeasurement]): measurement process
-            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
-                to use. If not specified, all samples are used.
-            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
-                returns the measurement statistic separately over each bin. If not
-                provided, the entire shot range is treated as a single bin.
-
-        Raises:
-            ValueError: if the measurement cannot be computed
-
-        Returns:
-            Union[float, dict, list[float]]: result of the measurement
-        """
-        self._samples = self.generate_samples()
-        result = super()._measure(measurement, shot_range, bin_size)
-        self._samples = None
-        return result
-
     def probability(self, wires=None, shot_range=None, bin_size=None):
         wires = wires or self.wires
 

From 75d3a3c8208355f309c316cd34d829e88715c4e1 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 3 Oct 2024 11:23:10 +0300
Subject: [PATCH 28/33] Implement review comments.'

---
 pennylane_ionq/device.py | 52 ++++++++++++-------------------
 tests/test_device.py     | 67 +++++++++++-----------------------------
 2 files changed, 38 insertions(+), 81 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 451a88c..cfd6fbd 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -80,8 +80,10 @@ class CircuitIndexNotSetException(Exception):
     """
 
     def __init__(self):
-        self.message = "Because multiple circuits have been submitted in this job, the index of the circuit \
-you want to acces must be first set via the set_current_circuit_index device method."
+        self.message = (
+            "Because multiple circuits have been submitted in this job, the index of the circuit "
+            "you want to access must be first set via the set_current_circuit_index device method."
+        )
         super().__init__(self.message)
 
 
@@ -141,9 +143,7 @@ def __init__(
         sharpen=False,
     ):
         if shots is None:
-            raise ValueError(
-                "The ionq device does not support analytic expectation values."
-            )
+            raise ValueError("The ionq device does not support analytic expectation values.")
 
         super().__init__(wires=wires, shots=shots)
         self._current_circuit_index = None
@@ -207,8 +207,7 @@ def batch_execute(self, circuits):
                 """Entry with args=(circuits=%s) called by=%s""",
                 circuits,
                 "::L".join(
-                    str(i)
-                    for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]
+                    str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]
                 ),
             )
 
@@ -247,9 +246,7 @@ def batch_execute(self, circuits):
 
         if self.tracker.active:
             for circuit in circuits:
-                shots_from_dev = (
-                    self._shots if not self.shot_vector else self._raw_shot_sequence
-                )
+                shots_from_dev = self._shots if not self.shot_vector else self._raw_shot_sequence
                 tape_resources = circuit.specs["resources"]
 
                 resources = Resources(  # temporary until shots get updated on tape !
@@ -279,9 +276,7 @@ def batch_apply(self, operations, circuit_index, **kwargs):
         rotations = kwargs.pop("rotations", [])
 
         if len(operations) == 0 and len(rotations) == 0:
-            warnings.warn(
-                "Circuit is empty. Empty circuits return failures. Submitting anyway."
-            )
+            warnings.warn("Circuit is empty. Empty circuits return failures. Submitting anyway.")
 
         for i, operation in enumerate(operations):
             if i > 0 and operation.name in {
@@ -308,16 +303,13 @@ def operations(self):
         return set(self._operation_map.keys())
 
     def apply(self, operations, **kwargs):
-
-        "Implementation of QubitDevice abstract method apply."
+        """Implementation of QubitDevice abstract method apply."""
 
         self.reset()
         rotations = kwargs.pop("rotations", [])
 
         if len(operations) == 0 and len(rotations) == 0:
-            warnings.warn(
-                "Circuit is empty. Empty circuits return failures. Submitting anyway."
-            )
+            warnings.warn("Circuit is empty. Empty circuits return failures. Submitting anyway.")
 
         for i, operation in enumerate(operations):
             if i > 0 and operation.name in {
@@ -421,9 +413,7 @@ def prob(self):
         # The IonQ API returns basis states using little-endian ordering.
         # Here, we rearrange the states to match the big-endian ordering
         # expected by PennyLane.
-        basis_states = (
-            int(bin(int(k))[2:].rjust(self.num_wires, "0")[::-1], 2) for k in histogram
-        )
+        basis_states = (int(bin(int(k))[2:].rjust(self.num_wires, "0")[::-1], 2) for k in histogram)
         idx = np.fromiter(basis_states, dtype=int)
 
         # convert the sparse probs into a probability array
@@ -438,10 +428,10 @@ def prob(self):
 
     def estimate_probability(self, wires=None, shot_range=None, bin_size=None):
         """Return the estimated probability of each computational basis state
-        using the generated samples. Overwrites the method in QubitDevice class
-        to accomodate batch submission of circuits. When invoking this method after
-        submitting circuits using any other method than QubitDevice.execute()
-        method, current circuit index should be set via set_current_circuit_index()
+        using the generated samples. Overwrites the method in ``QubitDevice`` class
+        to accommodate batch submission of circuits. When invoking this method after
+        submitting circuits using any other method than ``QubitDevice.execute()``
+        method, current circuit index should be set via ``set_current_circuit_index()``
         method to the circuit which is targeted out of one or potentially several
         circuits submitted in a batch. This is not very good design but the
         alternative would be to create huge amounts of code duplication with respect
@@ -471,10 +461,10 @@ def sample(
         bin_size=None,
         counts=False,
     ):
-        """Return samples of an observable. Overwrites the method in QubitDevice class
-        to accomodate batch submission of circuits. When invoking this method after submitting
-        circuits using any other method than QubitDevice.execute() method, current circuit
-        index should be set via set_current_circuit_index() method pointing to index of the
+        """Return samples of an observable. Overwrites the method in ``QubitDevice`` class
+        to accommodate batch submission of circuits. When invoking this method after submitting
+        circuits using any other method than ``QubitDevice.execute()`` method, current circuit
+        index should be set via ``set_current_circuit_index()`` method pointing to index of the
         circuit which is targeted out of one or potentially several circuits submitted in a batch.
         This is not very good design but the alternative would be to create huge amounts of
         code duplication with respect to the Pennylane base code.
@@ -508,9 +498,7 @@ def probability(self, wires=None, shot_range=None, bin_size=None):
         if shot_range is None and bin_size is None:
             return self.marginal_prob(self.prob, wires)
 
-        return self.estimate_probability(
-            wires=wires, shot_range=shot_range, bin_size=bin_size
-        )
+        return self.estimate_probability(wires=wires, shot_range=shot_range, bin_size=bin_size)
 
 
 class SimulatorDevice(IonQDevice):
diff --git a/tests/test_device.py b/tests/test_device.py
index 5e1eb1c..8561bdc 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -59,15 +59,11 @@ def test_generate_samples_qpu_device(self, wires, histogram):
 
         unique_outcomes1 = np.unique(sample1, axis=0)
         unique_outcomes2 = np.unique(sample2, axis=0)
-        assert np.all(
-            unique_outcomes1 == unique_outcomes2
-        )  # possible outcomes are the same
+        assert np.all(unique_outcomes1 == unique_outcomes2)  # possible outcomes are the same
 
         sorted_outcomes1 = np.sort(sample1, axis=0)
         sorted_outcomes2 = np.sort(sample2, axis=0)
-        assert np.all(
-            sorted_outcomes1 == sorted_outcomes2
-        )  # set of outcomes is the same
+        assert np.all(sorted_outcomes1 == sorted_outcomes2)  # set of outcomes is the same
 
 
 class TestDeviceIntegration:
@@ -107,9 +103,7 @@ def test_failedcircuit(self, monkeypatch):
         monkeypatch.setattr(
             requests, "post", lambda url, timeout, data, headers: (url, data, headers)
         )
-        monkeypatch.setattr(
-            ResourceManager, "handle_response", lambda self, response: None
-        )
+        monkeypatch.setattr(ResourceManager, "handle_response", lambda self, response: None)
         monkeypatch.setattr(Job, "is_complete", False)
         monkeypatch.setattr(Job, "is_failed", True)
 
@@ -124,17 +118,13 @@ def test_shots(self, shots, monkeypatch, mocker, tol):
         monkeypatch.setattr(
             requests, "post", lambda url, timeout, data, headers: (url, data, headers)
         )
-        monkeypatch.setattr(
-            ResourceManager, "handle_response", lambda self, response: None
-        )
+        monkeypatch.setattr(ResourceManager, "handle_response", lambda self, response: None)
         monkeypatch.setattr(Job, "is_complete", True)
 
         def fake_response(self, resource_id=None, params=None):
             """Return fake response data"""
             fake_json = {"0": 1}
-            setattr(
-                self.resource, "data", type("data", tuple(), {"value": fake_json})()
-            )
+            setattr(self.resource, "data", type("data", tuple(), {"value": fake_json})())
 
         monkeypatch.setattr(ResourceManager, "get", fake_response)
 
@@ -150,26 +140,20 @@ def circuit():
         circuit()
         assert json.loads(spy.call_args[1]["data"])["shots"] == shots
 
-    @pytest.mark.parametrize(
-        "error_mitigation", [None, {"debias": True}, {"debias": False}]
-    )
+    @pytest.mark.parametrize("error_mitigation", [None, {"debias": True}, {"debias": False}])
     def test_error_mitigation(self, error_mitigation, monkeypatch, mocker):
         """Test that shots are correctly specified when submitting a job to the API."""
 
         monkeypatch.setattr(
             requests, "post", lambda url, timeout, data, headers: (url, data, headers)
         )
-        monkeypatch.setattr(
-            ResourceManager, "handle_response", lambda self, response: None
-        )
+        monkeypatch.setattr(ResourceManager, "handle_response", lambda self, response: None)
         monkeypatch.setattr(Job, "is_complete", True)
 
         def fake_response(self, resource_id=None, params=None):
             """Return fake response data"""
             fake_json = {"0": 1}
-            setattr(
-                self.resource, "data", type("data", tuple(), {"value": fake_json})()
-            )
+            setattr(self.resource, "data", type("data", tuple(), {"value": fake_json})())
 
         monkeypatch.setattr(ResourceManager, "get", fake_response)
 
@@ -190,10 +174,7 @@ def circuit():
         spy = mocker.spy(requests, "post")
         circuit()
         if error_mitigation is not None:
-            assert (
-                json.loads(spy.call_args[1]["data"])["error_mitigation"]
-                == error_mitigation
-            )
+            assert json.loads(spy.call_args[1]["data"])["error_mitigation"] == error_mitigation
         else:
             with pytest.raises(KeyError, match="error_mitigation"):
                 json.loads(spy.call_args[1]["data"])["error_mitigation"]
@@ -237,9 +218,7 @@ def test_prob_no_results(self, d):
         dev = qml.device(d, wires=1, shots=1)
         assert dev.prob is None
 
-    @pytest.mark.parametrize(
-        "backend", ["harmony", "aria-1", "aria-2", "forte-1", None]
-    )
+    @pytest.mark.parametrize("backend", ["harmony", "aria-1", "aria-2", "forte-1", None])
     def test_backend_initialization(self, backend):
         """Test that the device initializes with the correct backend."""
         dev = qml.device(
@@ -250,7 +229,7 @@ def test_backend_initialization(self, backend):
         )
         assert dev.backend == backend
 
-    def test_recording_when_pennylane_tracker_active(self):
+    def test_recording_when_pennylane_tracker_active(self, requires_api):
         """Test recording device execution history via pennnylane tracker class."""
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
         dev.tracker = qml.Tracker()
@@ -273,7 +252,7 @@ def test_recording_when_pennylane_tracker_active(self):
         assert dev.tracker.history["resources"][0].shots.total_shots == 1024
         assert len(dev.tracker.history["results"]) == 2
 
-    def test_not_recording_when_pennylane_tracker_not_active(self):
+    def test_not_recording_when_pennylane_tracker_not_active(self, requires_api):
         """Test recording device not executed when tracker is inactive."""
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
         dev.tracker = qml.Tracker()
@@ -285,7 +264,7 @@ def test_not_recording_when_pennylane_tracker_not_active(self):
         dev.batch_execute([tape1])
         assert dev.tracker.history == {}
 
-    def test_warning_on_empty_circuit(self):
+    def test_warning_on_empty_circuit(self, requires_api):
         """Test warning are shown when circuit is empty."""
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
         with qml.tape.QuantumTape() as tape1:
@@ -299,9 +278,7 @@ def test_warning_on_empty_circuit(self):
     @mock.patch("logging.Logger.isEnabledFor", return_value=True)
     @mock.patch("logging.Logger.debug")
     def test_batch_execute_logging_when_enabled(
-        self,
-        mock_logging_debug_method,
-        mock_logging_is_enabled_for_method,
+        self, mock_logging_debug_method, mock_logging_is_enabled_for_method, requires_api
     ):
         """Test logging invoked in batch_execute method."""
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
@@ -392,9 +369,7 @@ def test_batch_execute_expectation_value(self, requires_api):
         assert results[0] == pytest.approx(-1, abs=0.1)
         assert results[1] == pytest.approx(0, abs=0.1)
 
-    def test_batch_execute_expectation_value_with_diagonalization_rotations(
-        self, requires_api
-    ):
+    def test_batch_execute_expectation_value_with_diagonalization_rotations(self, requires_api):
         """Test batch_execute method when computing expectation value of an
         observable that requires rotations for diagonalization."""
         dev = SimulatorDevice(wires=(0,), gateset="qis", shots=1024)
@@ -534,13 +509,9 @@ def __init__(self, state: str):
                 wires = list(range(len(state)))
                 super().__init__(wires=wires)
 
-            def process_samples(
-                self, samples, wire_order, shot_range=None, bin_size=None
-            ):
+            def process_samples(self, samples, wire_order, shot_range=None, bin_size=None):
                 counts_mp = qml.counts(wires=self._wires)
-                counts = counts_mp.process_samples(
-                    samples, wire_order, shot_range, bin_size
-                )
+                counts = counts_mp.process_samples(samples, wire_order, shot_range, bin_size)
                 return float(counts.get(self.state, 0))
 
             def process_counts(self, counts, wire_order):
@@ -735,9 +706,7 @@ class StopExecute(Exception):
         def mock_submit_job(*args):
             raise StopExecute()
 
-        mocker.patch(
-            "pennylane_ionq.device.SimulatorDevice._submit_job", mock_submit_job
-        )
+        mocker.patch("pennylane_ionq.device.SimulatorDevice._submit_job", mock_submit_job)
         dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
 
         with qml.tape.QuantumTape() as tape1:

From 0b92fa42c0cb338a3725ba56f27c870cf8ecbbfa Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 3 Oct 2024 11:51:52 +0300
Subject: [PATCH 29/33] Implement more review comments.'

---
 pennylane_ionq/device.py | 66 ------------------------------------
 tests/test_device.py     | 73 ++--------------------------------------
 2 files changed, 2 insertions(+), 137 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index cfd6fbd..a0483f4 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -426,72 +426,6 @@ def prob(self):
 
         return prob_array
 
-    def estimate_probability(self, wires=None, shot_range=None, bin_size=None):
-        """Return the estimated probability of each computational basis state
-        using the generated samples. Overwrites the method in ``QubitDevice`` class
-        to accommodate batch submission of circuits. When invoking this method after
-        submitting circuits using any other method than ``QubitDevice.execute()``
-        method, current circuit index should be set via ``set_current_circuit_index()``
-        method to the circuit which is targeted out of one or potentially several
-        circuits submitted in a batch. This is not very good design but the
-        alternative would be to create huge amounts of code duplication with respect
-        to the Pennylane base code.
-
-        Args:
-            wires (Iterable[Number, str], Number, str, Wires): wires to calculate
-                marginal probabilities for. Wires not provided are traced out of one or the system.
-            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
-                to use. If not specified, all samples are used.
-            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
-                returns the measurement statistic separately over each bin. If not
-                provided, the entire shot range is treated as a single bin.
-
-        Returns:
-            array[float]: list of the probabilities
-        """
-        self._samples = self.generate_samples()
-        result = super().estimate_probability(wires, shot_range, bin_size)
-        self._samples = None
-        return result
-
-    def sample(
-        self,
-        observable,
-        shot_range=None,
-        bin_size=None,
-        counts=False,
-    ):
-        """Return samples of an observable. Overwrites the method in ``QubitDevice`` class
-        to accommodate batch submission of circuits. When invoking this method after submitting
-        circuits using any other method than ``QubitDevice.execute()`` method, current circuit
-        index should be set via ``set_current_circuit_index()`` method pointing to index of the
-        circuit which is targeted out of one or potentially several circuits submitted in a batch.
-        This is not very good design but the alternative would be to create huge amounts of
-        code duplication with respect to the Pennylane base code.
-
-        Args:
-            observable (Observable): the observable to sample
-            shot_range (tuple[int]): 2-tuple of integers specifying the range of samples
-                to use. If not specified, all samples are used.
-            bin_size (int): Divides the shot range into bins of size ``bin_size``, and
-                returns the measurement statistic separately over each bin. If not
-                provided, the entire shot range is treated as a single bin.
-            counts (bool): whether counts (``True``) or raw samples (``False``)
-                should be returned
-
-        Raises:
-            EigvalsUndefinedError: if no information is available about the
-                eigenvalues of the observable
-
-        Returns:
-            Union[array[float], dict, list[dict]]: samples in an array of
-            dimension ``(shots,)`` or counts
-        """
-        self._samples = self.generate_samples()
-        result = super().sample(observable, shot_range, bin_size, counts)
-        self._samples = None
-        return result
-
     def probability(self, wires=None, shot_range=None, bin_size=None):
         wires = wires or self.wires
 
diff --git a/tests/test_device.py b/tests/test_device.py
index 8561bdc..26fa8c4 100755
--- a/tests/test_device.py
+++ b/tests/test_device.py
@@ -302,7 +302,7 @@ def test_batch_execute_probabilities_raises(self, requires_api):
         with pytest.raises(
             CircuitIndexNotSetException,
             match="Because multiple circuits have been submitted in this job, the index of the circuit \
-you want to acces must be first set via the set_current_circuit_index device method.",
+you want to access must be first set via the set_current_circuit_index device method.",
         ):
             dev.probability()
 
@@ -382,75 +382,6 @@ def test_batch_execute_expectation_value_with_diagonalization_rotations(self, re
         assert results[0] == pytest.approx(1, abs=0.1)
         assert results[1] == pytest.approx(0, abs=0.1)
 
-    def test_batch_execute_sample_raises(self, requires_api):
-        """Test invoking sample() method raises exception if circuit index not
-        previously set when multiple circuits are submitted in one job.
-        """
-        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
-        with qml.tape.QuantumTape() as tape1:
-            GPI(0.5, wires=[0])
-            qml.probs(wires=[0])
-        dev.batch_execute([tape1, tape1])
-        with pytest.raises(
-            CircuitIndexNotSetException,
-            match="Because multiple circuits have been submitted in this job, the index of the circuit \
-you want to acces must be first set via the set_current_circuit_index device method.",
-        ):
-            dev.sample(qml.PauliZ(0))
-
-    def test_batch_execute_sample(self, requires_api):
-        """Test batch_execute method when sampling."""
-        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
-        with qml.tape.QuantumTape() as tape1:
-            GPI(0, wires=[0])
-            qml.sample(qml.PauliZ(0))
-        with qml.tape.QuantumTape() as tape2:
-            GPI2(0, wires=[0])
-            qml.sample(qml.PauliZ(0))
-        results = dev.batch_execute([tape1, tape2])
-        dev.set_current_circuit_index(0)
-        results0 = dev.sample(qml.PauliZ(0))
-        dev.set_current_circuit_index(1)
-        results1 = dev.sample(qml.PauliZ(0))
-        assert set(results[0]) == set(results0) == {-1}
-        assert set(results[1]) == set(results1) == {-1, 1}
-
-    def test_batch_execute_estimate_probability_raises(self, requires_api):
-        """Test invoking estimate_probability() method raises exception if circuit index not
-        previously set when multiple circuits are submitted in one job.
-        """
-        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
-        with qml.tape.QuantumTape() as tape1:
-            GPI(0, wires=[0])
-            qml.probs(wires=[0])
-        with qml.tape.QuantumTape() as tape2:
-            GPI2(0, wires=[0])
-            qml.probs(wires=[0])
-        dev.batch_execute([tape1, tape2])
-        with pytest.raises(
-            CircuitIndexNotSetException,
-            match="Because multiple circuits have been submitted in this job, the index of the circuit \
-you want to acces must be first set via the set_current_circuit_index device method.",
-        ):
-            dev.estimate_probability()
-
-    def test_batch_execute_estimate_probability(self, requires_api):
-        """Test batch_execute method with invoking estimate_probability method."""
-        dev = SimulatorDevice(wires=(0,), gateset="native", shots=1024)
-        with qml.tape.QuantumTape() as tape1:
-            GPI(0, wires=[0])
-            qml.sample(qml.PauliZ(0))
-        with qml.tape.QuantumTape() as tape2:
-            GPI2(0, wires=[0])
-            qml.sample(qml.PauliZ(0))
-        dev.batch_execute([tape1, tape2])
-        dev.set_current_circuit_index(0)
-        prob0 = dev.estimate_probability()
-        dev.set_current_circuit_index(1)
-        prob1 = dev.estimate_probability()
-        np.testing.assert_array_almost_equal(prob0, [0.0, 1.0], decimal=1)
-        np.testing.assert_array_almost_equal(prob1, [0.5, 0.5], decimal=1)
-
     def test_batch_execute_invoking_prob_property_raises(self, requires_api):
         """Test invoking prob device property raises exception if circuit index not
         previously set when multiple circuits are submitted in one job.
@@ -466,7 +397,7 @@ def test_batch_execute_invoking_prob_property_raises(self, requires_api):
         with pytest.raises(
             CircuitIndexNotSetException,
             match="Because multiple circuits have been submitted in this job, the index of the circuit \
-you want to acces must be first set via the set_current_circuit_index device method.",
+you want to access must be first set via the set_current_circuit_index device method.",
         ):
             dev.prob
 

From 142c4cf646e9faba492fbd2fa86f875704634f0b Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 3 Oct 2024 12:25:07 +0300
Subject: [PATCH 30/33] Remove unused imports.

---
 pennylane_ionq/device.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index a0483f4..8eead26 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -26,9 +26,7 @@
 from pennylane.devices import QubitDevice
 
 from pennylane.measurements import (
-    SampleMeasurement,
     Shots,
-    StateMeasurement,
 )
 from pennylane.resource import Resources
 

From 0aafd7950024d74577a9abd76092d97e582d8734 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Thu, 3 Oct 2024 13:18:20 +0300
Subject: [PATCH 31/33] Remove unused imports.

---
 pennylane_ionq/device.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 8eead26..0be1e63 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -16,7 +16,6 @@
 """
 import inspect
 import logging
-from typing import Union
 import warnings
 from time import sleep
 

From 28babc1f39a858e425963f04f9d6b1be4cc3ecc9 Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 9 Oct 2024 11:02:28 +0300
Subject: [PATCH 32/33] Remove checks for BasisState, QubitStateVector and
 StatePrep because these are not supported by IonQ devices.

---
 pennylane_ionq/device.py | 19 +------------------
 1 file changed, 1 insertion(+), 18 deletions(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 0be1e63..1149b72 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -21,7 +21,6 @@
 
 import numpy as np
 
-from pennylane import DeviceError
 from pennylane.devices import QubitDevice
 
 from pennylane.measurements import (
@@ -200,7 +199,7 @@ def batch_execute(self, circuits):
             list[array[float]]: list of measured value(s)
         """
         if logger.isEnabledFor(logging.DEBUG):
-            logger.debug(
+            logger.debug(   # pragma: no cover
                 """Entry with args=(circuits=%s) called by=%s""",
                 circuits,
                 "::L".join(
@@ -276,14 +275,6 @@ def batch_apply(self, operations, circuit_index, **kwargs):
             warnings.warn("Circuit is empty. Empty circuits return failures. Submitting anyway.")
 
         for i, operation in enumerate(operations):
-            if i > 0 and operation.name in {
-                "BasisState",
-                "QubitStateVector",
-                "StatePrep",
-            }:
-                raise DeviceError(
-                    f"The operation {operation.name} is only supported at the beginning of a circuit."
-                )
             self._apply_operation(operation, circuit_index)
 
         # diagonalize observables
@@ -309,14 +300,6 @@ def apply(self, operations, **kwargs):
             warnings.warn("Circuit is empty. Empty circuits return failures. Submitting anyway.")
 
         for i, operation in enumerate(operations):
-            if i > 0 and operation.name in {
-                "BasisState",
-                "QubitStateVector",
-                "StatePrep",
-            }:
-                raise DeviceError(
-                    f"The operation {operation.name} is only supported at the beginning of a circuit."
-                )
             self._apply_operation(operation)
 
         # diagonalize observables

From eae3bcb9d77d807c4cfd04e3d54952c7ed6cb20f Mon Sep 17 00:00:00 2001
From: Radu Marginean <radu.marg@gmail.com>
Date: Wed, 9 Oct 2024 17:50:36 +0300
Subject: [PATCH 33/33] Fix code formatting.

---
 pennylane_ionq/device.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane_ionq/device.py b/pennylane_ionq/device.py
index 1149b72..926c6a7 100644
--- a/pennylane_ionq/device.py
+++ b/pennylane_ionq/device.py
@@ -199,7 +199,7 @@ def batch_execute(self, circuits):
             list[array[float]]: list of measured value(s)
         """
         if logger.isEnabledFor(logging.DEBUG):
-            logger.debug(   # pragma: no cover
+            logger.debug(  # pragma: no cover
                 """Entry with args=(circuits=%s) called by=%s""",
                 circuits,
                 "::L".join(