Repository: PennyLaneAI/pennylane
Pull request #5195 (2024-02-13 14:14:03)

Add broadcasting support to mitigate_with_zne
**Context:**
A QNode transformed with `mitigate_with_zne` does not accept parameters with a leading dimension

**Description of the Change:**
1. Replace `qml.math.dot` with `qml.math.tensordot` in `_polyfit`
2. Test case added for parameter broadcasting with `mitigate_with_zne`

**Benefits:**
BugFix

**Related GitHub Issues:**
https://github.com/PennyLaneAI/pennylane/issues/5100

**Related Shortcut Story:**
[sc-55298]

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://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5195?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Comparison is base [(`eb6298d`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/eb6298dbea57115f26c2dfe0daaac8c59720fa6a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68% compared to head [(`4f8ff59`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5195?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.68%.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5195      +/-   ##
==========================================
- Coverage   99.68%   99.68%   -0.01%     
==========================================
  Files         396      396              
  Lines       36468    36196     -272     
==========================================
- Hits        36354    36081     -273     
- Misses        114      115       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5195?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?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 b72be8aa7..765271edf 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -494,6 +494,9 @@
   operators have a valid `pauli_rep` property.
   [(#5177)](https://github.com/PennyLaneAI/pennylane/pull/5177)
 
+* A `QNode` transformed with `mitigate_with_zne` now accepts batch parameters.
+  [(#5195)](https://github.com/PennyLaneAI/pennylane/pull/5195)
+
 * The matrix of an empty `PauliSentence` instance is now correct (all-zeros).
   Further, matrices of empty `PauliWord` and `PauliSentence` instances can be turned to matrices now.
   [(#5188)](https://github.com/PennyLaneAI/pennylane/pull/5188)
diff --git a/pennylane/transforms/mitigate.py b/pennylane/transforms/mitigate.py
index d88f02e4f..9b2c1d847 100644
--- a/pennylane/transforms/mitigate.py
+++ b/pennylane/transforms/mitigate.py
@@ -261,7 +261,7 @@ def _polyfit(x, y, order):
     # i.e. coeffs = (X.T @ X)**-1 X.T @ y see https://en.wikipedia.org/wiki/Polynomial_regression
     c = qml.math.linalg.pinv(qml.math.transpose(X) @ X)
     c = c @ qml.math.transpose(X)
-    c = qml.math.dot(c, y)
+    c = qml.math.tensordot(c, y, axes=1)
     c = qml.math.transpose(qml.math.transpose(c) / scale)
     return c
 


Test patch:
diff --git a/tests/transforms/test_mitigate.py b/tests/transforms/test_mitigate.py
index 19fe155e1..5b350c31f 100644
--- a/tests/transforms/test_mitigate.py
+++ b/tests/transforms/test_mitigate.py
@@ -179,6 +179,30 @@ class TestMitigateWithZNE:
         assert args[0][0] == scale_factors
         assert np.allclose(args[0][1], np.mean(np.reshape(random_results, (3, 2)), axis=1))
 
+    def test_braodcasting(self):
+        """Tests that mitigate_with_zne supports batch arguments"""
+
+        batch_size = 2
+
+        @qml.qnode(dev_noisy)
+        def original_qnode(inputs):
+            qml.AmplitudeEmbedding(features=inputs, wires=range(2), normalize=True)
+            return [qml.expval(qml.PauliZ(wires=i)) for i in range(2)]
+
+        expanded_qnode = qml.transforms.broadcast_expand(original_qnode)
+
+        mitigated_qnode_orig = qml.transforms.mitigate_with_zne(
+            original_qnode, [1, 2, 3], fold_global, richardson_extrapolate
+        )
+        mitigated_qnode_expanded = qml.transforms.mitigate_with_zne(
+            expanded_qnode, [1, 2, 3], fold_global, richardson_extrapolate
+        )
+        rng = np.random.default_rng(seed=18954959)
+        inputs = rng.uniform(0, 1, size=(batch_size, 2**2))
+        result_orig = mitigated_qnode_orig(inputs)
+        result_expanded = mitigated_qnode_expanded(inputs)
+        assert qml.math.allclose(result_orig, result_expanded)
+
 
 @pytest.fixture
 def skip_if_no_mitiq_support():