qml.while_loop

while_loop(cond_fn)

A qjit() compatible for-loop for PennyLane programs. When used without qjit(), this function will fall back to a standard Python for loop.

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.

The semantics of while_loop are given by the following Python pseudocode:

def while_loop(cond_fn, body_fn, *args):
    while cond_fn(*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

CompileError – if the compiler is not installed

See also

for_loop(), qjit()

Example

dev = qml.device("lightning.qubit", wires=1)

@qml.qnode(dev)
def circuit(x: float):

    @qml.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
    loop_rx(x)

    return qml.expval(qml.Z(0))

>>> circuit(1.6)
-0.02919952

while_loop is also qjit() compatible; when used with the qjit() decorator, the while loop will not be unrolled, and instead will be captured as-is during compilation and executed during runtime:

>>> qml.qjit(circuit)(1.6)
Array(-0.02919952, dtype=float64)

code/api/pennylane.while_loop

 

Download Python script

 

Download Notebook

 

View on GitHub