qml.capture.make_plxpr

make_plxpr(func, static_argnums=(), **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

Kwargs:

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).

Returns

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

Return type

Callable

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,) }

code/api/pennylane.capture.make_plxpr

 

Download Python script

 

Download Notebook

 

View on GitHub