for_loop¶
- for_loop(lower_bound, upper_bound, step)[source]¶
A
qjit()
compatible for-loop decorator for PennyLane/Catalyst.Note
Catalyst can automatically convert Python for loop statements for you. Requires setting
autograph=True
, see theqjit()
function or documentation page for more details.This for-loop representation is a functional version of the traditional for-loop, similar to
jax.cond.fori_loop
. That is, any variables that are modified across iterations need to be provided as inputs/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
for_loop
are given by the following Python pseudo-code:def for_loop(lower_bound, upper_bound, step, loop_fn, *args): for i in range(lower_bound, upper_bound, step): args = loop_fn(i, *args) return args
Unlike
jax.cond.fori_loop
, the step can be negative if it is known at tracing time (i.e. constant). If a non-constant negative step is used, the loop will produce no iterations.- Parameters
lower_bound (int) – starting value of the iteration index
upper_bound (int) – (exclusive) upper bound of the iteration index
step (int) – increment applied to the iteration index at the end of each iteration
- Returns
A wrapper around the loop body function. Note that the loop body function must always have the iteration index as its first argument, which can be used arbitrarily inside the loop body. As the value of the index across iterations is handled automatically by the provided loop bounds, it must not be returned from the function.
- Return type
Callable[[int, …], …]
Example
dev = qml.device("lightning.qubit", wires=1) @qjit @qml.qnode(dev) def circuit(n: int, x: float): def loop_rx(i, x): # perform some work and update (some of) the arguments qml.RX(x, wires=0) # update the value of x for the next iteration return jnp.sin(x) # apply the for loop final_x = for_loop(0, n, 1)(loop_rx)(x) return qml.expval(qml.PauliZ(0)), final_x
>>> circuit(7, 1.6) [array(0.97926626), array(0.55395718)]