Proper compatible batching

With default.qubit and torch interface, I was able to batch inputs like this since version ~0.3X

qml.AngleEmbedding(inputs[:, x_idx: x_idx + self.n_qubits], wires=xwires, rotation='X')
qml.AngleEmbedding(inputs[:, x_idx + self.n_qubits: x_idx + 2 * self.n_qubits], wires=xwires, rotation='Y')
qml.AngleEmbedding(inputs[:, x_idx + 2 * self.n_qubits: x_idx + 3 * self.n_qubits], wires=xwires, rotation='Z')

This approach does not work with qiskit.aer and lightning.gpu. What is the current recommended practice for batching ? torch.vmap? Is there one that is compatible with both?

Hi @coot , welcome back!

PennyLane now automatically handles batching (since a while back in fact). We refer to this in the docs as parameter broadcasting. You can find some info on it here. Basically you code up your circuit as though you were only passing one input and then when you call that circuit you add a new dimension in the data (the batch dimension). So for qml.AngleEmbedding your inputs (when you call the circuit) can now be a 2D array instead of a 1D array. It should now be easier!
Note that a couple of things don’t support broadcasting so if you run into issues please check the docs for that function, follow the suggestions in the error message if you can, or come to the forum to see other similar questions and answers.

Let us know if this helps!