Add batching support to the `BasisStatePreparation` template

[antalszava]

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 BasisStatePreparation template.

This follows the addition in https://github.com/PennyLaneAI/pennylane/pull/1710.

Implementation

We would like to support the following example:

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)

This can be done by modifying the BasisStatePreparation 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

[eddddddy]

Hello, I submitted a pull request for this issue. Let me know if there's anything I missed.