Operations¶

Catalyst allows you to usequantum operationsavailable in PennyLane either via native support by the runtime or PennyLane’s decomposition rules.Theqml.adjoint() andqml.ctrl() functions inPennyLane are also supported via the decomposition mechanism in Catalyst. For example,

@qml.qnode(qml.device("lightning.qubit",wires=2))defcircuit():qml.Rot(0.3,0.4,0.5,wires=0)qml.adjoint(qml.SingleExcitation(jnp.pi/3,wires=[0,1]))returnqml.state()

In addition, you can qjit mostPennyLane templates to easily construct and evaluatemore complex quantum circuits.

Observables¶

The Catalyst has support forPennyLane observables.

For example, the following circuit is a QJIT compatible function that calculates the expectation value ofa tensor product of aqml.PauliX,qml.Hadamard andqml.Hermitian observables.

@qml.qnode(qml.device("lightning.qubit",wires=3))defcircuit(x,y):qml.RX(x,0)qml.RX(y,1)qml.CNOT([0,2])qml.CNOT([1,2])h_matrix=jnp.array([[complex(1.0,0.0),complex(2.0,0.0)],[complex(2.0,0.0),complex(-1.0,0.0)]])returnqml.expval(qml.PauliX(0)@qml.Hadamard(1)@qml.Hermitian(h_matrix,2))

Measurements¶

Most PennyLanemeasurement processesare supported in Catalyst, although not all features are supported for all measurement types.

For bothqml.sample() andqml.counts() omitting the wiresparameters produces samples on all declared qubits in the same format as in PennyLane.

Counts are returned a bit differently, namely as a pair of arrays representing a dictionary from basis statesto the number of observed samples. We thus have to do a bit of extra work to display them nicely.Note that the basis states are represented in their equivalent binary integer representation, inside of afloat data type. This way they are compatible with eigenvalue sampling, but this may change in the future.

@qjit@qml.qnode(qml.device("lightning.qubit",wires=2,shots=1000))defcounts():qml.Rot(0.1,0.2,0.3,wires=[0])returnqml.counts(wires=[0])basis_states,counts=counts()

>>>{format(int(state),'01b'):countforstate,countinzip(basis_states,counts)}{'0': 985, '1': 15}

You can specify the number of shots to be used in sample-based measurements when you create a device.qml.sample() andqml.counts() willautomatically use the device’sshots parameter when performing measurements.In the following example, the number of shots is set to\(500\) in the device instantiation.

@qjit@qml.qnode(qml.device("lightning.qubit",wires=3,shots=500))defcircuit(params):qml.RX(params[0],wires=0)qml.RX(params[1],wires=1)qml.RZ(params[2],wires=2)return(qml.sample(),qml.counts(),qml.expval(qml.PauliZ(0)),qml.var(qml.PauliZ(0)),qml.probs(wires=[0,1]),qml.state(),)

>>>circuit([0.3,0.5,0.7])(Array([[0, 0, 0],        [0, 0, 0],        [0, 0, 0],        ...,        [0, 0, 0],        [0, 0, 0],        [0, 0, 0]], dtype=int64), (Array([0, 1, 2, 3, 4, 5, 6, 7], dtype=int64),  Array([453,   0,  31,   0,  16,   0,   0,   0], dtype=int64)), Array(0.936, dtype=float64), Array(0.138816, dtype=float64), Array([0.926, 0.048, 0.026, 0.   ], dtype=float64), Array([ 0.89994966-0.32850727j,  0.        +0.j        ,        -0.08388168-0.22979488j,  0.        +0.j        ,        -0.04964902-0.13601409j,  0.        +0.j        ,        -0.0347301 +0.01267748j,  0.        +0.j        ],      dtype=complex128))

The PennyLane projective mid-circuit measurement is also supported in Catalyst.measure() is a QJIT compatible mid-circuit measurement for Catalyst that onlyrequires a list of wires that the measurement process acts on.

fromcatalystimportmeasure

In the following example,m will be equal toTrue if wire\(0\) is rotated by\(180\) degrees.

@qjit@qml.qnode(qml.device("lightning.qubit",wires=2))defcircuit(x):qml.RX(x,wires=0)m=measure(wires=0)returnm

>>>circuit(jnp.pi)Array(True, dtype=bool)>>>circuit(0.0)Array(False, dtype=bool)

Just-in-time¶

In just-in-time (JIT), the compilation is triggered at the call site the first timethe quantum function is executed. For example,circuit is compiled as early as the first call.

@qjit@qml.qnode(qml.device("lightning.qubit",wires=2))defcircuit(theta):qml.Hadamard(wires=0)qml.RX(theta,wires=1)qml.CNOT(wires=[0,1])returnqml.expval(qml.PauliZ(wires=1))

>>>circuit(0.5)# the first call, compilation occurs hereArray(0., dtype=float64)>>>circuit(0.5)# the precompiled quantum function is calledArray(0., dtype=float64)

Ahead-of-time¶

