Repository: PennyLaneAI/pennylane
Pull request #4767 (2023-10-31 21:33:14)

Raise error if broadcasting with `MottonenStatePreparation` and `BasisStatePreparation` `compute_decomposition`
**Context:**
`MottonenStatePreparation` and `BasisStatePreparation` fails right now when decomposing a broadcasted state vector because different state vectors have different decompositions.

**Description of the Change:**
* Raise an error in `MottonenStatePreparation.compute_decomposition` and `BasisStatePreparation.compute_decomposition` if the `batch_size` if not `None` and suggest users to use `broadcast_expand`.

**Benefits:**
No failures or wrong results when decomposing Mottonen or BasisState

**Possible Drawbacks:**

**Related GitHub Issues:**
#4753 

mudit2812: [sc-48950]
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4767?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 [(`8aee680`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/8aee6805684cffe3fc0bb7b5c641f94ad7301859?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64% compared to head [(`9e4b9d0`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4767?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) 99.64%.

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


```diff
@@            Coverage Diff             @@
##           master    #4767      +/-   ##
==========================================
- Coverage   99.64%   99.64%   -0.01%     
==========================================
  Files         381      381              
  Lines       34308    34052     -256     
==========================================
- Hits        34187    33930     -257     
- Misses        121      122       +1     
```


| [Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4767?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/templates/state\_preparations/basis.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4767?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdGF0ZV9wcmVwYXJhdGlvbnMvYmFzaXMucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/templates/state\_preparations/mottonen.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4767?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdGF0ZV9wcmVwYXJhdGlvbnMvbW90dG9uZW4ucHk=) | `100.00% <100.00%> (ø)` | |

... and [42 files with indirect coverage changes](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4767/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4767?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 6a70dd5b1..7c5fd3ece 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -69,6 +69,12 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `MottonenStatePreparation` now raises an error if decomposing a broadcasted state vector.
+  [(#4767)](https://github.com/PennyLaneAI/pennylane/pull/4767)
+
+* `BasisStatePreparation` now raises an error if decomposing a broadcasted state vector.
+  [(#4767)](https://github.com/PennyLaneAI/pennylane/pull/4767)
+
 * Gradient transforms now work with overridden shot vectors and default qubit.
   [(#4795)](https://github.com/PennyLaneAI/pennylane/pull/4795)
 
@@ -99,6 +105,7 @@ Lillian Frederiksen,
 Ankit Khandelwal,
 Christina Lee,
 Anurav Modak,
+Mudit Pandey,
 Matthew Silverman,
 David Wierichs,
 Justin Woodring,
diff --git a/pennylane/templates/state_preparations/basis.py b/pennylane/templates/state_preparations/basis.py
index 550b5c10a..0b309e383 100644
--- a/pennylane/templates/state_preparations/basis.py
+++ b/pennylane/templates/state_preparations/basis.py
@@ -54,6 +54,7 @@ class BasisStatePreparation(Operation):
     num_params = 1
     num_wires = AnyWires
     grad_method = None
+    ndim_params = (1,)
 
     def __init__(self, basis_state, wires, id=None):
         basis_state = qml.math.stack(basis_state)
@@ -109,6 +110,13 @@ class BasisStatePreparation(Operation):
         [PauliX(wires=['a']),
         PauliX(wires=['b'])]
         """
+        if len(qml.math.shape(basis_state)) > 1:
+            raise ValueError(
+                "Broadcasting with BasisStatePreparation is not supported. Please use the "
+                "qml.transforms.broadcast_expand transform to use broadcasting with "
+                "BasisStatePreparation."
+            )
+
         if not qml.math.is_abstract(basis_state):
             op_list = []
             for wire, state in zip(wires, basis_state):
diff --git a/pennylane/templates/state_preparations/mottonen.py b/pennylane/templates/state_preparations/mottonen.py
index 1818bd7ba..e8440ff83 100644
--- a/pennylane/templates/state_preparations/mottonen.py
+++ b/pennylane/templates/state_preparations/mottonen.py
@@ -286,6 +286,7 @@ class MottonenStatePreparation(Operation):
 
     num_wires = AnyWires
     grad_method = None
+    ndim_params = (1,)
 
     def __init__(self, state_vector, wires, id=None):
         # check if the `state_vector` param is batched
@@ -347,6 +348,12 @@ class MottonenStatePreparation(Operation):
         CNOT(wires=['a', 'b']),
         CNOT(wires=['a', 'b'])]
         """
+        if len(qml.math.shape(state_vector)) > 1:
+            raise ValueError(
+                "Broadcasting with MottonenStatePreparation is not supported. Please use the "
+                "qml.transforms.broadcast_expand transform to use broadcasting with "
+                "MottonenStatePreparation."
+            )
 
         a = qml.math.abs(state_vector)
         omega = qml.math.angle(state_vector)


Test patch:
diff --git a/tests/templates/test_state_preparations/test_basis_state_prep.py b/tests/templates/test_state_preparations/test_basis_state_prep.py
index b63b01a33..24b58e906 100644
--- a/tests/templates/test_state_preparations/test_basis_state_prep.py
+++ b/tests/templates/test_state_preparations/test_basis_state_prep.py
@@ -153,6 +153,18 @@ class TestDecomposition:
         assert np.allclose(res1, res2, atol=tol, rtol=0)
         assert np.allclose(state1, state2, atol=tol, rtol=0)
 
+    def test_batched_decomposition_fails(self):
+        """Test that attempting to decompose a BasisStatePreparation operation with
+        broadcasting raises an error."""
+        state = np.array([[1, 0], [1, 1]])
+
+        op = qml.BasisStatePreparation(state, wires=[0, 1])
+        with pytest.raises(ValueError, match="Broadcasting with BasisStatePreparation"):
+            _ = op.decomposition()
+
+        with pytest.raises(ValueError, match="Broadcasting with BasisStatePreparation"):
+            _ = qml.BasisStatePreparation.compute_decomposition(state, qml.wires.Wires([0, 1]))
+
 
 class TestInputs:
     """Test inputs and pre-processing."""
diff --git a/tests/templates/test_state_preparations/test_mottonen_state_prep.py b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
index f7ac7bb38..e27977c10 100644
--- a/tests/templates/test_state_preparations/test_mottonen_state_prep.py
+++ b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
@@ -289,6 +289,18 @@ class TestDecomposition:
         assert np.allclose(res1, res2, atol=tol, rtol=0)
         assert np.allclose(state1, state2, atol=tol, rtol=0)
 
+    def test_batched_decomposition_fails(self):
+        """Test that attempting to decompose a MottonenStatePreparation operation with
+        broadcasting raises an error."""
+        state = np.array([[1 / 2, 1 / 2, 1 / 2, 1 / 2], [0.0, 0.0, 0.0, 1.0]])
+
+        op = qml.MottonenStatePreparation(state, wires=[0, 1])
+        with pytest.raises(ValueError, match="Broadcasting with MottonenStatePreparation"):
+            _ = op.decomposition()
+
+        with pytest.raises(ValueError, match="Broadcasting with MottonenStatePreparation"):
+            _ = qml.MottonenStatePreparation.compute_decomposition(state, qml.wires.Wires([0, 1]))
+
 
 class TestInputs:
     """Test inputs and pre-processing."""
diff --git a/tests/transforms/test_batch_input.py b/tests/transforms/test_batch_input.py
index d30a12cc0..c642cf15a 100644
--- a/tests/transforms/test_batch_input.py
+++ b/tests/transforms/test_batch_input.py
@@ -201,6 +201,43 @@ def test_mottonenstate_preparation(mocker):
     assert np.allclose(res, indiv_res)
 
 
+def test_basis_state_preparation(mocker):
+    """Test that batching works for BasisStatePreparation"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @partial(qml.batch_input, argnum=0)
+    @qml.qnode(dev, interface="autograd")
+    def circuit(data, weights):
+        qml.templates.BasisStatePreparation(data, wires=[0, 1, 2])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 1, 2])
+
+    batch_size = 3
+
+    # create a batched input statevector
+    data = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]], requires_grad=False)
+
+    # weights is not batched
+    weights = np.random.random((10, 3, 3), requires_grad=True)
+
+    spy = mocker.spy(circuit.device, "execute")
+    res = circuit(data, weights)
+    assert res.shape == (batch_size, 2**3)
+    assert len(spy.call_args[0][0]) == batch_size
+
+    # check the results against individually executed circuits (no batching)
+    @qml.qnode(dev)
+    def circuit2(data, weights):
+        qml.templates.BasisStatePreparation(data, wires=[0, 1, 2])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 1, 2])
+
+    indiv_res = []
+    for state in data:
+        indiv_res.append(circuit2(state, weights))
+    assert np.allclose(res, indiv_res)
+
+
 def test_qubit_state_prep(mocker):
     """Test that batching works for StatePrep"""
 
diff --git a/tests/transforms/test_batch_params.py b/tests/transforms/test_batch_params.py
index 36accf7bb..14e213846 100644
--- a/tests/transforms/test_batch_params.py
+++ b/tests/transforms/test_batch_params.py
@@ -279,14 +279,14 @@ def test_shot_vector():
     @qml.batch_params
     @qml.qnode(dev)
     def circuit(data, x, weights):
-        qml.templates.AmplitudeEmbedding(data, wires=[0, 1, 2], normalize=True)
+        qml.templates.AngleEmbedding(data, wires=[0, 1, 2])
         qml.RX(x, wires=0)
         qml.RY(0.2, wires=1)
         qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
         return qml.probs(wires=[0, 2])
 
     batch_size = 6
-    data = np.random.random((batch_size, 8))
+    data = np.random.random((batch_size, 3))
     x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
     weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
 
@@ -306,14 +306,14 @@ def test_multi_returns_shot_vector():
     @qml.batch_params
     @qml.qnode(dev)
     def circuit(data, x, weights):
-        qml.templates.AmplitudeEmbedding(data, wires=[0, 1, 2], normalize=True)
+        qml.templates.AngleEmbedding(data, wires=[0, 1, 2])
         qml.RX(x, wires=0)
         qml.RY(0.2, wires=1)
         qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
         return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 2])
 
     batch_size = 6
-    data = np.random.random((batch_size, 8))
+    data = np.random.random((batch_size, 3))
     x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
     weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)