draft

pennylane/_qubit_device.py

@@ -263,8 +263,6 @@ def execute(self, circuit, **kwargs):
         if self.shots is not None or circuit.is_sampled:
             self._samples = self.generate_samples()
 
-        multiple_sampled_jobs = circuit.is_sampled and self._has_partitioned_shots()
-
         # compute the required statistics
         if not self.analytic and self._shot_vector is not None:
 
@@ -277,24 +275,38 @@ def execute(self, circuit, **kwargs):
                     circuit.observables, shot_range=[s1, s2], bin_size=shot_tuple.shots
                 )
 
+                counts_shot_vector = isinstance(r[0], list) and isinstance(r[0][0], dict)
+
                 if qml.math._multi_dispatch(r) == "jax":  # pylint: disable=protected-access
                     r = r[0]
                 elif not isinstance(r[0], dict):
-                    # Measurement types except for Counts
-                    r = qml.math.squeeze(r)
-                if isinstance(r, (np.ndarray, list)) and r.shape and isinstance(r[0], dict):
+                    # TODO: what if multi-measure and dict is not the first?
+
+                    if counts_shot_vector:
+                        r = r
+                    else:
+                        # Measurement types except for Counts
+                        r = qml.math.squeeze(r)
+
+                r_is_seq = (isinstance(r, np.ndarray) and r.shape or isinstance(r, list))
+                has_dict = r_is_seq and (isinstance(r[0], dict) or (isinstance(r[0], list) and isinstance(r[0][0], dict)))
+
+                #if ((isinstance(r, np.ndarray) and r.shape) or isinstance(r, list)) and isinstance(r[0], dict):
+                if isinstance(r, (np.ndarray)) and r.shape and isinstance(r[0], dict):
                     # This happens when measurement type is Counts
                     results.append(r)
+
+                elif counts_shot_vector:
+                    # TODO: if shot vector: good that we extend?
+                    results.extend(r)
+
                 elif shot_tuple.copies > 1:
                     results.extend(r.T)
                 else:
                     results.append(r.T)
 
                 s1 = s2
-
-            if not multiple_sampled_jobs:
-                # Can only stack single element outputs
-                results = qml.math.stack(results)
+            print('final results after appending', results, type(results))
 
         else:
             results = self.statistics(circuit.observables)
@@ -306,10 +318,19 @@ def execute(self, circuit, **kwargs):
             if len(circuit.measurements) == 1:
                 if circuit.measurements[0].return_type is qml.measurements.State:
                     # State: assumed to only be allowed if it's the only measurement
-                    results = self._asarray(results, dtype=self.C_DTYPE)
+
+                    if self._has_partitioned_shots():
+                        #TODO: revisit finite shots and State: do we disallow it?
+                        results = tuple([self._asarray(r, dtype=self.C_DTYPE) for r in results])
+                    else:
+                        results = self._asarray(results, dtype=self.C_DTYPE)
+
                 elif circuit.measurements[0].return_type is not qml.measurements.Counts:
                     # Measurements with expval, var or probs
-                    results = self._asarray(results, dtype=self.R_DTYPE)
+                    if self._has_partitioned_shots():
+                        results = tuple([self._asarray(r, dtype=self.R_DTYPE) for r in results])
+                    else:
+                        results = self._asarray(results, dtype=self.R_DTYPE)
 
             elif all(
                 ret in (qml.measurements.Expectation, qml.measurements.Variance)
@@ -325,7 +346,8 @@ def execute(self, circuit, **kwargs):
 
             results = self._asarray(results)
         else:
-            results = tuple(self._asarray(r) for r in results)
+            results = tuple(r for r in results)
+        print('still: final results after appending?', results, type(results))
 
         # increment counter for number of executions of qubit device
         self._num_executions += 1
@@ -487,9 +509,9 @@ def statistics(self, observables, shot_range=None, bin_size=None):
                 )
 
             elif obs.return_type is Counts:
-                results.append(
-                    self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=True)
-                )
+                r = self.sample(obs, shot_range=shot_range, bin_size=bin_size, counts=True)
+                print("raw result: ", r)
+                results.append(r)
 
             elif obs.return_type is Probability:
                 results.append(

pennylane/interfaces/autograd.py

@@ -109,6 +109,7 @@ def _execute(
     with qml.tape.Unwrap(*tapes):
         res, jacs = execute_fn(tapes, **gradient_kwargs)
 
+    print("Result now: ", res, type(res))
     for i, r in enumerate(res):
 
         if isinstance(r, np.ndarray):
@@ -125,6 +126,8 @@ def _execute(
 
         elif isinstance(res[i], tuple):
             res[i] = tuple(np.tensor(r) for r in res[i])
+        elif tapes[i]._qfunc_output.return_type is qml.measurements.Counts:
+            continue
 
         else:
             res[i] = qml.math.toarray(res[i])

pennylane/interfaces/execution.py

@@ -316,7 +316,10 @@ def cost_fn(params, x):
 
     if isinstance(cache, bool) and cache:
         # cache=True: create a LRUCache object
-        cache = LRUCache(maxsize=cachesize, getsizeof=lambda x: qml.math.shape(x)[0])
+
+        # TODO: changed from qml.math.shape to len because shape converts to
+        # arrays under the hood -> may create ragged array (best solution?)
+        cache = LRUCache(maxsize=cachesize, getsizeof=lambda x: len(x))
         setattr(cache, "_persistent_cache", False)
 
     batch_execute = set_shots(device, override_shots)(device.batch_execute)

pennylane/qnode.py

@@ -625,6 +625,7 @@ def __call__(self, *args, **kwargs):
             override_shots=override_shots,
             **self.execute_kwargs,
         )
+        print("result of execute: ", res)
 
         if autograd.isinstance(res, (tuple, list)) and len(res) == 1:
             # If a device batch transform was applied, we need to 'unpack'
@@ -646,14 +647,11 @@ def __call__(self, *args, **kwargs):
 
         self._update_original_device()
 
-        if isinstance(self._qfunc_output, Sequence) or (
-            self.tape.is_sampled and self.device._has_partitioned_shots()
-        ):
+        if isinstance(self._qfunc_output, Sequence) or self.device._has_partitioned_shots() or self._qfunc_output.return_type is qml.measurements.Counts:
+            # Shot vectors outputs are also sequences
             return res
-        if self._qfunc_output.return_type is qml.measurements.Counts:
-            # return a dictionary with counts not as a single-element array
-            return res[0]
 
+        # Squeeze arraylike outputs
         return qml.math.squeeze(res)