An alternative is to trigger the compilation without specifying any concrete values for the functionparameters. This works by specifying the argument signature right in the function definition, whichwill trigger compilation “ahead-of-time” (AOT) before the program is executed. We can use both builtinPython scalar types, as well as the specialShapedArray type that JAX uses to represent the shapeand data type of a tensor:

fromjax.coreimportShapedArray@qjit# compilation happens at definition@qml.qnode(qml.device("lightning.qubit",wires=2))defcircuit(x:complex,z:ShapedArray(shape=(3,),dtype=jnp.float64)):theta=jnp.abs(x)qml.RY(theta,wires=0)qml.Rot(z[0],z[1],z[2],wires=0)returnqml.state()

>>>circuit(0.2j,jnp.array([0.3,0.6,0.9]))# calls precompiled functionArray([0.75634905-0.52801002j, 0. +0.j,   0.35962678+0.14074839j, 0. +0.j], dtype=complex128)

At this stage the compilation already happened, so the execution ofcircuit calls the compiled function directly onthe first call, resulting in faster initial execution. Note that implicit type promotion for most datatypes are allowedin the compilation as long as it doesn’t lead to a loss of data.

Conditionals¶

cond() is a functional version of the traditional if-else conditional for Catalyst.This means that each execution path, aTrue branch and aFalse branch, is provided as aseparate function. Both functions will be traced during compilation, but only one of them the will beexecuted at runtime, depending of the value of a Boolean predicate. The JAX equivalent is thejax.lax.cond function, but this version is optimized to work with quantum programs in PennyLane.

Note thatcond() can also be used outside of theqjit() context for better interoperability with PennyLane.

Values produced inside the scope of a conditional can be returned to the outside context, butthe return type signature of each branch must be identical. If no values are returned, theFalse branch is optional. Refer to the example below to learn more about the syntax of thisdecorator.

@cond(predicate:bool)defconditional_fn():# do something when the predicate is truereturn"optionally return some value"@conditional_fn.otherwisedefconditional_fn():# optionally define an alternative execution pathreturn"if provided, return types need to be identical in both branches"ret_val=conditional_fn()# must invoke the defined function

Loops¶

for_loop() andwhile_loop() are functional versions of the traditional for- andwhile-loop for Catalyst. That is, any variables that are modified across iterations need to beprovided as inputs and outputs to the loop body function. Input arguments contain the value of avariable at the start of an iteration, while output arguments contain the value at the end of theiteration. The outputs are then fed back as inputs to the next iteration. The final iteration valuesare also returned from the transformed function.

for_loop() andwhile_loop() can also be interpreted without needing to compile its surrounding context.

The for-loop statement:

Thefor_loop() executes a fixed number of iterations as indicated via the values specifiedin its header: alower_bound, anupper_bound, and astep size.

The loop body function must always have the iteration index (in the below examplei) as itsfirst argument and its value can be used arbitrarily inside the loop body. As the value of the indexacross iterations is handled automatically by the provided loop bounds, it must not be returned fromthe body function.

@for_loop(lower_bnd,upper_bnd,step)defloop_body(i,*args):# code to be executed over index i starting# from lower_bnd to upper_bnd - 1 by stepreturnargsfinal_args=loop_body(init_args)

The semantics offor_loop() are given by the following Python implementation:

foriinrange(lower_bnd,upper_bnd,step):args=body_fn(i,*args)

The while-loop statement:

Thewhile_loop(), on the other hand, is able to execute an arbitrary number of iterations,until the condition function specified in its header returnsFalse.

The loop condition is evaluated every iteration and can be any callable with an identical signatureas the loop body function. The return type of the condition function must be a Boolean.

@while_loop(lambda*args:"some condition")defloop_body(*args):# perform some work and update (some of) the argumentsreturnargsfinal_args=loop_body(init_args)

Optimizers¶

You can develop your own optimization algorithm using thegrad() method, control-flow operators that arecompatible with QJIT, or by utilizing differentiable optimizers inOptax.

For example, you can useoptax.sgd in a QJIT workflow to calculatethe gradient descent optimizer. The following example shows a simple use case of thisfeature in Catalyst.

Theoptax.sgd gets a smooth function of the formgd_fun(params,*args,**kwargs)and calculates either just the value or both the value and gradient of the function depending onthe value ofvalue_and_grad argument. To optimize params iteratively, you later need to usejax.lax.fori_loop to loop over the gradient descent steps.

importoptaxfromjax.laximportfori_loopdev=qml.device("lightning.qubit",wires=1)@qml.qnode(dev)defcircuit(param):qml.Hadamard(0)qml.RY(param,wires=0)returnqml.expval(qml.PauliZ(0))@qjitdefworkflow():defgd_fun(param):diff=grad(circuit,argnums=0)returncircuit(param),diff(param)opt=optax.sgd(learning_rate=0.4)defgd_update(i,args):param,state=args_,gradient=gd_fun(param)(updates,state)=opt.update(gradient,state)param=optax.apply_updates(param,updates)return(param,state)param=0.1state=opt.init(param)(param,_)=fori_loop(0,100,gd_update,(param,state))returnparam

>>>workflow()Array(1.57079633, dtype=float64)