qml.capture.make_plxpr

make_plxpr(func, static_argnums=(), autograph=True, **kwargs)

Takes a function and returns a Callable that, when called, produces a PLxPR representing the function with the given args.

This function relies on jax.make_jaxpr as part of creating the representation. Any keyword arguments passed to make_plxpr that are not directly used in the function will be passed to make_jaxpr.

Parameters

func (Callable) – the Callable to be captured

Keyword Arguments

• static_argnums (Union(int, Sequence[int])) – optional, an int or collection of ints that specify which positional arguments to treat as static (trace- and compile-time constant).

• autograph (bool) – whether to use AutoGraph to convert Python control flow to native PennyLane control flow. Defaults to True.

Returns

function that, when called, returns the PLxPR representation of func for the specified inputs.

Return type

Callable

Note

More details on using AutoGraph are provided under Usage Details.

There are some limitations and sharp bits regarding AutoGraph; to better understand supported behaviour and limitations, see https://docs.pennylane.ai/en/stable/development/autograph.html

Example

qml.capture.enable()

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

@qml.qnode(dev)
def circ(x):
    qml.RX(x, 0)
    qml.Hadamard(0)
    return qml.expval(qml.X(0))

plxpr = qml.capture.make_plxpr(circ)(1.2)

>>> print(plxpr)
{ lambda ; a:f32[]. let
    b:f32[] = qnode[
        device=<default.qubit device (wires=1) at 0x152a6f010>
        n_consts=0
        qfunc_jaxpr={ lambda ; c:f32[]. let
            _:AbstractOperator() = RX[n_wires=1] c 0
            _:AbstractOperator() = Hadamard[n_wires=1] 0
            d:AbstractOperator() = PauliX[n_wires=1] 0
            e:AbstractMeasurement(n_wires=None) = expval_obs d
          in (e,) }
        qnode=<QNode: device='<default.qubit device (wires=1) at 0x152a6f010>', interface='auto', diff_method='best'>
        qnode_kwargs={'diff_method': 'best', 'grad_on_execution': 'best', 'cache': False, 'cachesize': 10000, 'max_diff': 1, 'device_vjp': False, 'mcm_method': None, 'postselect_mode': None}
        shots=Shots(total=None)
    ] a
  in (b,) }

Usage Details

The autograph argument is True by default, converting Pythonic control flow to PennyLane supported control flow. This requires the diastatic-malt package, a standalone fork of the AutoGraph module in TensorFlow (official documentation ).

Note

There are some limitations and sharp bits regarding AutoGraph; to better understand supported behaviour and limitations, see https://docs.pennylane.ai/en/stable/development/autograph.html

On its own, capture of standard Python control flow is not supported:

def fn(x):
    if x > 5:
        return x+1
    return x+2

For this function, capture doesn’t work without autograph:

>>> plxpr_fn = qml.capture.make_plxpr(fn, autograph=False)
>>> plxpr = plxpr_fn(3)
TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].

With AutoGraph, the control flow is automatically converted to the native PennyLane control flow implementation, and succeeds:

>>> plxpr_fn = qml.capture.make_plxpr(fn)
>>> plxpr = plxpr_fn(3)
>>> plxpr
{ lambda ; a:i64[]. let
    b:bool[] = gt a 5
    _:bool[] c:i64[] = cond[
      args_slice=slice(4, None, None)
      consts_slices=[slice(2, 3, None), slice(3, 4, None)]
      jaxpr_branches=[{ lambda a:i64[]; . let  in (True, a) }, { lambda a:i64[]; . let b:i64[] = add a 2 in (True, b) }]
    ] b True a a
  in (c,) }

We can evaulate this to get the results:

>>> jax.core.eval_jaxpr(plxpr.jaxpr, plxpr.consts, 2)
[Array(4, dtype=int64, weak_type=True)]

>>> jax.core.eval_jaxpr(plxpr.jaxpr, plxpr.consts, 7)
[Array(8, dtype=int64, weak_type=True)]

code/api/pennylane.capture.make_plxpr

 

Download Python script

 

Download Notebook

 

View on GitHub