Repository: PennyLaneAI/pennylane

Issue:
Add batching support to the AmplitudeEmbedding template
### Feature details

A recent PR, #1710, added support for creating PennyLane QNodes with *batches* of input parameters via [`@qml.batch_params`](https://pennylane.readthedocs.io/en/latest/code/api/pennylane.batch_params.html). Under the hood, a separate circuit per batch dimension is created, and as a result this approach is both hardware and simulator compatible:

```python
dev = qml.device("default.qubit", wires=2)

@qml.batch_params
@qml.beta.qnode(dev)
def circuit(x, weights):
    qml.RX(x, wires=0)
    qml.RY(0.2, wires=1)
    qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1, 2])
    return qml.expval(qml.Hadamard(0))

def cost(x, weights):
    return np.sum(circuit(x, weights))

batch_size = 3
x = np.linspace(0.1, 0.5, batch_size)
weights = np.random.random((batch_size, 10, 3, 3))

>>> circuit(x, weights)
[-0.30773348  0.23135516  0.13086565]
>>> cost_fn(x, weights)
-0.8581269507766536
>>> qml.grad(cost_fn)(x, weights)[0]
[ 0.23235464  0.00928953 -0.30083487]
```

One such template that would be great to support in batch mode is [`qml.templates.AmplitudeEmbedding`](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.templates.embeddings.AmplitudeEmbedding.html). In order to add support, this template needs to be modified to recognize the first dimension of the input parameter as a _batch_ dimension.

### Implementation

We would like to support the following example:

