Repository: PennyLaneAI/pennylane

Issue:
Add batching support to the `MottonenStatePreparation` template
### Feature details

Following up on the issue description from https://github.com/PennyLaneAI/pennylane/issues/1745 we would like to add batch support to the [`MottonenStatePreparation` template](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.templates.state_preparations.MottonenStatePreparation.html).

This follows the addition in https://github.com/PennyLaneAI/pennylane/pull/1710.

### Implementation

We would like to support the following example:

```python
import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit", wires=3)

@qml.batch_params
@qml.beta.qnode(dev)
def circuit(data, weights):
    qml.templates.MottonenStatePreparation(data, wires=[0, 1, 2])
    qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
    return qml.expval(qml.PauliX(0))

batch_size = 3

# create a batched input statevector
data = np.random.random((batch_size, 8))
data /= np.linalg.norm(data, axis=1).reshape(-1, 1)  # normalize

weights = np.random.random((batch_size, 10, 3, 3))
circuit(data, weights)
```

This can be done by modifying the `MottonenStatePreparation` source code as needed. In addition, a test should be added to `tests/transforms/test_batch_params.py`.

### How important would you say this feature is?

2: Somewhat important. Needed this quarter.

### Additional information

_No response_


Discussion:

Resolving patch:
diff --git a/pennylane/templates/state_preparations/mottonen.py b/pennylane/templates/state_preparations/mottonen.py
--- a/pennylane/templates/state_preparations/mottonen.py
+++ b/pennylane/templates/state_preparations/mottonen.py
@@ -80,15 +80,15 @@ def _compute_theta(alpha):
     Returns:
         (tensor_like): rotation angles theta
     """
-    ln = alpha.shape[0]
-    k = np.log2(alpha.shape[0])
+    ln = alpha.shape[-1]
+    k = np.log2(ln)
 
     M_trans = np.zeros(shape=(ln, ln))
     for i in range(len(M_trans)):
         for j in range(len(M_trans[0])):
             M_trans[i, j] = _matrix_M_entry(j, i)
 
-    theta = qml.math.dot(M_trans, alpha)
+    theta = qml.math.dot(M_trans, alpha.T).T
 
     return theta / 2 ** k
 
@@ -120,8 +120,8 @@ def _uniform_rotation_dagger(gate, alpha, control_wires, target_wire):
     gray_code_rank = len(control_wires)
 
     if gray_code_rank == 0:
-        if theta[0] != 0.0:
-            gate(theta[0], wires=[target_wire])
+        if qml.math.all(theta[..., 0] != 0.0):
+            gate(theta[..., 0], wires=[target_wire])
         return
 
     code = gray_code(gray_code_rank)
@@ -133,8 +133,8 @@ def _uniform_rotation_dagger(gate, alpha, control_wires, target_wire):
     ]
 
     for i, control_index in enumerate(control_indices):
-        if theta[i] != 0.0:
-            gate(theta[i], wires=[target_wire])
+        if qml.math.all(theta[..., i] != 0.0):
+            gate(theta[..., i], wires=[target_wire])
         qml.CNOT(wires=[control_wires[control_index], target_wire])
 
 
@@ -163,11 +163,11 @@ def _get_alpha_z(omega, n, k):
         for j in range(1, 2 ** (n - k) + 1)
     ]
 
-    term1 = qml.math.take(omega, indices=indices1)
-    term2 = qml.math.take(omega, indices=indices2)
+    term1 = qml.math.take(omega, indices=indices1, axis=-1)
+    term2 = qml.math.take(omega, indices=indices2, axis=-1)
     diff = (term1 - term2) / 2 ** (k - 1)
 
-    return qml.math.sum(diff, axis=1)
+    return qml.math.sum(diff, axis=-1)
 
 
 def _get_alpha_y(a, n, k):
@@ -190,12 +190,12 @@ def _get_alpha_y(a, n, k):
         [(2 * (j + 1) - 1) * 2 ** (k - 1) + l for l in range(2 ** (k - 1))]
         for j in range(2 ** (n - k))
     ]
-    numerator = qml.math.take(a, indices=indices_numerator)
-    numerator = qml.math.sum(qml.math.abs(numerator) ** 2, axis=1)
+    numerator = qml.math.take(a, indices=indices_numerator, axis=-1)
+    numerator = qml.math.sum(qml.math.abs(numerator) ** 2, axis=-1)
 
     indices_denominator = [[j * 2 ** k + l for l in range(2 ** k)] for j in range(2 ** (n - k))]
-    denominator = qml.math.take(a, indices=indices_denominator)
-    denominator = qml.math.sum(qml.math.abs(denominator) ** 2, axis=1)
+    denominator = qml.math.take(a, indices=indices_denominator, axis=-1)
+    denominator = qml.math.sum(qml.math.abs(denominator) ** 2, axis=-1)
 
     # Divide only where denominator is zero, else leave initial value of zero.
     # The equation guarantees that the numerator is also zero in the corresponding entries.
@@ -250,20 +250,31 @@ class MottonenStatePreparation(Operation):
 
     def __init__(self, state_vector, wires, do_queue=True, id=None):
 
-        shape = qml.math.shape(state_vector)
+        # check if the `state_vector` param is batched
+        batched = len(qml.math.shape(state_vector)) > 1
 
-        if len(shape) != 1:
-            raise ValueError(f"State vector must be a one-dimensional vector; got shape {shape}.")
+        state_batch = state_vector if batched else [state_vector]
 
-        n_amplitudes = shape[0]
-        if n_amplitudes != 2 ** len(qml.wires.Wires(wires)):
-            raise ValueError(
-                f"State vector must be of length {2 ** len(wires)} or less; got length {n_amplitudes}."
-            )
+        # apply checks to each state vector in the batch
+        for i, state in enumerate(state_batch):
+            shape = qml.math.shape(state)
 
-        norm = qml.math.sum(qml.math.abs(state_vector) ** 2)
-        if not qml.math.allclose(norm, 1.0, atol=1e-3):
-            raise ValueError("State vector has to be of norm 1.0, got {}".format(norm))
+            if len(shape) != 1:
+                raise ValueError(
+                    f"State vectors must be one-dimensional; vector {i} has shape {shape}."
+                )
+
+            n_amplitudes = shape[0]
+            if n_amplitudes != 2 ** len(qml.wires.Wires(wires)):
+                raise ValueError(
+                    f"State vectors must be of length {2 ** len(wires)} or less; vector {i} has length {n_amplitudes}."
+                )
+
+            norm = qml.math.sum(qml.math.abs(state) ** 2)
+            if not qml.math.allclose(norm, 1.0, atol=1e-3):
+                raise ValueError(
+                    f"State vectors have to be of norm 1.0, vector {i} has norm {norm}"
+                )
 
         super().__init__(state_vector, wires=wires, do_queue=do_queue, id=id)
 


Test patch:
diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py
--- a/tests/math/test_functions.py
+++ b/tests/math/test_functions.py
@@ -1215,6 +1215,19 @@ def test_array_indexing(self, t):
         res = fn.take(t, indices)
         assert fn.allclose(res, [1, 3, 4, 5, 2])
 
+    def test_array_indexing_autograd(self):
+        """Test that indexing with a sequence properly extracts
+        the elements from the flattened tensor"""
+        t = np.array([[[1, 2], [3, 4], [-1, 1]], [[5, 6], [0, -1], [2, 1]]])
+        indices = [0, 2, 3, 6, -2]
+
+        def cost_fn(t):
+            return np.sum(fn.take(t, indices))
+
+        grad = qml.grad(cost_fn)(t)
+        expected = np.array([[[1, 0], [1, 1], [0, 0]], [[1, 0], [0, 0], [1, 0]]])
+        assert fn.allclose(grad, expected)
+
     @pytest.mark.parametrize("t", take_data)
     def test_multidimensional_indexing(self, t):
         """Test that indexing with a multi-dimensional sequence properly extracts
