while_loop¶
- while_loop(cond_fn)[source]¶
A
qjit()
compatible while-loop decorator for PennyLane/Catalyst.This decorator provides a functional version of the traditional while loop, similar to
jax.lax.while_loop
. That is, any variables that are modified across iterations need to be provided as inputs and outputs to the loop body function:Input arguments contain the value of a variable at the start of an iteration
Output arguments contain the value at the end of the iteration. The outputs are then fed back as inputs to the next iteration.
The final iteration values are also returned from the transformed function.
This form of control flow can also be called from the Python interpreter without needing to use
qjit()
.The semantics of
while_loop
are given by the following Python pseudo-code:def while_loop(cond_fun, body_fun, *args): while cond_fun(*args): args = body_fn(*args) return args
- Parameters
cond_fn (Callable) – the condition function in the while loop
- Returns
A wrapper around the while-loop function.
- Return type
Callable
- Raises
TypeError – Invalid return type of the condition expression.
Example
dev = qml.device("lightning.qubit", wires=1) @qjit @qml.qnode(dev) def circuit(x: float): @while_loop(lambda x: x < 2.0) def loop_rx(x): # perform some work and update (some of) the arguments qml.RX(x, wires=0) return x ** 2 # apply the while loop final_x = loop_rx(x) return qml.expval(qml.PauliZ(0)), final_x
>>> circuit(1.6) [array(-0.02919952), array(2.56)]