# Copyright 2024 Xanadu Quantum Technologies Inc.# Licensed under the Apache License, Version 2.0 (the "License");# you may not use this file except in compliance with the License.# You may obtain a copy of the License at# http://www.apache.org/licenses/LICENSE-2.0# Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License."""Contains the switches to (de)activate the capturing mechanism, and astatus reporting function on whether it is enabled or not."""fromcollections.abcimportCallablefromcontextlibimportcontextmanagerhas_jax=Truetry:importjax# pylint: disable=unused-importexceptImportError:has_jax=Falsedef_make_switches()->tuple[Callable[[],None],Callable[[],None],Callable[[],bool]]:r"""Create three functions, corresponding to an activation switch, a deactivation switch and a status query, in that order. .. note:: While the internal variable is named in some context, this function can be used to make switches for any context. """_FEATURE_ENABLED=False# since this changes what happens with tracing, we need to turn the behaviour# off by default to preserve our ability to jit pennylane circuits.defenable_fn()->None:"""Enable the capturing mechanism of hybrid quantum-classical programs in a PennyLane Program Representation (plxpr)."""ifnothas_jax:raiseImportError("plxpr requires JAX to be installed.")nonlocal_FEATURE_ENABLED_FEATURE_ENABLED=Truedefdisable_fn()->None:"""Disable the capturing mechanism of hybrid quantum-classical programs in a PennyLane Program Representation (plxpr)."""nonlocal_FEATURE_ENABLED_FEATURE_ENABLED=Falsedefstatus_fn()->bool:"""Return whether the capturing mechanism of hybrid quantum-classical programs in a PennyLane Program Representation (plxpr) is enabled."""nonlocal_FEATURE_ENABLEDreturn_FEATURE_ENABLEDreturnenable_fn,disable_fn,status_fnenable,disable,enabled=_make_switches()
[docs]@contextmanagerdefpause():"""Temporarily stop program capture. >>> def f(): ... with qml.capture.pause(): ... qml.X(0) ... return 2 >>> jax.make_jaxpr(f)() { lambda ; . let in (2,) } """disable()try:yieldfinally:enable()