Parameter broadcasting of initial states via `AmplitudeEmbedding`

[Qottmann]

Not sure if this is a bug or just not implemented yet. I want to use parameter broadcasting with different initial states via qml.AmplitudeEmbedding. Here is an example without parameter broadcasting:

dev = qml.device('default.qubit', wires=3)
state = np.ones(2**3)/np.sqrt(2**3)
@qml.qnode(dev)
def circuit_rx(x, z, state):
    qml.AmplitudeEmbedding(features=state, wires=range(3))
    qml.RX(x, wires=0)
    qml.RZ(z, wires=0)
    qml.RY(0.3, wires=0)
    return qml.probs(wires=0)
circuit_rx(0.5, 0.4, state)

I would expect the following to work with parameter broadcasting:

states = np.array([state for _ in range(10)])
circuit_rx(0.5, 0.4, states)

but I get

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
/home/qottmann/Qottmann/Xanadu/Sandbox.ipynb Cell 15 in <cell line: 2>()
      1 states = np.array([state for _ in range(10)])
----> 2 circuit_rx(0.5, 0.4, states)

File ~/Qottmann/Xanadu/pennylane/pennylane/qnode.py:619, in QNode.__call__(self, *args, **kwargs)
    612 using_custom_cache = (
    613     hasattr(cache, "__getitem__")
    614     and hasattr(cache, "__setitem__")
    615     and hasattr(cache, "__delitem__")
    616 )
    617 self._tape_cached = using_custom_cache and self.tape.hash in cache
--> 619 res = qml.execute(
    620     [self.tape],
    621     device=self.device,
    622     gradient_fn=self.gradient_fn,
    623     interface=self.interface,
    624     gradient_kwargs=self.gradient_kwargs,
    625     override_shots=override_shots,
    626     **self.execute_kwargs,
    627 )
    629 if autograd.isinstance(res, (tuple, list)) and len(res) == 1:
    630     # If a device batch transform was applied, we need to 'unpack'
    631     # the returned tuple/list to a float.
   (...)
    638     # TODO: find a more explicit way of determining that a batch transform
    639     # was applied.
    641     res = res[0]

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/execution.py:344, in execute(tapes, device, gradient_fn, interface, mode, gradient_kwargs, cache, cachesize, max_diff, override_shots, expand_fn, max_expansion, device_batch_transform)
    340     return batch_fn(res)
    342 if gradient_fn == "backprop" or interface is None:
    343     return batch_fn(
--> 344         qml.interfaces.cache_execute(
    345             batch_execute, cache, return_tuple=False, expand_fn=expand_fn
    346         )(tapes)
    347     )
    349 # the default execution function is batch_execute
    350 execute_fn = qml.interfaces.cache_execute(batch_execute, cache, expand_fn=expand_fn)

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/execution.py:172, in cache_execute.<locals>.wrapper(tapes, **kwargs)
    168         return (res, []) if return_tuple else res
    170 else:
    171     # execute all unique tapes that do not exist in the cache
--> 172     res = fn(execution_tapes.values(), **kwargs)
    174 final_res = []
    176 for i, tape in enumerate(tapes):

File ~/Qottmann/Xanadu/pennylane/pennylane/interfaces/execution.py:97, in cache_execute.<locals>.fn(tapes, **kwargs)
     95 def fn(tapes, **kwargs):  # pylint: disable=function-redefined
     96     tapes = [expand_fn(tape) for tape in tapes]
---> 97     return original_fn(tapes, **kwargs)

File ~/anaconda3/envs/pennylane/lib/python3.8/contextlib.py:75, in ContextDecorator.__call__.<locals>.inner(*args, **kwds)
     72 @wraps(func)
     73 def inner(*args, **kwds):
     74     with self._recreate_cm():
---> 75         return func(*args, **kwds)

File ~/Qottmann/Xanadu/pennylane/pennylane/_qubit_device.py:362, in QubitDevice.batch_execute(self, circuits)
    357 for circuit in circuits:
    358     # we need to reset the device here, else it will
    359     # not start the next computation in the zero state
    360     self.reset()
--> 362     res = self.execute(circuit)
    363     results.append(res)
    365 if self.tracker.active:

File ~/Qottmann/Xanadu/pennylane/pennylane/_qubit_device.py:260, in QubitDevice.execute(self, circuit, **kwargs)
    257 self.check_validity(circuit.operations, circuit.observables)
    259 # apply all circuit operations
--> 260 self.apply(circuit.operations, rotations=circuit.diagonalizing_gates, **kwargs)
    262 # generate computational basis samples
    263 if self.shots is not None or circuit.is_sampled:

File ~/Qottmann/Xanadu/pennylane/pennylane/devices/default_qubit.py:221, in DefaultQubit.apply(self, operations, rotations, **kwargs)
    215     raise DeviceError(
    216         f"Operation {operation.name} cannot be used after other Operations have already been applied "
    217         f"on a {self.short_name} device."
    218     )
    220 if isinstance(operation, QubitStateVector):
--> 221     self._apply_state_vector(operation.parameters[0], operation.wires)
    222 elif isinstance(operation, BasisState):
    223     self._apply_basis_state(operation.parameters[0], operation.wires)

File ~/Qottmann/Xanadu/pennylane/pennylane/devices/default_qubit.py:620, in DefaultQubit._apply_state_vector(self, state, device_wires)
    617 n_state_vector = state.shape[0]
    619 if len(qml.math.shape(state)) != 1 or n_state_vector != 2 ** len(device_wires):
--> 620     raise ValueError("State vector must be of length 2**wires.")
    622 norm = qml.math.linalg.norm(state, ord=2)
    623 if not qml.math.is_abstract(norm):

ValueError: State vector must be of length 2**wires.

Is there an easy fix for that @dwierichs ? And am I right in assuming this does some sort of vectorization under the hood and I can expect some speed ups for that? This would be nice for a QML demo, since for-looping through many state examples takes quite a lot of time.

[dwierichs]

This is indeed just not implemented yet. But I have been working on the templates a bit. Let me get back to you tomorrow with a fix :)

And yes, in some cases using broadcasting this way will speed up your QNode :)

[dwierichs]

The above mentioned #2810 introduces broadcasting to templates. Note that the speedup only can be achieved once #2627 is merged! If it is required immediately, one could checkout the branch of #2627 and copy the changes to AmplitudeEmbedding over to that branch (as they are quite minimal). Otherwise, a bit more patience would be required, I'm afraid :grimacing:

[dwierichs]

I think this should be resolved now and can be closed? @Qottmann