```python
dev = qml.device("default.qubit", wires=3)

@qml.batch_params
@qml.beta.qnode(dev)
def circuit(data, weights):
    qml.templates.AmplitudeEmbedding(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 `AmplitudeEmbedding` source code as needed. In addition, a test should be added to `tests/transforms/test_batch_params.py`.

Note that some hardware devices do not support `AmplitudeEmbedding` directly, and instead decompose via the [`qml.templates.MottonenStatePreparation`](https://pennylane.readthedocs.io/en/stable/code/api/pennylane.templates.state_preparations.MottonenStatePreparation.html) template, so adding support to this template would also be an added bonus.

### How important would you say this feature is?

2: Somewhat important. Needed this quarter.

### Additional information

_No response_


Discussion:
Hi Josh, I'd be happy to take a go at this!
@dime10 great! Don't hesitate to open a work in progress PR here, and ask any questions as they arise (you can tag me directly via my GitHub handle) 🙂 
Hi @josh146, after reading some of the relevant code I had some questions:

The way I understand it, the AmplitudeEmbedding template provides some preprocessing on the inputs and is then expanded into a tape consisting solely of a state preparation routine (QubitStateVector). The argument to QubitStateVector would then be a batched argument (i.e. the 2 dimensional features parameter). However, as this tape is converted into multiple tapes with individual parameters by the batch transform, QubitStateVector (or other operations in the template expansion) need not be aware of the batching. Is that accurate?

1. Are we only looking to support batching on the features parameter (for this template), or others as well (e.g. different padding for each circuit in a batch)?

2. Do we want to assert that each parameter in a batch has the same shape (e.g. [[0,2], [1,2], [2,2]]) or allow different shapes with padding enabled (e.g. [[0,2], [1], [2,2]] --> [[0,2], [1,0], [2,2]])?
Hi @dime10,

> need not be aware of the batching. Is that accurate?

Yes, your understanding of the process is spot on 🙂 

> Are we only looking to support batching on the features parameter (for this template), or others as well

This is a good question — I’m inclined to say let’s only add support for the `features` parameter for now. I imagine this is the main use case (wanting to execute a particular quantum algorithm with a ‘batch’ of input states), with the padding remaining the same for all states in the batch.

> Do we want to assert that each parameter in a batch has the same shape 

For simplicity, I think it makes sense here to assume that all states in the batch have the same shape. Another advantage is that it avoids ragged arrays; the batch of parameters will always be a well-defined array or tensor.

Let me know if that helps, or if you have any other questions!

Resolving patch:
diff --git a/pennylane/templates/embeddings/amplitude.py b/pennylane/templates/embeddings/amplitude.py
--- a/pennylane/templates/embeddings/amplitude.py
+++ b/pennylane/templates/embeddings/amplitude.py
@@ -162,6 +162,7 @@ def expand(self):
     def _preprocess(features, wires, pad_with, normalize):
         """Validate and pre-process inputs as follows:
 
+        * If features is batched, the processing that follows is applied to each feature set in the batch.
         * Check that the features tensor is one-dimensional.
         * If pad_with is None, check that the first dimension of the features tensor
           has length :math:`2^n` where :math:`n` is the number of qubits. Else check that the
@@ -170,41 +171,49 @@ def _preprocess(features, wires, pad_with, normalize):
           features tensor.
         """
 
-        shape = qml.math.shape(features)
+        # check if features is batched
+        batched = len(qml.math.shape(features)) > 1
 
-        # check shape
-        if len(shape) != 1:
-            raise ValueError(f"Features must be a one-dimensional tensor; got shape {shape}.")
+        features_batch = features if batched else [features]
 
-        n_features = shape[0]
-        if pad_with is None and n_features != 2 ** len(wires):
-            raise ValueError(
-                f"Features must be of length {2 ** len(wires)}; got length {n_features}. "
-                f"Use the 'pad' argument for automated padding."
-            )
-
-        if pad_with is not None and n_features > 2 ** len(wires):
-            raise ValueError(
-                f"Features must be of length {2 ** len(wires)} or "
-                f"smaller to be padded; got length {n_features}."
-            )
+        # apply pre-processing to each features tensor in the batch
+        for i, feature_set in enumerate(features_batch):
+            shape = qml.math.shape(feature_set)
 
-        # pad
-        if pad_with is not None and n_features < 2 ** len(wires):
-            padding = [pad_with] * (2 ** len(wires) - n_features)
-            features = qml.math.concatenate([features, padding], axis=0)
+            # check shape
+            if len(shape) != 1:
+                raise ValueError(f"Features must be a one-dimensional tensor; got shape {shape}.")
 
-        # normalize
-        norm = qml.math.sum(qml.math.abs(features) ** 2)
+            n_features = shape[0]
+            if pad_with is None and n_features != 2 ** len(wires):
+                raise ValueError(
+                    f"Features must be of length {2 ** len(wires)}; got length {n_features}. "
+                    f"Use the 'pad' argument for automated padding."
+                )
 
-        if not qml.math.allclose(norm, 1.0, atol=TOLERANCE):
-            if normalize or pad_with:
-                features = features / np.sqrt(norm)
-            else:
+            if pad_with is not None and n_features > 2 ** len(wires):
                 raise ValueError(
-                    f"Features must be a vector of length 1.0; got length {norm}."
-                    "Use 'normalize=True' to automatically normalize."
+                    f"Features must be of length {2 ** len(wires)} or "
+                    f"smaller to be padded; got length {n_features}."
                 )
 
-        features = qml.math.cast(features, np.complex128)
-        return features
+            # pad
+            if pad_with is not None and n_features < 2 ** len(wires):
+                padding = [pad_with] * (2 ** len(wires) - n_features)
+                feature_set = qml.math.concatenate([feature_set, padding], axis=0)
+
+            # normalize
+            norm = qml.math.sum(qml.math.abs(feature_set) ** 2)
+
+            if not qml.math.allclose(norm, 1.0, atol=TOLERANCE):
+                if normalize or pad_with:
+                    feature_set = feature_set / np.sqrt(norm)
+                else:
+                    raise ValueError(
+                        f"Features must be a vector of norm 1.0; got norm {norm}."
+                        "Use 'normalize=True' to automatically normalize."
+                    )
+
+            features_batch[i] = qml.math.cast(feature_set, np.complex128)
+
+        return features_batch if batched else features_batch[0]


Test patch:
diff --git a/tests/templates/test_embeddings/test_amplitude.py b/tests/templates/test_embeddings/test_amplitude.py
--- a/tests/templates/test_embeddings/test_amplitude.py
+++ b/tests/templates/test_embeddings/test_amplitude.py
@@ -125,11 +125,11 @@ def circuit(x=None):
             )
             return [qml.expval(qml.PauliZ(i)) for i in range(n_qubits)]
 
-        with pytest.raises(ValueError, match="Features must be a vector of length"):
+        with pytest.raises(ValueError, match="Features must be a vector of norm"):
             circuit(x=not_nrmlzd)
 
     def test_throws_exception_if_features_wrong_shape(self):
-        """Checks exception if features has more than one dimension."""
+        """Checks exception if features has more than two dimensions."""
 
         n_qubits = 2
         dev = qml.device("default.qubit", wires=n_qubits)
@@ -140,7 +140,7 @@ def circuit(x=None):
             return qml.expval(qml.PauliZ(0))
 
         with pytest.raises(ValueError, match="Features must be a one-dimensional (tensor|vector)"):
-            circuit(x=[[1.0, 0.0], [0.0, 0.0]])
+            circuit(x=[[[1.0, 0.0], [0.0, 0.0]], [[1.0, 0.0], [0.0, 0.0]]])
 
     @pytest.mark.parametrize("inpt", NOT_ENOUGH_FEATURES + TOO_MANY_FEATURES)
     def test_throws_exception_if_fewer_features_than_amplitudes(self, inpt):
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
@@ -27,18 +27,20 @@ def test_simple_circuit(mocker):
 
     @qml.batch_params
     @qml.qnode(dev)
-    def circuit(x, weights):
+    def circuit(data, x, weights):
+        qml.templates.AmplitudeEmbedding(data, wires=[0, 1, 2], normalize=True)
         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 = 3
+    data = np.random.random((batch_size, 8))
     x = np.linspace(0.1, 0.5, batch_size, requires_grad=True)
     weights = np.ones((batch_size, 10, 3, 3), requires_grad=True)
 
     spy = mocker.spy(circuit.device, "batch_execute")
-    res = circuit(x, weights)
+    res = circuit(data, x, weights)
     assert res.shape == (batch_size, 1, 4)
     assert len(spy.call_args[0][0]) == batch_size