[BUG] Qiskit devices don't support vmap parameter broadcasting about pennylane HOT 3 CLOSED

@josh146 Yes it does.

Minimal non-working example:

dev = qml.device('lightning.qubit', wires=10)

@qml.qnode(dev)
def circuit(x):
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))

jax.vmap(circuit)(jax.numpy.array([0.5, 0.6, 0.7]))

from pennylane.

Thanks for opening this issue @lauracappelli . We'll try and get a fix in for the next release coming out on March 5th.

The problematic line of code is this:

This works for default.qubit because it does actually support native broadcasting. Most device don't, and use qml.transforms.broadcast_expand during the preprocessing step (Device.preprocess or Device.batch_transform). The problem is that jax is trying to add a batch dimension after we already handled any native broadcasting.

In the short-term, we just update the above line to be:

    out = jax.pure_callback(pure_callback_wrapper, shape_dtype_structs, params, vectorized=device.name == "default.qubit")

Though in the longer term, we should rethink jax.vmap and where we handle parameter broadcasting in pennylane.

from pennylane.

@albi3ro would this bug also impact lightning.qubit?

from pennylane.