diff --git a/tests/templates/test_state_preparations/test_mottonen_state_prep.py b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
--- a/tests/templates/test_state_preparations/test_mottonen_state_prep.py
+++ b/tests/templates/test_state_preparations/test_mottonen_state_prep.py
@@ -294,7 +294,7 @@ def test_error_state_vector_not_normalized(self, state_vector, wires):
         """Tests that the correct error messages is raised if
         the given state vector is not normalized."""
 
-        with pytest.raises(ValueError, match="State vector has to be of norm"):
+        with pytest.raises(ValueError, match="State vectors have to be of norm"):
             qml.MottonenStatePreparation(state_vector, wires)
 
     # fmt: off
@@ -308,21 +308,21 @@ def test_error_num_entries(self, state_vector, wires):
         the number of entries in the given state vector does not match
         with the number of wires in the system."""
 
-        with pytest.raises(ValueError, match="State vector must be of (length|shape)"):
+        with pytest.raises(ValueError, match="State vectors must be of (length|shape)"):
             qml.MottonenStatePreparation(state_vector, wires)
 
     @pytest.mark.parametrize(
         "state_vector",
         [
-            ([[0, 0, 1, 0]]),
-            ([[0, 1], [1, 0], [0, 0], [0, 0]]),
+            ([[[0, 0, 1, 0]]]),
+            ([[[0, 1], [1, 0], [0, 0], [0, 0]]]),
         ],
     )
     def test_exception_wrong_shape(self, state_vector):
         """Verifies that exception is raised if the
         number of dimensions of features is incorrect."""
 
-        with pytest.raises(ValueError, match="State vector must be a one-dimensional"):
+        with pytest.raises(ValueError, match="State vectors must be one-dimensional"):
             qml.MottonenStatePreparation(state_vector, 2)
 
     def test_id(self):
diff --git a/tests/transforms/test_batch_params.py b/tests/transforms/test_batch_params.py
--- a/tests/transforms/test_batch_params.py
+++ b/tests/transforms/test_batch_params.py
@@ -65,6 +65,42 @@ def circuit(data):
     assert len(spy.call_args[0][0]) == batch_size
 
 
+def test_mottonenstate_preparation(mocker):
+    """Test that batching works for MottonenStatePreparation"""
+    dev = qml.device("default.qubit", wires=3)
+
+    @qml.batch_params
+    @qml.qnode(dev)
+    def circuit(data, weights):
+        qml.templates.MottonenStatePreparation(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.random.random((batch_size, 2 ** 3))
+    data /= np.linalg.norm(data, axis=1).reshape(-1, 1)  # normalize
+    weights = np.random.random((batch_size, 10, 3, 3))
+
+    spy = mocker.spy(circuit.device, "batch_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.MottonenStatePreparation(data, wires=[0, 1, 2])
+        qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
+        return qml.probs(wires=[0, 1, 2])
+
+    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"""