qml.capture¶

This module implements PennyLane’s capturing mechanism for hybrid quantum-classical programs.

Warning

This module is experimental and will change significantly in the future.

`disable`()	Disable the capturing mechanism of hybrid quantum-classical programs in a PennyLane Program Representation (plxpr).
`enable`()	Enable the capturing mechanism of hybrid quantum-classical programs in a PennyLane Program Representation (plxpr).
`enabled`()	Return whether the capturing mechanism of hybrid quantum-classical programs in a PennyLane Program Representation (plxpr) is enabled.
`pause`()	Temporarily stop program capture.
`determine_abstracted_axes`(args)	Compute the abstracted axes and extract the abstract shapes from the arguments.
`expand_plxpr_transforms`(f)	Function for applying transforms to plxpr.
`eval_jaxpr`(jaxpr, consts, *args)	A version of `jax.core.eval_jaxpr` that can handle creating arrays with dynamic shapes.
`run_autograph`(fn)	Decorator that converts the given function into graph form.
`disable_autograph`	Context decorator that disables AutoGraph for the given function/context.
`PlxprInterpreter`()	A base class for defining plxpr interpreters.
`FlatFn`(f[, in_tree])	Wrap a function so that it caches the pytree shape of the output into the `out_tree` property, so that the results can be repacked later.
`make_plxpr`(func[, static_argnums, autograph])	Takes a function and returns a `Callable` that, when called, produces a PLxPR representing the function with the given args.
`register_custom_staging_rule`(primitive, ...)	Register a custom staging rule for a primitive, where the output should match the variables retrieved by `get_outvars_from_params`.

The primitives submodule offers easy access to objects with jax dependencies such as primitives and abstract types. It is not available with import pennylane, but the contents can be accessed via manual import from pennylane.capture.primitives import *.

`AbstractOperator`()	An operator captured into plxpr.
`AbstractMeasurement`(abstract_eval[, ...])	An abstract measurement.
`adjoint_transform_prim`	A subclass for JAX's Primitive that differentiates between different classes of primitives.
`cond_prim`	A subclass for JAX's Primitive that differentiates between different classes of primitives.
`ctrl_transform_prim`	A subclass for JAX's Primitive that differentiates between different classes of primitives.
`for_loop_prim`	A subclass for JAX's Primitive that differentiates between different classes of primitives.
`qnode_prim`	A subclass for JAX's Primitive that differentiates between different classes of primitives.
`while_loop_prim`	A subclass for JAX's Primitive that differentiates between different classes of primitives.

See also:

plxpr_to_tape(plxpr, consts, *args[, shots])

Convert a plxpr into a tape.

To activate and deactivate the new PennyLane program capturing mechanism, use the switches qml.capture.enable and qml.capture.disable. Whether or not the capturing mechanism is currently being used can be queried with qml.capture.enabled. By default, the mechanism is disabled:

>>> import pennylane as qml
>>> qml.capture.enabled()
False
>>> qml.capture.enable()
>>> qml.capture.enabled()
True
>>> qml.capture.disable()
>>> qml.capture.enabled()
False

Custom Operator Behaviour

Any operator that inherits from Operator gains a default ability to be captured in a Jaxpr. Any positional argument is bound as a tracer, wires are processed out into individual tracers, and any keyword arguments are passed as keyword metadata.

class MyOp1(qml.operation.Operator):

    def __init__(self, arg1, wires, key=None):
        super().__init__(arg1, wires=wires)

def qfunc(a):
    MyOp1(a, wires=(0,1), key="a")

qml.capture.enable()
print(jax.make_jaxpr(qfunc)(0.1))

{ lambda ; a:f32[]. let
    _:AbstractOperator() = MyOp1[key=a n_wires=2] a 0 1
in () }

But an operator developer may need to override custom behavior for calling cls._primitive.bind (where cls indicates the class) if:

The operator does not accept wires, like SymbolicOp or CompositeOp.
The operator needs to enforce a data / metadata distinction, like PauliRot.

In such cases, the operator developer can override cls._primitive_bind_call, which will be called when constructing a new class instance instead of type.__call__. For example,

class JustMetadataOp(qml.operation.Operator):

    def __init__(self, metadata):
        super().__init__(wires=[])
        self._metadata = metadata

    @classmethod
    def _primitive_bind_call(cls, metadata):
        return cls._primitive.bind(metadata=metadata)


def qfunc():
    JustMetadataOp("Y")

qml.capture.enable()
print(jax.make_jaxpr(qfunc)())

{ lambda ; . let _:AbstractOperator() = JustMetadataOp[metadata=Y]  in () }

As you can see, the input "Y", while being passed as a positional argument, is converted to metadata within the custom _primitive_bind_call method.

If needed, developers can also override the implementation method of the primitive like was done with Controlled. Controlled needs to do so to handle packing and unpacking the control wires.

class MyCustomOp(qml.operation.Operator):
    pass

@MyCustomOp._primitive.def_impl
def _(*args, **kwargs):
    return type.__call__(MyCustomOp, *args, **kwargs)

code/qml_capture

Download Python script

Download Notebook

View on GitHub

qml.capture¶

Contents