diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 64470febd79..784196a4a91 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -117,6 +117,9 @@ [-4.20735506e-01, 4.20735506e-01]]) ``` +* `qml.adjoint` can be used with the `qml.qjit` decorator. + [(#4725)](https://github.com/PennyLaneAI/pennylane/pull/4725) + * `qml.ctrl` can be used with the `qml.qjit` decorator. [(#4726)](https://github.com/PennyLaneAI/pennylane/pull/4726) diff --git a/pennylane/compiler/__init__.py b/pennylane/compiler/__init__.py index 2e95b798f37..0f3dfaaccf4 100644 --- a/pennylane/compiler/__init__.py +++ b/pennylane/compiler/__init__.py @@ -63,6 +63,8 @@ .. autosummary:: :toctree: api + ~adjoint + ~ctrl ~grad ~jacobian diff --git a/pennylane/ops/op_math/adjoint.py b/pennylane/ops/op_math/adjoint.py index 486ef22b249..e9fff121aac 100644 --- a/pennylane/ops/op_math/adjoint.py +++ b/pennylane/ops/op_math/adjoint.py @@ -21,6 +21,8 @@ from pennylane.operation import Observable, Operation, Operator from pennylane.queuing import QueuingManager from pennylane.tape import make_qscript +from pennylane.compiler import compiler +from pennylane.compiler.compiler import CompileError from .symbolicop import SymbolicOp @@ -28,6 +30,7 @@ # pylint: disable=no-member def adjoint(fn, lazy=True): """Create the adjoint of an Operator or a function that applies the adjoint of the provided function. + :func:`~.qjit` compatible. Args: fn (function or :class:`~.operation.Operator`): A single operator or a quantum function that @@ -36,6 +39,7 @@ def adjoint(fn, lazy=True): Keyword Args: lazy=True (bool): If the transform is behaving lazily, all operations are wrapped in a ``Adjoint`` class and handled later. If ``lazy=False``, operation-specific adjoint decompositions are first attempted. + Setting ``lazy=False`` is not supported when used with :func:`~.qjit`. Returns: (function or :class:`~.operation.Operator`): If an Operator is provided, returns an Operator that is the adjoint. @@ -44,7 +48,17 @@ def adjoint(fn, lazy=True): .. note:: - The adjoint and inverse are identical for unitary gates, but not in general. For example, quantum channels and observables may have different adjoint and inverse operators. + The adjoint and inverse are identical for unitary gates, but not in general. For example, quantum channels and + observables may have different adjoint and inverse operators. + + .. note:: + + When used with :func:`~.qjit`, this function only supports the Catalyst compiler. + See :func:`catalyst.adjoint` for more details. + + Please see the Catalyst :doc:`quickstart guide `, + as well as the :doc:`sharp bits and debugging tips ` + page for an overview of the differences between Catalyst and PennyLane. .. note:: @@ -101,6 +115,34 @@ def circuit(a): >>> print(qml.draw(circuit)(0.2)) 0: ──RX(0.20)──SX──SX†──RX(0.20)†─┤ + **Example with compiler** + + The adjoint used in a compilation context can be applied on control flow. + + .. code-block:: python + + dev = qml.device("lightning.qubit", wires=1) + + @qml.qjit + @qml.qnode(dev) + def workflow(theta, n, wires): + def func(): + @qml.for_loop(0, n, 1) + def loop_fn(i): + qml.RX(theta, wires=wires) + + loop_fn() + qml.adjoint(func)() + return qml.probs() + + >>> workflow(jnp.pi/2, 3, 0) + [1.00000000e+00 7.39557099e-32] + + .. warning:: + + The Catalyst adjoint function does not support performing the adjoint + of quantum functions that contain mid-circuit measurements. + .. details:: :title: Lazy Evaluation @@ -117,6 +159,12 @@ def circuit(a): Adjoint(S)(wires=[0]) """ + if active_jit := compiler.active_compiler(): + if lazy is False: + raise CompileError("Setting lazy=False is not supported with qjit.") + available_eps = compiler.AvailableCompilers.names_entrypoints + ops_loader = available_eps[active_jit]["ops"].load() + return ops_loader.adjoint(fn) if isinstance(fn, Operator): return Adjoint(fn) if lazy else _single_op_eager(fn, update_queue=True) if not callable(fn): diff --git a/tests/test_compiler.py b/tests/test_compiler.py index e049880522e..c9d404688fe 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -17,6 +17,7 @@ # pylint: disable=import-outside-toplevel import pytest import pennylane as qml +from pennylane.compiler.compiler import CompileError from pennylane import numpy as np @@ -209,6 +210,45 @@ def circuit(x: float): result_header = "func.func private @circuit(%arg0: tensor) -> tensor" assert result_header in mlir_str + def test_qjit_adjoint(self): + """Test JIT compilation with adjoint support""" + dev = qml.device("lightning.qubit", wires=2) + + @qml.qjit + @qml.qnode(device=dev) + def workflow_cl(theta, wires): + def func(): + qml.RX(theta, wires=wires) + + qml.adjoint(func)() + return qml.probs() + + @qml.qnode(device=dev) + def workflow_pl(theta, wires): + def func(): + qml.RX(theta, wires=wires) + + qml.adjoint(func)() + return qml.probs() + + assert jnp.allclose(workflow_cl(0.1, [1]), workflow_pl(0.1, [1])) + + def test_qjit_adjoint_lazy(self): + """Test that Lazy kwarg is not supported.""" + dev = qml.device("lightning.qubit", wires=2) + + @qml.qjit + @qml.qnode(device=dev) + def workflow(theta, wires): + def func(): + qml.RX(theta, wires=wires) + + qml.adjoint(func, lazy=False)() + return qml.probs() + + with pytest.raises(CompileError, match="Setting lazy=False is not supported with qjit."): + workflow(0.1, [1]) + def test_control(self): """Test that control works with qjit.""" dev = qml.device("lightning.qubit", wires=2)