Repository: PennyLaneAI/pennylane
Pull request #1897 (2021-11-15 22:27:58)

`ExpvalCost` when `optimize=True` with shots batch
**Context:**
`ExpvalCost` returns wrong values when `optimize=True` because of the shape of results due to shots (e.g (8000,5)).

**Description of the Change:**
Change the function to take into account the shape of the results.

**Benefits:**
It now returns the correct values.

**Related GitHub Issues:**
Closes https://github.com/PennyLaneAI/pennylane/issues/1480

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#1897](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (e642931) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/927c4c67a1e1f9c79a5d7d4e89c496fac418657f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (927c4c6) will **decrease** coverage by `58.04%`.
> The diff coverage is `0.00%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/graphs/tree.svg?width=650&height=150&src=pr&token=9S3YzlUgGk&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@             Coverage Diff             @@
##           master    #1897       +/-   ##
===========================================
- Coverage   98.98%   40.93%   -58.05%     
===========================================
  Files         219      219               
  Lines       16513    16525       +12     
===========================================
- Hits        16345     6765     -9580     
- Misses        168     9760     +9592     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/vqe/vqe.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3ZxZS92cWUucHk=) | `42.59% <0.00%> (-57.41%)` | :arrow_down: |
| [pennylane/qnn/keras.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9rZXJhcy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/qnn/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi90b3JjaC5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvdG9yY2gucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvamF4LnB5) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/beta/devices/\_\_init\_\_.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2JldGEvZGV2aWNlcy9fX2luaXRfXy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/beta/devices/numpy\_ops.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2JldGEvZGV2aWNlcy9udW1weV9vcHMucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/batch/torch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvdG9yY2gucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/batch/autograd.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvYXV0b2dyYWQucHk=) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| [pennylane/interfaces/batch/tensorflow.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvYmF0Y2gvdGVuc29yZmxvdy5weQ==) | `0.00% <0.00%> (-100.00%)` | :arrow_down: |
| ... and [176 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [927c4c6...e642931](https://codecov.io/gh/PennyLaneAI/pennylane/pull/1897?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 4bb1354b0..458b97bd2 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -58,6 +58,10 @@
 
 <h3>Bug fixes</h3>
 
+* `ExpvalCost` now returns corrects results shape when `optimize=True` with 
+  shots batch.
+  [(#1897)](https://github.com/PennyLaneAI/pennylane/pull/1897)
+  
 * `qml.circuit_drawer.MPLDrawer` was slightly modified to work with
   matplotlib version 3.5.
   [(#1899)](https://github.com/PennyLaneAI/pennylane/pull/1899)
@@ -76,5 +80,5 @@
 
 This release contains contributions from (in alphabetical order): 
 
-Guillermo Alonso-Linaje, Olivia Di Matteo, Jalani Kanem, Shumpei Kobayashi, Christina Lee, Alejandro Montanez, Maria Schuld, Jay Soni
-
+Guillermo Alonso-Linaje, Olivia Di Matteo, Jalani Kanem, Shumpei Kobayashi, Christina Lee, Alejandro Montanez,
+Romain Moyard, Maria Schuld, Jay Soni
diff --git a/pennylane/vqe/vqe.py b/pennylane/vqe/vqe.py
index 4638d6560..5e0317de4 100644
--- a/pennylane/vqe/vqe.py
+++ b/pennylane/vqe/vqe.py
@@ -177,10 +177,24 @@ class ExpvalCost:
 
             def cost_fn(*qnode_args, **qnode_kwargs):
                 """Combine results from grouped QNode executions with grouped coefficients"""
-                total = 0
-                for o, c in zip(obs_groupings, coeffs_groupings):
-                    res = circuit(*qnode_args, obs=o, **qnode_kwargs)
-                    total += sum([r * c_ for r, c_ in zip(res, c)])
+                if device.shot_vector:
+                    shots_batch = 0
+
+                    for i in device.shot_vector:
+                        shots_batch += i[1]
+
+                    total = [0] * shots_batch
+
+                    for o, c in zip(obs_groupings, coeffs_groupings):
+                        res = circuit(*qnode_args, obs=o, **qnode_kwargs)
+
+                        for i, batch_res in enumerate(res):
+                            total[i] += sum(batch_res * c)
+                else:
+                    total = 0
+                    for o, c in zip(obs_groupings, coeffs_groupings):
+                        res = circuit(*qnode_args, obs=o, **qnode_kwargs)
+                        total += sum([r * c_ for r, c_ in zip(res, c)])
                 return total
 
             self.cost_fn = cost_fn


Test patch:
diff --git a/tests/test_vqe.py b/tests/test_vqe.py
index 88aea2cea..6da5c20cb 100644
--- a/tests/test_vqe.py
+++ b/tests/test_vqe.py
@@ -370,7 +370,8 @@ class TestVQE:
 
     @pytest.mark.slow
     @pytest.mark.parametrize("interface", ["tf", "torch", "autograd"])
-    def test_optimize(self, interface, tf_support, torch_support):
+    @pytest.mark.parametrize("shots", [None, [(8000, 5)], [(8000, 5), (9000, 4)]])
+    def test_optimize(self, interface, tf_support, torch_support, shots):
         """Test that an ExpvalCost with observable optimization gives the same result as another
         ExpvalCost without observable optimization."""
         if interface == "tf" and not tf_support:
@@ -378,7 +379,7 @@ class TestVQE:
         if interface == "torch" and not torch_support:
             pytest.skip("This test requires Torch")
 
-        dev = qml.device("default.qubit", wires=4)
+        dev = qml.device("default.qubit", wires=4, shots=shots)
         hamiltonian = big_hamiltonian
 
         cost = qml.ExpvalCost(
@@ -412,7 +413,7 @@ class TestVQE:
         assert exec_opt == 5  # Number of groups in the Hamiltonian
         assert exec_no_opt == 15
 
-        assert np.allclose(c1, c2)
+        assert np.allclose(c1, c2, atol=1e-1)
 
     @pytest.mark.parametrize("interface", ["tf", "torch", "autograd"])
     def test_optimize_multiple_terms(self, interface, tf_support, torch_support):