Repository: PennyLaneAI/pennylane
Pull request #2091 (2022-01-16 18:04:11)

Add batching support to BasisStatePreparation template
Also add an additional test in tests/transforms/test_batch_params.py
------------------------------------------------------------------------------------------------------------

**Context:**
Part of the effort to support batching in circuits. See [#1745](https://github.com/PennyLaneAI/pennylane/issues/1745)

**Description of the Change:** 
 Allow a batch dimension in the BasisStatePreparation template.
Now the user would be able to run code like the following:
```python
import pennylane as qml
from pennylane import numpy as np

wires = 3
dev = qml.device("default.qubit", wires=wires)

@qml.batch_params
@qml.qnode(dev)
def circuit(data, weights):
    qml.templates.BasisStatePreparation(data, wires=[0, 1, 2])
    qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
    return qml.expval(qml.PauliX(0))

batch_size = 2

# create a batched input statevector
data = np.ones( (batch_size, wires))
circuit(data, weights)
```

**Benefits:**
Add this support to the BasisStatePreparation template.

**Possible Drawbacks:**
None

**Related GitHub Issues:**
https://github.com/PennyLaneAI/pennylane/issues/2074


codecov[bot]: # [Codecov](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#2091](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (0287b29) into [master](https://codecov.io/gh/PennyLaneAI/pennylane/commit/1275736f790ced1d778858ed383448d4a43a4cdd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (1275736) will **decrease** coverage by `2.14%`.
> The diff coverage is `91.66%`.

[![Impacted file tree graph](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/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/2091?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI)

```diff
@@            Coverage Diff             @@
##           master    #2091      +/-   ##
==========================================
- Coverage   99.18%   97.04%   -2.15%     
==========================================
  Files         226      226              
  Lines       17373    17376       +3     
==========================================
- Hits        17232    16863     -369     
- Misses        141      513     +372     
```


| [Impacted Files](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091?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://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RlbXBsYXRlcy9zdGF0ZV9wcmVwYXJhdGlvbnMvYmFzaXMucHk=) | `96.55% <91.66%> (-3.45%)` | :arrow_down: |
| [pennylane/interfaces/batch/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/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/interfaces/jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2ludGVyZmFjZXMvamF4LnB5) | `0.00% <0.00%> (-97.44%)` | :arrow_down: |
| [pennylane/devices/default\_qubit\_jax.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdF9qYXgucHk=) | `0.00% <0.00%> (-96.63%)` | :arrow_down: |
| [pennylane/qnn/cost.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3Fubi9jb3N0LnB5) | `42.85% <0.00%> (-57.15%)` | :arrow_down: |
| [pennylane/transforms/classical\_jacobian.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL3RyYW5zZm9ybXMvY2xhc3NpY2FsX2phY29iaWFuLnB5) | `77.08% <0.00%> (-22.92%)` | :arrow_down: |
| [pennylane/math/utils.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvdXRpbHMucHk=) | `81.81% <0.00%> (-18.19%)` | :arrow_down: |
| [pennylane/math/multi\_dispatch.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvbXVsdGlfZGlzcGF0Y2gucHk=) | `87.14% <0.00%> (-12.86%)` | :arrow_down: |
| [pennylane/math/is\_independent.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21hdGgvaXNfaW5kZXBlbmRlbnQucHk=) | `88.50% <0.00%> (-11.50%)` | :arrow_down: |
| [pennylane/collections/sum.py](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2NvbGxlY3Rpb25zL3N1bS5weQ==) | `89.47% <0.00%> (-10.53%)` | :arrow_down: |
| ... and [13 more](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091/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/2091?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/2091?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI). Last update [1275736...0287b29](https://codecov.io/gh/PennyLaneAI/pennylane/pull/2091?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).

antalszava: Hi @eddddddy, thank you for posting your solution! :slightly_smiling_face: We'll be having a look soon and getting back to you with comments.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 830471f99..61eafdef0 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -326,6 +326,9 @@
   [(#2062)](https://github.com/PennyLaneAI/pennylane/pull/2062)
   [(#2063)](https://github.com/PennyLaneAI/pennylane/pull/2063)
 
+* `qml.BasisStatePreparation` now supports the `batch_params` decorator.
+  [(#2091)](https://github.com/PennyLaneAI/pennylane/pull/2091)
+
 * Added a new `multi_dispatch` decorator that helps ease the definition of new functions
   inside PennyLane. The decorator is used throughout the math module, demonstrating use cases.
   [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084)
@@ -364,7 +367,7 @@
 
   Previously, `qml.jacobian` would attempt to stack the Jacobian for multiple
   QNode arguments, which succeeded whenever the arguments have the same shape.
-  In this case, the stacked Jacobian would also be transposed, leading to the 
+  In this case, the stacked Jacobian would also be transposed, leading to the
   output shape `(*reverse_QNode_args_shape, *reverse_output_shape, num_QNode_args)`
 
   If no stacking and transposing occurs, the output shape instead is a `tuple`
@@ -458,6 +461,5 @@
 
 This release contains contributions from (in alphabetical order):
 
-Juan Miguel Arrazola, Ali Asadi, Esther Cruz, Christina Lee, Olivia Di Matteo, Diego Guala, Josh Izaac,
-Ankit Khandelwal, Korbinian Kottmann, Jay Soni, Antal Száva, David Wierichs, Shaoming Zhang
-
+Juan Miguel Arrazola, Ali Asadi, Esther Cruz, Christina Lee, Olivia Di Matteo, Diego Guala, Edward Jiang, Josh Izaac,
+Ankit Khandelwal, Korbinian Kottmann, Jay Soni, Antal Száva, David Wierichs, Shaoming Zhang
\ No newline at end of file
diff --git a/pennylane/templates/state_preparations/basis.py b/pennylane/templates/state_preparations/basis.py
index 92a0c4a4f..20f806342 100644
--- a/pennylane/templates/state_preparations/basis.py
+++ b/pennylane/templates/state_preparations/basis.py
@@ -52,20 +52,29 @@ class BasisStatePreparation(Operation):
 
     def __init__(self, basis_state, wires, do_queue=True, id=None):
 
-        shape = qml.math.shape(basis_state)
+        # check if the `basis_state` param is batched
+        batched = len(qml.math.shape(basis_state)) > 1
 
-        if len(shape) != 1:
-            raise ValueError(f"Basis state must be one-dimensional; got shape {shape}.")
+        state_batch = basis_state if batched else [basis_state]
 
-        n_bits = shape[0]
-        if n_bits != len(wires):
-            raise ValueError(f"Basis state must be of length {len(wires)}; got length {n_bits}.")
+        for i, state in enumerate(state_batch):
+            shape = qml.math.shape(state)
 
-        # we can extract a list here, because embedding is not differentiable
-        basis_state = list(qml.math.toarray(basis_state))
+            if len(shape) != 1:
+                raise ValueError(
+                    f"Basis states must be one-dimensional; state {i} has shape {shape}."
+                )
 
-        if not all(bit in [0, 1] for bit in basis_state):
-            raise ValueError(f"Basis state must only consist of 0s and 1s; got {basis_state}")
+            n_bits = shape[0]
+            if n_bits != len(wires):
+                raise ValueError(
+                    f"Basis states must be of length {len(wires)}; state {i} has length {n_bits}."
+                )
+
+            if not all(bit in [0, 1] for bit in state):
+                raise ValueError(
+                    f"Basis states must only consist of 0s and 1s; state {i} is {basis_state}"
+                )
 
         super().__init__(basis_state, wires=wires, do_queue=do_queue, id=id)
 


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 b542b9223..2ac99dd64 100644
--- a/tests/templates/test_state_preparations/test_basis_state_prep.py
+++ b/tests/templates/test_state_preparations/test_basis_state_prep.py
@@ -112,7 +112,7 @@ class TestInputs:
         """Tests that the correct error message is raised when the number
         of qubits does not match the number of wires."""
 
-        with pytest.raises(ValueError, match="Basis state must be of (shape|length)"):
+        with pytest.raises(ValueError, match="Basis states must be of (shape|length)"):
             qml.BasisStatePreparation(basis_state, wires)
 
     # fmt: off
@@ -125,7 +125,7 @@ class TestInputs:
         """Tests that the correct error messages is raised when
         the basis state contains numbers different from 0 and 1."""
 
-        with pytest.raises(ValueError, match="Basis state must only (contain|consist)"):
+        with pytest.raises(ValueError, match="Basis states must only (contain|consist)"):
             qml.BasisStatePreparation(basis_state, wires)
 
     def test_exception_wrong_dim(self):
@@ -138,15 +138,15 @@ class TestInputs:
             qml.BasisStatePreparation(basis_state, wires=range(2))
             return qml.expval(qml.PauliZ(0))
 
-        with pytest.raises(ValueError, match="Basis state must be one-dimensional"):
-            basis_state = np.array([[0, 1]])
+        with pytest.raises(ValueError, match="Basis states must be one-dimensional"):
+            basis_state = np.array([[[0, 1]]])
             circuit(basis_state)
 
-        with pytest.raises(ValueError, match="Basis state must be of length"):
+        with pytest.raises(ValueError, match="Basis states must be of length"):
             basis_state = np.array([0, 1, 0])
             circuit(basis_state)
 
-        with pytest.raises(ValueError, match="Basis state must only consist of"):
+        with pytest.raises(ValueError, match="Basis states must only consist of"):
             basis_state = np.array([0, 2])
             circuit(basis_state)
 
diff --git a/tests/transforms/test_batch_params.py b/tests/transforms/test_batch_params.py
index daa722587..a44846177 100644
--- a/tests/transforms/test_batch_params.py
+++ b/tests/transforms/test_batch_params.py
@@ -121,6 +121,41 @@ 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=4)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data, weights):
+        qml.templates.BasisStatePreparation(data, wires=[0, 1, 2, 3])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2, 3])
+        return qml.probs(wires=[0, 1, 2, 3])
+
+    batch_size = 5
+
+    # create random batched basis states
+    data = np.random.randint(2, size=(batch_size, 4))
+    weights = np.random.random((batch_size, 10, 4, 3))
+
+    spy = mocker.spy(circuit.device, "batch_execute")
+    res = circuit(data, weights)
+    assert res.shape == (batch_size, 2 ** 4)
+    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, 3])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2, 3])
+        return qml.probs(wires=[0, 1, 2, 3])
+
+    indiv_res = []
+    for state, weight in zip(data, weights):
+        indiv_res.append(circuit2(state, weight))
+    assert np.allclose(res, indiv_res)
+
+
 @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"])
 def test_autograd(diff_method, tol):
     """Test derivatives when using autograd"""