diff --git a/doc/code/qml.rst b/doc/code/qml.rst index fa32159fefb..c568ea5ab26 100644 --- a/doc/code/qml.rst +++ b/doc/code/qml.rst @@ -6,4 +6,4 @@ qml .. automodapi:: pennylane :no-heading: :include-all-objects: - :skip: Version, SimpleSpec, plugin_devices, plugin_converters, default_config, reload, version_info, defaultdict + :skip: plugin_converters, default_config, version_info, defaultdict diff --git a/doc/development/deprecations.rst b/doc/development/deprecations.rst index a1ef7ce7277..7f469e13359 100644 --- a/doc/development/deprecations.rst +++ b/doc/development/deprecations.rst @@ -9,6 +9,12 @@ deprecations are listed below. Pending deprecations -------------------- +* All of the legacy devices (any with the name ``default.qubit.{autograd,torch,tf,jax,legacy}``) are deprecated. Use ``default.qubit`` instead, + as it supports backpropagation for the many backends the legacy devices support. + + - Deprecated in v0.38 + - Will be removed in v0.39 + * The ``decomp_depth`` argument in ``qml.device`` is deprecated. - Deprecated in v0.38 diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index de11157582f..8dc02dca2b6 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -21,9 +21,6 @@ into multiple single-term measurements on a single tape, providing better support for simulators [(#5987)](https://github.com/PennyLaneAI/pennylane/pull/5987) - that can handle non-commuting observables but don't natively support multi-term observables. - [(#5884)](https://github.com/PennyLaneAI/pennylane/pull/5884) - * `SProd.terms` now flattens out the terms if the base is a multi-term observable. [(#5885)](https://github.com/PennyLaneAI/pennylane/pull/5885) @@ -37,6 +34,9 @@

Improvements πŸ› 

+* Added the `compute_sparse_matrix` method for `qml.ops.qubit.BasisStateProjector`. + [(#5790)](https://github.com/PennyLaneAI/pennylane/pull/5790) + * `StateMP.process_state` defines rules in `cast_to_complex` for complex casting, avoiding a superfluous state vector copy in Lightning simulations [(#5995)](https://github.com/PennyLaneAI/pennylane/pull/5995) @@ -60,6 +60,51 @@ * Molecules and Hamiltonians can now be constructed for all the elements present in the periodic table. [(#5821)](https://github.com/PennyLaneAI/pennylane/pull/5821) +* `qml.for_loop` and `qml.while_loop` now fallback to standard Python control + flow if `@qjit` is not present, allowing the same code to work with and without + `@qjit` without any rewrites. + [(#6014)](https://github.com/PennyLaneAI/pennylane/pull/6014) + + ```python + dev = qml.device("lightning.qubit", wires=3) + + @qml.qnode(dev) + def circuit(x, n): + + @qml.for_loop(0, n, 1) + def init_state(i): + qml.Hadamard(wires=i) + + init_state() + + @qml.for_loop(0, n, 1) + def apply_operations(i, x): + qml.RX(x, wires=i) + + @qml.for_loop(i + 1, n, 1) + def inner(j): + qml.CRY(x**2, [i, j]) + + inner() + return jnp.sin(x) + + apply_operations(x) + return qml.probs() + ``` + + ```pycon + >>> print(qml.draw(circuit)(0.5, 3)) + 0: ──H──RX(0.50)─╭●────────╭●─────────────────────────────────────── Probs + 1: ──H───────────╰RY(0.25)─│──────────RX(0.48)─╭●─────────────────── Probs + 2: ──H─────────────────────╰RY(0.25)───────────╰RY(0.23)──RX(0.46)── Probs + >>> circuit(0.5, 3) + array([0.125 , 0.125 , 0.09949758, 0.15050242, 0.07594666, + 0.11917543, 0.08942104, 0.21545687]) + >>> qml.qjit(circuit)(0.5, 3) + Array([0.125 , 0.125 , 0.09949758, 0.15050242, 0.07594666, + 0.11917543, 0.08942104, 0.21545687], dtype=float64) + ``` + * The `qubit_observable` function is modified to return an ascending wire order for molecular Hamiltonians. [(#5950)](https://github.com/PennyLaneAI/pennylane/pull/5950) @@ -136,6 +181,10 @@ Instead, use `pennylane.gradients.classical_fisher` and `pennylane.gradients.quantum_fisher`. [(#5985)](https://github.com/PennyLaneAI/pennylane/pull/5985) +* The legacy devices `default.qubit.{autograd,torch,tf,jax,legacy}` are deprecated. + Instead, use `default.qubit` as it now supports backpropagation through the several backends. + [(#5997)](https://github.com/PennyLaneAI/pennylane/pull/5997) +

Documentation πŸ“

* Improves the docstring for `QuantumScript.expand` and `qml.tape.tape.expand_tape`. @@ -177,10 +226,12 @@ Pietropaolo Frisoni, Emiliano Godinez, Austin Huang, Renke Huang, +Josh Izaac, Soran Jahangiri, Christina Lee, William Maxwell, Vincent Michaud-Rioux, +Anurav Modak, Mudit Pandey, Erik Schultheis, nate stemen. diff --git a/pennylane/__init__.py b/pennylane/__init__.py index f81c1b90083..3f0c484fc4f 100644 --- a/pennylane/__init__.py +++ b/pennylane/__init__.py @@ -15,13 +15,9 @@ This is the top level module from which all basic functions and classes of PennyLane can be directly imported. """ -from importlib import reload, metadata -from sys import version_info -import warnings import numpy as _np -from semantic_version import SimpleSpec, Version from pennylane.boolean_fn import BooleanFn import pennylane.numpy @@ -155,6 +151,8 @@ import pennylane.noise from pennylane.noise import NoiseModel +from pennylane.devices.device_constructor import device, refresh_devices + # Look for an existing configuration file default_config = Configuration("config.toml") @@ -176,275 +174,10 @@ def __getattr__(name): return pennylane.ops.LinearCombination return pennylane.ops.Hamiltonian - raise AttributeError(f"module 'pennylane' has no attribute '{name}'") - - -def _get_device_entrypoints(): - """Returns a dictionary mapping the device short name to the - loadable entrypoint""" - entries = ( - metadata.entry_points()["pennylane.plugins"] - if version_info[:2] == (3, 9) - # pylint:disable=unexpected-keyword-arg - else metadata.entry_points(group="pennylane.plugins") - ) - return {entry.name: entry for entry in entries} - - -def refresh_devices(): - """Scan installed PennyLane plugins to refresh the device list.""" - - # This function does not return anything; instead, it has a side effect - # which is to update the global plugin_devices variable. - - # We wish to retain the behaviour of a global plugin_devices dictionary, - # as re-importing metadata can be a very slow operation on systems - # with a large number of installed packages. - global plugin_devices # pylint:disable=global-statement - - reload(metadata) - plugin_devices = _get_device_entrypoints() - - -# get list of installed devices -plugin_devices = _get_device_entrypoints() - - -# pylint: disable=protected-access -def device(name, *args, **kwargs): - r""" - Load a device and return the instance. - - This function is used to load a particular quantum device, - which can then be used to construct QNodes. - - PennyLane comes with support for the following devices: - - * :mod:`'default.qubit' `: a simple - state simulator of qubit-based quantum circuit architectures. - - * :mod:`'default.mixed' `: a mixed-state - simulator of qubit-based quantum circuit architectures. - - * ``'lightning.qubit'``: a more performant state simulator of qubit-based - quantum circuit architectures written in C++. - - * :mod:`'default.qutrit' `: a simple - state simulator of qutrit-based quantum circuit architectures. - - * :mod:`'default.qutrit.mixed' `: a - mixed-state simulator of qutrit-based quantum circuit architectures. - - * :mod:`'default.gaussian' `: a simple simulator - of Gaussian states and operations on continuous-variable circuit architectures. - - * :mod:`'default.clifford' `: an efficient - simulator of Clifford circuits. - - * :mod:`'default.tensor' `: a simulator - of quantum circuits based on tensor networks. - - Additional devices are supported through plugins β€” see - the `available plugins `_ for more - details. To list all currently installed devices, run - :func:`qml.about `. - - Args: - name (str): the name of the device to load - wires (int): the number of wires (subsystems) to initialise - the device with. Note that this is optional for certain - devices, such as ``default.qubit`` - - Keyword Args: - config (pennylane.Configuration): a PennyLane configuration object - that contains global and/or device specific configurations. - custom_decomps (Dict[Union(str, Operator), Callable]): Custom - decompositions to be applied by the device at runtime. - decomp_depth (int): For when custom decompositions are specified, - the maximum expansion depth used by the expansion function. - - .. warning:: - - The ``decomp_depth`` argument is deprecated and will be removed in version 0.39. - - All devices must be loaded by specifying their **short-name** as listed above, - followed by the **wires** (subsystems) you wish to initialize. The ``wires`` - argument can be an integer, in which case the wires of the device are addressed - by consecutive integers: - - .. code-block:: python - - dev = qml.device('default.qubit', wires=5) - - def circuit(): - qml.Hadamard(wires=1) - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[3, 4]) - ... - - The ``wires`` argument can also be a sequence of unique numbers or strings, specifying custom wire labels - that the user employs to address the wires: - - .. code-block:: python + if name == "plugin_devices": + return pennylane.devices.device_constructor.plugin_devices - dev = qml.device('default.qubit', wires=['ancilla', 'q11', 'q12', -1, 1]) - - def circuit(): - qml.Hadamard(wires='q11') - qml.Hadamard(wires=['ancilla']) - qml.CNOT(wires=['q12', -1]) - ... - - On some newer devices, such as ``default.qubit``, the ``wires`` argument can be omitted altogether, - and instead the wires will be computed when executing a circuit depending on its contents. - - >>> dev = qml.device("default.qubit") - - Most devices accept a ``shots`` argument which specifies how many circuit executions - are used to estimate stochastic return values. As an example, ``qml.sample()`` measurements - will return as many samples as specified in the shots argument. The shots argument can be - changed on a per-call basis using the built-in ``shots`` keyword argument. Note that the - ``shots`` argument can be a single integer or a list of shot values. - - .. code-block:: python - - dev = qml.device('default.qubit', wires=1, shots=10) - - @qml.qnode(dev) - def circuit(a): - qml.RX(a, wires=0) - return qml.sample(qml.Z(0)) - - >>> circuit(0.8) # 10 samples are returned - array([ 1, 1, 1, 1, -1, 1, 1, -1, 1, 1]) - >>> circuit(0.8, shots=[3, 4, 4]) # default is overwritten for this call - (array([1, 1, 1]), array([ 1, -1, 1, 1]), array([1, 1, 1, 1])) - >>> circuit(0.8) # back to default of 10 samples - array([ 1, -1, 1, 1, -1, 1, 1, 1, 1, 1]) - - When constructing a device, we may optionally pass a dictionary of custom - decompositions to be applied to certain operations upon device execution. - This is useful for enabling support of gates on devices where they would normally - be unsupported. - - For example, suppose we are running on an ion trap device which does not - natively implement the CNOT gate, but we would still like to write our - circuits in terms of CNOTs. On a ion trap device, CNOT can be implemented - using the ``IsingXX`` gate. We first define a decomposition function - (such functions have the signature ``decomposition(*params, wires)``): - - .. code-block:: python - - def ion_trap_cnot(wires, **_): - return [ - qml.RY(np.pi/2, wires=wires[0]), - qml.IsingXX(np.pi/2, wires=wires), - qml.RX(-np.pi/2, wires=wires[0]), - qml.RY(-np.pi/2, wires=wires[0]), - qml.RY(-np.pi/2, wires=wires[1]) - ] - - Next, we create a device, and a QNode for testing. When constructing the - QNode, we can set the expansion strategy to ``"device"`` to ensure the - decomposition is applied and will be viewable when we draw the circuit. - Note that custom decompositions should accept keyword arguments even when - it is not used. - - .. code-block:: python - - # As the CNOT gate normally has no decomposition, we can use default.qubit - # here for expository purposes. - dev = qml.device( - 'default.qubit', wires=2, custom_decomps={"CNOT" : ion_trap_cnot} - ) - - @qml.qnode(dev, expansion_strategy="device") - def run_cnot(): - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.X(1)) - - >>> print(qml.draw(run_cnot)()) - 0: ──RY(1.57)─╭IsingXX(1.57)──RX(-1.57)──RY(-1.57)── - 1: ───────────╰IsingXX(1.57)──RY(-1.57)───────────── - - Some devices may accept additional arguments. For instance, - ``default.gaussian`` accepts the keyword argument ``hbar``, to set - the convention used in the commutation relation :math:`[\x,\p]=i\hbar` - (by default set to 2). - - Please refer to the documentation for the individual devices to see any - additional arguments that might be required or supported. - """ - if name not in plugin_devices: - # Device does not exist in the loaded device list. - # Attempt to refresh the devices, in case the user - # installed the plugin during the current Python session. - refresh_devices() - - if name in plugin_devices: - options = {} - - # load global configuration settings if available - config = kwargs.get("config", default_config) - - if config: - # combine configuration options with keyword arguments. - # Keyword arguments take preference, followed by device options, - # followed by plugin options, followed by global options. - options.update(config["main"]) - options.update(config[name.split(".")[0] + ".global"]) - options.update(config[name]) - - # Pop the custom decomposition keyword argument; we will use it here - # only and not pass it to the device. - custom_decomps = kwargs.pop("custom_decomps", None) - decomp_depth = kwargs.pop("decomp_depth", None) - - if decomp_depth is not None: - warnings.warn( - "The decomp_depth argument is deprecated and will be removed in version 0.39. ", - PennyLaneDeprecationWarning, - ) - else: - decomp_depth = 10 - - kwargs.pop("config", None) - options.update(kwargs) - - # loads the device class - plugin_device_class = plugin_devices[name].load() - - if hasattr(plugin_device_class, "pennylane_requires") and Version( - version() - ) not in SimpleSpec(plugin_device_class.pennylane_requires): - raise DeviceError( - f"The {name} plugin requires PennyLane versions {plugin_device_class.pennylane_requires}, " - f"however PennyLane version {__version__} is installed." - ) - - # Construct the device - dev = plugin_device_class(*args, **options) - - # Once the device is constructed, we set its custom expansion function if - # any custom decompositions were specified. - - if custom_decomps is not None: - if isinstance(dev, pennylane.devices.LegacyDevice): - custom_decomp_expand_fn = pennylane.transforms.create_decomp_expand_fn( - custom_decomps, dev, decomp_depth=decomp_depth - ) - dev.custom_expand(custom_decomp_expand_fn) - else: - custom_decomp_preprocess = ( - pennylane.transforms.tape_expand._create_decomp_preprocessing( - custom_decomps, dev, decomp_depth=decomp_depth - ) - ) - dev.preprocess = custom_decomp_preprocess - - return dev - - raise DeviceError(f"Device {name} does not exist. Make sure the required plugin is installed.") + raise AttributeError(f"module 'pennylane' has no attribute '{name}'") def version(): diff --git a/pennylane/compiler/qjit_api.py b/pennylane/compiler/qjit_api.py index 140bca19b4b..8d26b37422d 100644 --- a/pennylane/compiler/qjit_api.py +++ b/pennylane/compiler/qjit_api.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """QJIT compatible quantum and compilation operations API""" +from collections.abc import Callable from .compiler import ( AvailableCompilers, @@ -377,20 +378,54 @@ def loop_rx(x): ops_loader = compilers[active_jit]["ops"].load() return ops_loader.while_loop(cond_fn) - raise CompileError("There is no active compiler package.") # pragma: no cover + # if there is no active compiler, simply interpret the while loop + # via the Python interpretor. + def _decorator(body_fn: Callable) -> Callable: + """Transform that will call the input ``body_fn`` until the closure variable ``cond_fn`` is met. + Args: + body_fn (Callable): -def for_loop(lower_bound, upper_bound, step): - """A :func:`~.qjit` compatible for-loop for PennyLane programs. + Closure Variables: + cond_fn (Callable): - .. note:: + Returns: + Callable: a callable with the same signature as ``body_fn`` and ``cond_fn``. + """ + return WhileLoopCallable(cond_fn, body_fn) - This function only supports the Catalyst compiler. See - :func:`catalyst.for_loop` for more details. + return _decorator - 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. + +class WhileLoopCallable: # pylint:disable=too-few-public-methods + """Base class to represent a while loop. This class + when called with an initial state will execute the while + loop via the Python interpreter. + + Args: + cond_fn (Callable): the condition function in the while loop + body_fn (Callable): the function that is executed within the while loop + """ + + def __init__(self, cond_fn, body_fn): + self.cond_fn = cond_fn + self.body_fn = body_fn + + def __call__(self, *init_state): + args = init_state + fn_res = args if len(args) > 1 else args[0] if len(args) == 1 else None + + while self.cond_fn(*args): + fn_res = self.body_fn(*args) + args = fn_res if len(args) > 1 else (fn_res,) if len(args) == 1 else () + + return fn_res + + +def for_loop(lower_bound, upper_bound, step): + """A :func:`~.qjit` compatible for-loop for PennyLane programs. When + used without :func:`~.qjit`, this function will fall back to a standard + Python for loop. This decorator provides a functional version of the traditional for-loop, similar to `jax.cond.fori_loop `__. @@ -430,19 +465,14 @@ def for_loop(lower_bound, upper_bound, step, loop_fn, *args): across iterations is handled automatically by the provided loop bounds, it must not be returned from the function. - Raises: - CompileError: if the compiler is not installed - .. seealso:: :func:`~.while_loop`, :func:`~.qjit` **Example** - .. code-block:: python dev = qml.device("lightning.qubit", wires=1) - @qml.qjit @qml.qnode(dev) def circuit(n: int, x: float): @@ -457,10 +487,24 @@ def loop_rx(i, x): # apply the for loop final_x = loop_rx(x) - return qml.expval(qml.Z(0)), final_x + return qml.expval(qml.Z(0)) >>> circuit(7, 1.6) - (array(0.97926626), array(0.55395718)) + array(0.97926626) + + ``for_loop`` is also :func:`~.qjit` compatible; when used with the + :func:`~.qjit` decorator, the for loop will not be unrolled, and instead + will be captured as-is during compilation and executed during runtime: + + >>> qml.qjit(circuit)(7, 1.6) + Array(0.97926626, dtype=float64) + + .. note:: + + Please see the Catalyst :doc:`quickstart guide `, + as well as the :doc:`sharp bits and debugging tips ` + page for an overview of using quantum just-in-time compilation. + """ if active_jit := active_compiler(): @@ -468,4 +512,59 @@ def loop_rx(i, x): ops_loader = compilers[active_jit]["ops"].load() return ops_loader.for_loop(lower_bound, upper_bound, step) - raise CompileError("There is no active compiler package.") # pragma: no cover + # if there is no active compiler, simply interpret the for loop + # via the Python interpretor. + def _decorator(body_fn): + """Transform that will call the input ``body_fn`` within a for loop defined by the closure variables lower_bound, upper_bound, and step. + + Args: + body_fn (Callable): The function called within the for loop. Note that the loop body + function must always have the iteration index as its first + argument, which can be used arbitrarily inside the loop body. As the value of the index + across iterations is handled automatically by the provided loop bounds, it must not be + returned from the function. + + Closure Variables: + lower_bound (int): starting value of the iteration index + upper_bound (int): (exclusive) upper bound of the iteration index + step (int): increment applied to the iteration index at the end of each iteration + + Returns: + Callable: a callable with the same signature as ``body_fn`` + """ + return ForLoopCallable(lower_bound, upper_bound, step, body_fn) + + return _decorator + + +class ForLoopCallable: # pylint:disable=too-few-public-methods + """Base class to represent a for loop. This class + when called with an initial state will execute the while + loop via the Python interpreter. + + Args: + lower_bound (int): starting value of the iteration index + upper_bound (int): (exclusive) upper bound of the iteration index + step (int): increment applied to the iteration index at the end of each iteration + body_fn (Callable): The function called within the for loop. Note that the loop body + function must always have the iteration index as its first + argument, which can be used arbitrarily inside the loop body. As the value of the index + across iterations is handled automatically by the provided loop bounds, it must not be + returned from the function. + """ + + def __init__(self, lower_bound, upper_bound, step, body_fn): + self.lower_bound = lower_bound + self.upper_bound = upper_bound + self.step = step + self.body_fn = body_fn + + def __call__(self, *init_state): + args = init_state + fn_res = args if len(args) > 1 else args[0] if len(args) == 1 else None + + for i in range(self.lower_bound, self.upper_bound, self.step): + fn_res = self.body_fn(i, *args) + args = fn_res if len(args) > 1 else (fn_res,) if len(args) == 1 else () + + return fn_res diff --git a/pennylane/devices/__init__.py b/pennylane/devices/__init__.py index f0750efbd5a..475a84aac0b 100644 --- a/pennylane/devices/__init__.py +++ b/pennylane/devices/__init__.py @@ -149,6 +149,7 @@ def execute(self, circuits, execution_config = qml.devices.DefaultExecutionConfi """ from .execution_config import ExecutionConfig, DefaultExecutionConfig, MCMConfig +from .device_constructor import device, refresh_devices from .device_api import Device from .default_qubit import DefaultQubit @@ -163,5 +164,15 @@ def execute(self, circuits, execution_config = qml.devices.DefaultExecutionConfi from .default_clifford import DefaultClifford from .default_tensor import DefaultTensor from .null_qubit import NullQubit +from .default_qutrit import DefaultQutrit from .default_qutrit_mixed import DefaultQutritMixed from .._device import Device as LegacyDevice +from .._device import DeviceError + + +# pylint: disable=undefined-variable +def __getattr__(name): + if name == "plugin_devices": + return device_constructor.plugin_devices + + raise AttributeError(f"module 'pennylane.devices' has no attribute '{name}'") diff --git a/pennylane/devices/default_qubit_autograd.py b/pennylane/devices/default_qubit_autograd.py index 5ad6b629625..abcc6e0452f 100644 --- a/pennylane/devices/default_qubit_autograd.py +++ b/pennylane/devices/default_qubit_autograd.py @@ -14,12 +14,15 @@ """This module contains an autograd implementation of the :class:`~.DefaultQubitLegacy` reference plugin. """ +import warnings + +from pennylane import PennyLaneDeprecationWarning from pennylane import numpy as pnp from pennylane.devices import DefaultQubitLegacy class DefaultQubitAutograd(DefaultQubitLegacy): - """Simulator plugin based on ``"default.qubit.legacy"``, written using Autograd. + r"""Simulator plugin based on ``"default.qubit.legacy"``, written using Autograd. **Short name:** ``default.qubit.autograd`` @@ -34,6 +37,9 @@ class DefaultQubitAutograd(DefaultQubitLegacy): pip install autograd + .. warning:: + This device is deprecated. Use :class:`~pennylane.devices.DefaultQubit` instead; for example through ``qml.device("default.qubit")``, which now supports backpropagation. + **Example** The ``default.qubit.autograd`` is designed to be used with end-to-end classical backpropagation @@ -104,6 +110,14 @@ def _const_mul(constant, array): return constant * array def __init__(self, wires, *, shots=None, analytic=None): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + PennyLaneDeprecationWarning, + ) + r_dtype = pnp.float64 c_dtype = pnp.complex128 super().__init__(wires, shots=shots, r_dtype=r_dtype, c_dtype=c_dtype, analytic=analytic) diff --git a/pennylane/devices/default_qubit_jax.py b/pennylane/devices/default_qubit_jax.py index f349b8ff7c9..0cae4827e41 100644 --- a/pennylane/devices/default_qubit_jax.py +++ b/pennylane/devices/default_qubit_jax.py @@ -14,6 +14,8 @@ """This module contains a jax implementation of the :class:`~.DefaultQubitLegacy` reference plugin. """ +import warnings + # pylint: disable=ungrouped-imports import numpy as np @@ -34,7 +36,7 @@ class DefaultQubitJax(DefaultQubitLegacy): - """Simulator plugin based on ``"default.qubit.legacy"``, written using jax. + r"""Simulator plugin based on ``"default.qubit.legacy"``, written using jax. **Short name:** ``default.qubit.jax`` @@ -49,6 +51,9 @@ class DefaultQubitJax(DefaultQubitLegacy): pip install jax jaxlib + .. warning:: + This device is deprecated. Use :class:`~pennylane.devices.DefaultQubit` instead; for example through ``qml.device("default.qubit")``, which now supports backpropagation. + **Example** The ``default.qubit.jax`` device is designed to be used with end-to-end classical backpropagation @@ -165,6 +170,14 @@ def circuit(): operations = DefaultQubitLegacy.operations.union({"ParametrizedEvolution"}) def __init__(self, wires, *, shots=None, prng_key=None, analytic=None): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + qml.PennyLaneDeprecationWarning, + ) + if jax.config.read("jax_enable_x64"): c_dtype = jnp.complex128 r_dtype = jnp.float64 diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py index 74da831f24e..f351b2aa455 100644 --- a/pennylane/devices/default_qubit_legacy.py +++ b/pennylane/devices/default_qubit_legacy.py @@ -20,6 +20,7 @@ """ import functools import itertools +import warnings from string import ascii_letters as ABC import numpy as np @@ -77,13 +78,13 @@ def _get_slice(index, axis, num_axes): # pylint: disable=unused-argument class DefaultQubitLegacy(QubitDevice): - """Default qubit device for PennyLane. + r"""Default qubit device for PennyLane. .. warning:: - This is the legacy implementation of DefaultQubit. It has been replaced by - ``qml.devices.DefaultQubit``, which can be accessed with the familiar constructor, - ``qml.device("default.qubit")``. + This is the legacy implementation of DefaultQubit and is deprecated. It has been replaced by + :class:`~pennylane.devices.DefaultQubit`, which can be accessed with the familiar constructor, + ``qml.device("default.qubit")``, and now supports backpropagation. This change will not alter device behaviour for most workflows, but may have implications for plugin developers and users who directly interact with device methods. Please consult @@ -206,6 +207,14 @@ class DefaultQubitLegacy(QubitDevice): def __init__( self, wires, *, r_dtype=np.float64, c_dtype=np.complex128, shots=None, analytic=None ): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + qml.PennyLaneDeprecationWarning, + ) + super().__init__(wires, shots, r_dtype=r_dtype, c_dtype=c_dtype, analytic=analytic) self._debugger = None diff --git a/pennylane/devices/default_qubit_tf.py b/pennylane/devices/default_qubit_tf.py index 48d7c60b788..2171aa7592e 100644 --- a/pennylane/devices/default_qubit_tf.py +++ b/pennylane/devices/default_qubit_tf.py @@ -15,6 +15,7 @@ reference plugin. """ import itertools +import warnings import numpy as np from packaging.version import Version @@ -42,7 +43,7 @@ class DefaultQubitTF(DefaultQubitLegacy): - """Simulator plugin based on ``"default.qubit.legacy"``, written using TensorFlow. + r"""Simulator plugin based on ``"default.qubit.legacy"``, written using TensorFlow. **Short name:** ``default.qubit.tf`` @@ -57,6 +58,9 @@ class DefaultQubitTF(DefaultQubitLegacy): pip install tensorflow>=2.0 + .. warning:: + This device is deprecated. Use :class:`~pennylane.devices.DefaultQubit` instead; for example through ``qml.device("default.qubit")``, which now supports backpropagation. + **Example** The ``default.qubit.tf`` is designed to be used with end-to-end classical backpropagation @@ -162,6 +166,14 @@ def _asarray(array, dtype=None): return res def __init__(self, wires, *, shots=None, analytic=None): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + qml.PennyLaneDeprecationWarning, + ) + r_dtype = tf.float64 c_dtype = tf.complex128 diff --git a/pennylane/devices/default_qubit_torch.py b/pennylane/devices/default_qubit_torch.py index e0bb276e513..3ed26025b60 100644 --- a/pennylane/devices/default_qubit_torch.py +++ b/pennylane/devices/default_qubit_torch.py @@ -34,6 +34,7 @@ import numpy as np +from pennylane import PennyLaneDeprecationWarning from pennylane.ops.qubit.attributes import diagonal_in_z_basis from . import DefaultQubitLegacy @@ -43,7 +44,7 @@ class DefaultQubitTorch(DefaultQubitLegacy): - """Simulator plugin based on ``"default.qubit.legacy"``, written using PyTorch. + r"""Simulator plugin based on ``"default.qubit.legacy"``, written using PyTorch. **Short name:** ``default.qubit.torch`` @@ -58,6 +59,10 @@ class DefaultQubitTorch(DefaultQubitLegacy): pip install torch>=1.8.0 + .. warning:: + This device is deprecated. Use :class:`~pennylane.devices.DefaultQubit` instead; for example through ``qml.device("default.qubit")``, which now supports backpropagation. + + **Example** The ``default.qubit.torch`` is designed to be used with end-to-end classical backpropagation @@ -165,6 +170,13 @@ def circuit(x): _ndim = staticmethod(lambda tensor: tensor.ndim) def __init__(self, wires, *, shots=None, analytic=None, torch_device=None): + warnings.warn( + f"Use of '{self.short_name}' is deprecated. Instead, use 'default.qubit', " + "which supports backpropagation. " + "If you experience issues, reach out to the PennyLane team on " + "the discussion forum: https://discuss.pennylane.ai/", + PennyLaneDeprecationWarning, + ) # Store if the user specified a Torch device. Otherwise the execute # method attempts to infer the Torch device from the gate parameters. self._torch_device_specified = torch_device is not None diff --git a/pennylane/devices/device_constructor.py b/pennylane/devices/device_constructor.py new file mode 100644 index 00000000000..107243f7718 --- /dev/null +++ b/pennylane/devices/device_constructor.py @@ -0,0 +1,291 @@ +# Copyright 2018-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. +""" +This module contains code for the main device construction delegation logic. +""" +import warnings +from importlib import metadata +from sys import version_info + +from semantic_version import SimpleSpec, Version + +import pennylane as qml + + +def _get_device_entrypoints(): + """Returns a dictionary mapping the device short name to the + loadable entrypoint""" + + entries = ( + metadata.entry_points()["pennylane.plugins"] + if version_info[:2] == (3, 9) + # pylint:disable=unexpected-keyword-arg + else metadata.entry_points(group="pennylane.plugins") + ) + return {entry.name: entry for entry in entries} + + +# get list of installed devices +plugin_devices = _get_device_entrypoints() + + +def refresh_devices(): + """Scan installed PennyLane plugins to refresh the device list.""" + + # This function does not return anything; instead, it has a side effect + # which is to update the global plugin_devices variable. + + # We wish to retain the behaviour of a global plugin_devices dictionary, + # as re-importing metadata can be a very slow operation on systems + # with a large number of installed packages. + + global plugin_devices # pylint:disable=global-statement + plugin_devices = _get_device_entrypoints() + + +# pylint: disable=protected-access +def device(name, *args, **kwargs): + r""" + Load a device and return the instance. + + This function is used to load a particular quantum device, + which can then be used to construct QNodes. + + PennyLane comes with support for the following devices: + + * :mod:`'default.qubit' `: a simple + state simulator of qubit-based quantum circuit architectures. + + * :mod:`'default.mixed' `: a mixed-state + simulator of qubit-based quantum circuit architectures. + + * ``'lightning.qubit'``: a more performant state simulator of qubit-based + quantum circuit architectures written in C++. + + * :mod:`'default.qutrit' `: a simple + state simulator of qutrit-based quantum circuit architectures. + + * :mod:`'default.qutrit.mixed' `: a + mixed-state simulator of qutrit-based quantum circuit architectures. + + * :mod:`'default.gaussian' `: a simple simulator + of Gaussian states and operations on continuous-variable circuit architectures. + + * :mod:`'default.clifford' `: an efficient + simulator of Clifford circuits. + + * :mod:`'default.tensor' `: a simulator + of quantum circuits based on tensor networks. + + Additional devices are supported through plugins β€” see + the `available plugins `_ for more + details. To list all currently installed devices, run + :func:`qml.about `. + + Args: + name (str): the name of the device to load + wires (int): the number of wires (subsystems) to initialise + the device with. Note that this is optional for certain + devices, such as ``default.qubit`` + + Keyword Args: + config (pennylane.Configuration): a PennyLane configuration object + that contains global and/or device specific configurations. + custom_decomps (Dict[Union(str, Operator), Callable]): Custom + decompositions to be applied by the device at runtime. + decomp_depth (int): For when custom decompositions are specified, + the maximum expansion depth used by the expansion function. + + .. warning:: + + The ``decomp_depth`` argument is deprecated and will be removed in version 0.39. + + All devices must be loaded by specifying their **short-name** as listed above, + followed by the **wires** (subsystems) you wish to initialize. The ``wires`` + argument can be an integer, in which case the wires of the device are addressed + by consecutive integers: + + .. code-block:: python + + dev = qml.device('default.qubit', wires=5) + + def circuit(): + qml.Hadamard(wires=1) + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[3, 4]) + ... + + The ``wires`` argument can also be a sequence of unique numbers or strings, specifying custom wire labels + that the user employs to address the wires: + + .. code-block:: python + + dev = qml.device('default.qubit', wires=['ancilla', 'q11', 'q12', -1, 1]) + + def circuit(): + qml.Hadamard(wires='q11') + qml.Hadamard(wires=['ancilla']) + qml.CNOT(wires=['q12', -1]) + ... + + On some newer devices, such as ``default.qubit``, the ``wires`` argument can be omitted altogether, + and instead the wires will be computed when executing a circuit depending on its contents. + + >>> dev = qml.device("default.qubit") + + Most devices accept a ``shots`` argument which specifies how many circuit executions + are used to estimate stochastic return values. As an example, ``qml.sample()`` measurements + will return as many samples as specified in the shots argument. The shots argument can be + changed on a per-call basis using the built-in ``shots`` keyword argument. Note that the + ``shots`` argument can be a single integer or a list of shot values. + + .. code-block:: python + + dev = qml.device('default.qubit', wires=1, shots=10) + + @qml.qnode(dev) + def circuit(a): + qml.RX(a, wires=0) + return qml.sample(qml.Z(0)) + + >>> circuit(0.8) # 10 samples are returned + array([ 1, 1, 1, 1, -1, 1, 1, -1, 1, 1]) + >>> circuit(0.8, shots=[3, 4, 4]) # default is overwritten for this call + (array([1, 1, 1]), array([ 1, -1, 1, 1]), array([1, 1, 1, 1])) + >>> circuit(0.8) # back to default of 10 samples + array([ 1, -1, 1, 1, -1, 1, 1, 1, 1, 1]) + + When constructing a device, we may optionally pass a dictionary of custom + decompositions to be applied to certain operations upon device execution. + This is useful for enabling support of gates on devices where they would normally + be unsupported. + + For example, suppose we are running on an ion trap device which does not + natively implement the CNOT gate, but we would still like to write our + circuits in terms of CNOTs. On a ion trap device, CNOT can be implemented + using the ``IsingXX`` gate. We first define a decomposition function + (such functions have the signature ``decomposition(*params, wires)``): + + .. code-block:: python + + def ion_trap_cnot(wires, **_): + return [ + qml.RY(np.pi/2, wires=wires[0]), + qml.IsingXX(np.pi/2, wires=wires), + qml.RX(-np.pi/2, wires=wires[0]), + qml.RY(-np.pi/2, wires=wires[0]), + qml.RY(-np.pi/2, wires=wires[1]) + ] + + Next, we create a device, and a QNode for testing. When constructing the + QNode, we can set the expansion strategy to ``"device"`` to ensure the + decomposition is applied and will be viewable when we draw the circuit. + Note that custom decompositions should accept keyword arguments even when + it is not used. + + .. code-block:: python + + # As the CNOT gate normally has no decomposition, we can use default.qubit + # here for expository purposes. + dev = qml.device( + 'default.qubit', wires=2, custom_decomps={"CNOT" : ion_trap_cnot} + ) + + @qml.qnode(dev, expansion_strategy="device") + def run_cnot(): + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.X(1)) + + >>> print(qml.draw(run_cnot)()) + 0: ──RY(1.57)─╭IsingXX(1.57)──RX(-1.57)──RY(-1.57)── + 1: ───────────╰IsingXX(1.57)──RY(-1.57)───────────── + + Some devices may accept additional arguments. For instance, + ``default.gaussian`` accepts the keyword argument ``hbar``, to set + the convention used in the commutation relation :math:`[\x,\p]=i\hbar` + (by default set to 2). + + Please refer to the documentation for the individual devices to see any + additional arguments that might be required or supported. + """ + if name not in plugin_devices: + # Device does not exist in the loaded device list. + # Attempt to refresh the devices, in case the user + # installed the plugin during the current Python session. + refresh_devices() + + if name in plugin_devices: + options = {} + + # load global configuration settings if available + config = kwargs.get("config", qml.default_config) + + if config: + # combine configuration options with keyword arguments. + # Keyword arguments take preference, followed by device options, + # followed by plugin options, followed by global options. + options.update(config["main"]) + options.update(config[name.split(".")[0] + ".global"]) + options.update(config[name]) + + # Pop the custom decomposition keyword argument; we will use it here + # only and not pass it to the device. + custom_decomps = kwargs.pop("custom_decomps", None) + decomp_depth = kwargs.pop("decomp_depth", None) + + if decomp_depth is not None: + warnings.warn( + "The decomp_depth argument is deprecated and will be removed in version 0.39. ", + qml.PennyLaneDeprecationWarning, + ) + else: + decomp_depth = 10 + + kwargs.pop("config", None) + options.update(kwargs) + + # loads the device class + plugin_device_class = plugin_devices[name].load() + + if hasattr(plugin_device_class, "pennylane_requires") and Version( + qml.version() + ) not in SimpleSpec(plugin_device_class.pennylane_requires): + raise qml.DeviceError( + f"The {name} plugin requires PennyLane versions {plugin_device_class.pennylane_requires}, " + f"however PennyLane version {qml.version()} is installed." + ) + + # Construct the device + dev = plugin_device_class(*args, **options) + + # Once the device is constructed, we set its custom expansion function if + # any custom decompositions were specified. + + if custom_decomps is not None: + if isinstance(dev, qml.devices.LegacyDevice): + custom_decomp_expand_fn = qml.transforms.create_decomp_expand_fn( + custom_decomps, dev, decomp_depth=decomp_depth + ) + dev.custom_expand(custom_decomp_expand_fn) + else: + custom_decomp_preprocess = qml.transforms.tape_expand._create_decomp_preprocessing( + custom_decomps, dev, decomp_depth=decomp_depth + ) + dev.preprocess = custom_decomp_preprocess + + return dev + + raise qml.DeviceError( + f"Device {name} does not exist. Make sure the required plugin is installed." + ) diff --git a/pennylane/devices/qubit/simulate.py b/pennylane/devices/qubit/simulate.py index e55741b029d..9017e300316 100644 --- a/pennylane/devices/qubit/simulate.py +++ b/pennylane/devices/qubit/simulate.py @@ -230,7 +230,6 @@ def measure_final_state(circuit, state, is_state_batched, **execution_kwargs) -> mid_measurements = execution_kwargs.get("mid_measurements", None) # analytic case - if not circuit.shots: if mid_measurements is not None: raise TypeError("Native mid-circuit measurements are only supported with finite shots.") @@ -243,7 +242,6 @@ def measure_final_state(circuit, state, is_state_batched, **execution_kwargs) -> ) # finite-shot case - rng = default_rng(rng) results = measure_with_samples( circuit.measurements, diff --git a/pennylane/drawer/tape_mpl.py b/pennylane/drawer/tape_mpl.py index d73e4bcc68f..2626cde639d 100644 --- a/pennylane/drawer/tape_mpl.py +++ b/pennylane/drawer/tape_mpl.py @@ -459,6 +459,7 @@ def tape_mpl( if update_style := (has_mpl and style != "rcParams"): restore_params = mpl.rcParams.copy() _set_style(style) + try: return _tape_mpl( tape, diff --git a/pennylane/gradients/pulse_gradient.py b/pennylane/gradients/pulse_gradient.py index fe42c0eab55..b3b154608c3 100644 --- a/pennylane/gradients/pulse_gradient.py +++ b/pennylane/gradients/pulse_gradient.py @@ -393,7 +393,7 @@ def stoch_pulse_grad( jax.config.update("jax_enable_x64", True) - dev = qml.device("default.qubit.jax") + dev = qml.device("default.qubit") def sin(p, t): return jax.numpy.sin(p * t) diff --git a/pennylane/gradients/pulse_gradient_odegen.py b/pennylane/gradients/pulse_gradient_odegen.py index cc1bd678da6..d01baaadd80 100644 --- a/pennylane/gradients/pulse_gradient_odegen.py +++ b/pennylane/gradients/pulse_gradient_odegen.py @@ -485,7 +485,7 @@ def pulse_odegen( .. code-block:: python - dev = qml.device("default.qubit.jax") + dev = qml.device("default.qubit") @qml.qnode(dev, interface="jax", diff_method=qml.gradients.pulse_odegen) def circuit(params): diff --git a/pennylane/ops/functions/evolve.py b/pennylane/ops/functions/evolve.py index c7d2722b9b2..96660aea77e 100644 --- a/pennylane/ops/functions/evolve.py +++ b/pennylane/ops/functions/evolve.py @@ -129,7 +129,10 @@ def evolve(*args, **kwargs): # pylint: disable=unused-argument import jax - dev = qml.device("default.qubit.jax", wires=4) + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit") + @jax.jit @qml.qnode(dev, interface="jax") def circuit(params): @@ -138,13 +141,13 @@ def circuit(params): >>> params = [1., 2., 3., 4.] >>> circuit(params) - Array(0.8627419, dtype=float32) + Array(0.86231063, dtype=float64) >>> jax.grad(circuit)(params) - [Array(50.690746, dtype=float32), - Array(-6.296886e-05, dtype=float32), - Array(-6.3341584e-05, dtype=float32), - Array(-7.052516e-05, dtype=float32)] + [Array(50.391273, dtype=float64), + Array(-9.42415807e-05, dtype=float64), + Array(-0.0001049, dtype=float64), + Array(-0.00010601, dtype=float64)] .. note:: In the example above, the decorator ``@jax.jit`` is used to compile this execution just-in-time. This means diff --git a/pennylane/ops/qubit/observables.py b/pennylane/ops/qubit/observables.py index b20e96dd478..9fd4ee43e65 100644 --- a/pennylane/ops/qubit/observables.py +++ b/pennylane/ops/qubit/observables.py @@ -584,6 +584,24 @@ def compute_diagonalizing_gates( """ return [] + @staticmethod + def compute_sparse_matrix(basis_state): # pylint: disable=arguments-differ,unused-argument + """ + Computes the sparse CSR matrix representation of the projector onto the basis state. + + Args: + basis_state (Iterable): The basis state as an iterable of integers (0 or 1). + + Returns: + scipy.sparse.csr_matrix: The sparse CSR matrix representation of the projector. + """ + + num_qubits = len(basis_state) + data = [1] + rows = [int("".join(str(bit) for bit in basis_state), 2)] + cols = rows + return csr_matrix((data, (rows, cols)), shape=(2**num_qubits, 2**num_qubits)) + class StateVectorProjector(Projector): r"""Observable corresponding to the state projector :math:`P=\ket{\phi}\bra{\phi}`, where diff --git a/pennylane/pulse/__init__.py b/pennylane/pulse/__init__.py index a5396032ccd..8bbafb5af6d 100644 --- a/pennylane/pulse/__init__.py +++ b/pennylane/pulse/__init__.py @@ -236,7 +236,9 @@ import jax - dev = qml.device("default.qubit.jax", wires=1) + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=1) @jax.jit @qml.qnode(dev, interface="jax") @@ -246,10 +248,10 @@ def circuit(params): >>> params = [1.2] >>> circuit(params) -Array(0.96632576, dtype=float32) +Array(0.96632722, dtype=float64) >>> jax.grad(circuit)(params) -[Array(2.3569832, dtype=float32)] +[Array(2.35694829, dtype=float64)] We can use the decorator ``jax.jit`` to compile this execution just-in-time. This means the first execution will typically take a little longer with the benefit that all following executions will be significantly faster. diff --git a/pennylane/pulse/convenience_functions.py b/pennylane/pulse/convenience_functions.py index 2dca820151f..dd4f71fabed 100644 --- a/pennylane/pulse/convenience_functions.py +++ b/pennylane/pulse/convenience_functions.py @@ -26,7 +26,7 @@ # pylint: disable=unused-argument def constant(scalar, time): - """Returns the given ``scalar``, for use in defining a :class:`~.ParametrizedHamiltonian` with a + r"""Returns the given ``scalar``, for use in defining a :class:`~.ParametrizedHamiltonian` with a trainable coefficient. Args: @@ -58,7 +58,12 @@ def constant(scalar, time): .. code-block:: python - dev = qml.device("default.qubit.jax", wires=1) + import jax + + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit") + @qml.qnode(dev, interface="jax") def circuit(params): qml.evolve(H)(params, t=2) @@ -67,10 +72,10 @@ def circuit(params): >>> params = jnp.array([5.0]) >>> circuit(params) - Array(0.40808904, dtype=float32) + Array(0.40808193, dtype=float64) >>> jax.grad(circuit)(params) - Array([-3.6517754], dtype=float32) + Array([-3.65178003], dtype=float64) """ return scalar diff --git a/pennylane/pulse/hardware_hamiltonian.py b/pennylane/pulse/hardware_hamiltonian.py index 8d7f33666d7..8139d884a8b 100644 --- a/pennylane/pulse/hardware_hamiltonian.py +++ b/pennylane/pulse/hardware_hamiltonian.py @@ -96,7 +96,9 @@ def drive(amplitude, phase, wires): .. code-block:: python3 - dev = qml.device("default.qubit.jax", wires=wires) + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=wires) @qml.qnode(dev, interface="jax") def circuit(params): @@ -107,7 +109,7 @@ def circuit(params): >>> circuit(params) Array(0.32495208, dtype=float64) >>> jax.grad(circuit)(params) - [Array(1.31956098, dtype=float64)] + [Array(1.31956098, dtype=float64, weak_type=True)] We can also create a Hamiltonian with multiple local drives. The following circuit corresponds to the evolution where an additional local drive that changes in time is acting on wires ``[0, 1]`` is added to the Hamiltonian: @@ -132,12 +134,12 @@ def circuit_local(params): params = (p_global, p_amp, p_phase) >>> circuit_local(params) - Array(-0.5334795, dtype=float64) + Array(0.37385014, dtype=float64) >>> jax.grad(circuit_local)(params) - (Array(0.01654573, dtype=float64), - [Array(-0.04422795, dtype=float64, weak_type=True), - Array(-0.51375441, dtype=float64, weak_type=True)], - Array(0.21901967, dtype=float64)) + (Array(-3.35835837, dtype=float64), + [Array(-3.35835837, dtype=float64, weak_type=True), + Array(-3.35835837, dtype=float64, weak_type=True)], + Array(0.1339487, dtype=float64)) .. details:: :title: Theoretical background @@ -189,7 +191,7 @@ def circuit_local(params): H_d = qml.pulse.drive(amplitude, phase, wires) # detuning term - H_z = qml.dot([-3*np.pi/4]*len(wires), [qml.Z(i) for i in wires]) + H_z = qml.dot([-3*jnp.pi/4]*len(wires), [qml.Z(i) for i in wires]) The total Hamiltonian of that evolution is given by @@ -202,7 +204,7 @@ def circuit_local(params): .. code-block:: python3 - dev = qml.device("default.qubit.jax", wires=wires) + dev = qml.device("default.qubit", wires=wires) @qml.qnode(dev, interface="jax") def circuit(params): qml.evolve(H_i + H_z + H_d)(params, t=[0, 10]) @@ -210,9 +212,9 @@ def circuit(params): >>> params = [2.4] >>> circuit(params) - Array(0.6962041, dtype=float64) + Array(0.96347734, dtype=float64) >>> jax.grad(circuit)(params) - [Array(1.75825695, dtype=float64)] + [Array(-0.4311521, dtype=float64, weak_type=True)] """ wires = Wires(wires) diff --git a/pennylane/pulse/parametrized_evolution.py b/pennylane/pulse/parametrized_evolution.py index 5c806f4f374..fc335b65a80 100644 --- a/pennylane/pulse/parametrized_evolution.py +++ b/pennylane/pulse/parametrized_evolution.py @@ -138,7 +138,9 @@ class ParametrizedEvolution(Operation): import jax - dev = qml.device("default.qubit.jax", wires=1) + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=1) @jax.jit @qml.qnode(dev, interface="jax") def circuit(params): @@ -147,10 +149,10 @@ def circuit(params): >>> params = [1.2] >>> circuit(params) - Array(0.96632576, dtype=float32) + Array(0.96632722, dtype=float64) >>> jax.grad(circuit)(params) - [Array(2.3569832, dtype=float32)] + [Array(2.35694829, dtype=float64)] .. note:: In the example above, the decorator ``@jax.jit`` is used to compile this execution just-in-time. This means @@ -224,7 +226,7 @@ def f2(p, t): .. code-block:: python - dev = qml.device("default.qubit.jax", wires=3) + dev = qml.device("default.qubit", wires=3) @qml.qnode(dev, interface="jax") def circuit1(params): @@ -246,11 +248,11 @@ def circuit2(params): >>> params = jnp.array([1., 2., 3.]) >>> circuit1(params) - Array(-0.01543971, dtype=float32) + Array(-0.01542578, dtype=float64) >>> params = jnp.concatenate([params, params]) # H1 + H2 requires 6 parameters! >>> circuit2(params) - Array(-0.78236955, dtype=float32) + Array(-0.78235162, dtype=float64) Here, ``circuit1`` is not executing the evolution of ``H1`` and ``H2`` simultaneously, but rather executing ``H1`` in the ``[0, 10]`` time window and then executing ``H2`` with the same time window, @@ -269,10 +271,10 @@ def circuit(params): return qml.expval(qml.Z(0) @ qml.Z(1) @ qml.Z(2)) >>> circuit(params) - Array(-0.78236955, dtype=float32) + Array(-0.78235162, dtype=float64) >>> jax.grad(circuit)(params) - Array([-4.8066125 , 3.703827 , -1.3297377 , -2.406232 , 0.6811726 , - -0.52277344], dtype=float32) + Array([-4.80708632, 3.70323783, -1.32958799, -2.40642477, 0.68105214, + -0.52269657], dtype=float64) Given that we used the same time window (``[0, 10]``), the results are the same as before. @@ -312,7 +314,7 @@ def circuit(params): .. code-block:: python - dev = qml.device("default.qubit.jax", wires=1) + dev = qml.device("default.qubit", wires=1) @qml.qnode(dev, interface="jax") def circuit(param, time): @@ -321,11 +323,11 @@ def circuit(param, time): >>> circuit(param, time) Array([[1. , 0. ], - [0.9897738 , 0.01022595], - [0.9599043 , 0.04009585], - [0.9123617 , 0.08763832], - [0.84996957, 0.15003097], - [0.7761489 , 0.22385144]], dtype=float32) + [0.98977406, 0.01022594], + [0.95990416, 0.04009584], + [0.91236167, 0.08763833], + [0.84996865, 0.15003133], + [0.77614817, 0.22385181]], dtype=float64) **Computing complementary time evolution** diff --git a/pennylane/pulse/rydberg.py b/pennylane/pulse/rydberg.py index 761c61ffdc8..2059d3e9fd5 100644 --- a/pennylane/pulse/rydberg.py +++ b/pennylane/pulse/rydberg.py @@ -86,7 +86,11 @@ def rydberg_interaction( .. code-block:: python - dev = qml.device("default.qubit.jax", wires=9) + import jax + + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=9) @qml.qnode(dev, interface="jax") def circuit(): @@ -94,7 +98,7 @@ def circuit(): return qml.expval(qml.Z(0)) >>> circuit() - Array(1., dtype=float32) + Array(1., dtype=float64) """ if wires is None: wires = list(range(len(register))) @@ -208,7 +212,11 @@ def rydberg_drive(amplitude, phase, detuning, wires): .. code-block:: python - dev = qml.device("default.qubit.jax", wires=wires) + import jax + + jax.config.update("jax_enable_x64", True) + + dev = qml.device("default.qubit", wires=wires) @qml.qnode(dev, interface="jax") def circuit(params): qml.evolve(H_i + H_d)(params, t=[0, 0.5]) diff --git a/pennylane/pulse/transmon.py b/pennylane/pulse/transmon.py index 2728133a07f..7262d54f6ce 100644 --- a/pennylane/pulse/transmon.py +++ b/pennylane/pulse/transmon.py @@ -351,6 +351,10 @@ def phase(phi0, t): .. code-block:: python3 + import jax + + jax.config.update("jax_enable_x64", True) + qubit_freqs = [5.1, 5., 5.3] connections = [[0, 1], [1, 2]] # qubits 0 and 1 are coupled, as are 1 and 2 g = [0.02, 0.05] @@ -364,7 +368,7 @@ def amp(max_amp, t): return max_amp * jnp.sin(t) ** 2 for q in range(3): H += qml.pulse.transmon_drive(amp, phase, freq, q) # Parametrized drive for each qubit - dev = qml.device("default.qubit.jax", wires=range(3)) + dev = qml.device("default.qubit", wires=range(3)) @jax.jit @qml.qnode(dev, interface="jax") diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py index 5900f056ecc..271929ad8c8 100644 --- a/pennylane/tape/qscript.py +++ b/pennylane/tape/qscript.py @@ -538,7 +538,7 @@ def data(self): @property def trainable_params(self): - """Store or return a list containing the indices of parameters that support + r"""Store or return a list containing the indices of parameters that support differentiability. The indices provided match the order of appearence in the quantum circuit. diff --git a/pennylane/workflow/set_shots.py b/pennylane/workflow/set_shots.py index bb463233fed..1bc7dff7f33 100644 --- a/pennylane/workflow/set_shots.py +++ b/pennylane/workflow/set_shots.py @@ -24,15 +24,14 @@ @contextlib.contextmanager def set_shots(device, shots): - """Context manager to temporarily change the shots - of a device. + r"""Context manager to temporarily change the shots of a device. This context manager can be used in two ways. As a standard context manager: >>> dev = qml.device("default.qubit.legacy", wires=2, shots=None) - >>> with set_shots(dev, shots=100): + >>> with qml.workflow.set_shots(dev, shots=100): ... print(dev.shots) 100 >>> print(dev.shots) @@ -40,7 +39,7 @@ def set_shots(device, shots): Or as a decorator that acts on a function that uses the device: - >>> set_shots(dev, shots=100)(lambda: dev.shots)() + >>> qml.workflow.set_shots(dev, shots=100)(lambda: dev.shots)() 100 """ if isinstance(device, qml.devices.Device): diff --git a/tests/conftest.py b/tests/conftest.py index 4c8eba9dc6d..3094a2fd784 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -19,6 +19,7 @@ import os import pathlib import sys +import warnings import numpy as np import pytest @@ -49,6 +50,23 @@ def set_numpy_seed(): yield +@pytest.fixture(scope="function", autouse=True) +def capture_legacy_device_deprecation_warnings(): + with warnings.catch_warnings(record=True) as recwarn: + warnings.simplefilter("always") + yield + + for w in recwarn: + if isinstance(w, qml.PennyLaneDeprecationWarning): + assert "Use of 'default.qubit." in str(w.message) + assert "is deprecated" in str(w.message) + assert "use 'default.qubit'" in str(w.message) + + for w in recwarn: + if "Use of 'default.qubit." not in str(w.message): + warnings.warn(message=w.message, category=w.category) + + @pytest.fixture(scope="session") def tol(): """Numerical tolerance for equality tests.""" @@ -81,7 +99,8 @@ def n_subsystems_fixture(request): @pytest.fixture(scope="session") def qubit_device(n_subsystems): - return qml.device("default.qubit.legacy", wires=n_subsystems) + with pytest.warns(qml.PennyLaneDeprecationWarning, match="Use of 'default.qubit.legacy'"): + return qml.device("default.qubit.legacy", wires=n_subsystems) @pytest.fixture(scope="function", params=[(np.float32, np.complex64), (np.float64, np.complex128)]) diff --git a/tests/devices/default_tensor/test_default_tensor.py b/tests/devices/default_tensor/test_default_tensor.py index 9fb622e3174..203680053d4 100644 --- a/tests/devices/default_tensor/test_default_tensor.py +++ b/tests/devices/default_tensor/test_default_tensor.py @@ -405,7 +405,7 @@ def test_jax(self, method): jax = pytest.importorskip("jax") dev = qml.device("default.tensor", wires=1, method=method) - ref_dev = qml.device("default.qubit.jax", wires=1) + ref_dev = qml.device("default.qubit", wires=1) def circuit(x): qml.RX(x[1], wires=0) diff --git a/tests/devices/test_default_qubit_legacy.py b/tests/devices/test_default_qubit_legacy.py index ec9fd6e65f5..be8e586adcc 100644 --- a/tests/devices/test_default_qubit_legacy.py +++ b/tests/devices/test_default_qubit_legacy.py @@ -2001,7 +2001,8 @@ class TestApplyOps: gates in DefaultQubitLegacy.""" state = np.arange(2**4, dtype=np.complex128).reshape((2, 2, 2, 2)) - dev = qml.device("default.qubit.legacy", wires=4) + with pytest.warns(qml.PennyLaneDeprecationWarning): + dev = qml.device("default.qubit.legacy", wires=4) single_qubit_ops = [ (qml.PauliX, dev._apply_x), diff --git a/tests/devices/test_default_qubit_legacy_broadcasting.py b/tests/devices/test_default_qubit_legacy_broadcasting.py index ebf21502971..14d741d0423 100644 --- a/tests/devices/test_default_qubit_legacy_broadcasting.py +++ b/tests/devices/test_default_qubit_legacy_broadcasting.py @@ -1619,7 +1619,8 @@ class TestApplyOpsBroadcasted: gates in DefaultQubitLegacy.""" broadcasted_state = np.arange(2**4 * 3, dtype=np.complex128).reshape((3, 2, 2, 2, 2)) - dev = qml.device("default.qubit.legacy", wires=4) + with pytest.warns(qml.PennyLaneDeprecationWarning): + dev = qml.device("default.qubit.legacy", wires=4) single_qubit_ops = [ (qml.PauliX, dev._apply_x), diff --git a/tests/test_device.py b/tests/devices/test_device.py similarity index 98% rename from tests/test_device.py rename to tests/devices/test_device.py index 2c15a799227..a5f6631eb90 100644 --- a/tests/test_device.py +++ b/tests/devices/test_device.py @@ -188,6 +188,12 @@ def get_device(wires=1): yield get_device +# pylint: disable=pointless-statement +def test_invalid_attribute_in_devices_raises_error(): + with pytest.raises(AttributeError, match="'pennylane.devices' has no attribute 'blabla'"): + qml.devices.blabla + + def test_gradients_record(): """Test that execute_and_gradients and gradient both track the number of gradients requested.""" @@ -642,7 +648,7 @@ def test_default_expand_with_initial_state(self, op, decomp): prep = [op] ops = [qml.AngleEmbedding(features=[0.1], wires=[0], rotation="Z"), op, qml.PauliZ(wires=2)] - dev = qml.device("default.qubit.legacy", wires=3) + dev = qml.device("default.mixed", wires=3) tape = qml.tape.QuantumTape(ops=prep + ops, measurements=[], shots=100) new_tape = dev.default_expand_fn(tape) @@ -998,7 +1004,14 @@ def test_outdated_API(self, monkeypatch): with monkeypatch.context() as m: m.setattr(qml, "version", lambda: "0.0.1") with pytest.raises(DeviceError, match="plugin requires PennyLane versions"): - qml.device("default.qubit.legacy", wires=0) + qml.device("default.mixed", wires=0) + + def test_plugin_devices_from_devices_triggers_getattr(self, mocker): + spied = mocker.spy(qml.devices, "__getattr__") + + qml.devices.plugin_devices + + spied.assert_called_once() def test_refresh_entrypoints(self, monkeypatch): """Test that new entrypoints are found by the refresh_devices function""" @@ -1011,6 +1024,7 @@ def test_refresh_entrypoints(self, monkeypatch): # reimporting PennyLane within the context sets qml.plugin_devices to {} reload(qml) + reload(qml.devices.device_constructor) # since there are no entry points, there will be no plugin devices assert not qml.plugin_devices @@ -1023,6 +1037,7 @@ def test_refresh_entrypoints(self, monkeypatch): # Test teardown: re-import PennyLane to revert all changes and # restore the plugin_device dictionary reload(qml) + reload(qml.devices.device_constructor) def test_hot_refresh_entrypoints(self, monkeypatch): """Test that new entrypoints are found by the device loader if not currently present""" @@ -1034,9 +1049,10 @@ def test_hot_refresh_entrypoints(self, monkeypatch): m.setattr(metadata, "entry_points", lambda **kwargs: retval) # reimporting PennyLane within the context sets qml.plugin_devices to {} - reload(qml) + reload(qml.devices) + reload(qml.devices.device_constructor) - m.setattr(qml, "refresh_devices", lambda: None) + m.setattr(qml.devices.device_constructor, "refresh_devices", lambda: None) assert not qml.plugin_devices # since there are no entry points, there will be no plugin devices @@ -1052,10 +1068,11 @@ def test_hot_refresh_entrypoints(self, monkeypatch): # Test teardown: re-import PennyLane to revert all changes and # restore the plugin_device dictionary reload(qml) + reload(qml.devices.device_constructor) def test_shot_vector_property(self): """Tests shot vector initialization.""" - dev = qml.device("default.qubit.legacy", wires=1, shots=[1, 3, 3, 4, 4, 4, 3]) + dev = qml.device("default.mixed", wires=1, shots=[1, 3, 3, 4, 4, 4, 3]) shot_vector = dev.shot_vector assert len(shot_vector) == 4 assert shot_vector[0].shots == 1 diff --git a/tests/drawer/test_draw.py b/tests/drawer/test_draw.py index c1ae86acf15..05d5efeb437 100644 --- a/tests/drawer/test_draw.py +++ b/tests/drawer/test_draw.py @@ -878,13 +878,11 @@ def circ(): class TestLevelExpansionStrategy: - @pytest.fixture( - params=[qml.device("default.qubit.legacy", wires=3), qml.devices.DefaultQubit()], - ) - def transforms_circuit(self, request): + @pytest.fixture + def transforms_circuit(self): @qml.transforms.merge_rotations @qml.transforms.cancel_inverses - @qml.qnode(request.param, diff_method="parameter-shift") + @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift") def circ(weights, order): qml.RandomLayers(weights, wires=(0, 1)) qml.Permute(order, wires=(0, 1, 2)) @@ -1024,11 +1022,7 @@ def circ(): assert draw(circ)() == expected -@pytest.mark.parametrize( - "device", - [qml.device("default.qubit.legacy", wires=2), qml.device("default.qubit", wires=2)], -) -def test_applied_transforms(device): +def test_applied_transforms(): """Test that any transforms applied to the qnode are included in the output.""" @qml.transform @@ -1037,7 +1031,7 @@ def just_pauli_x(_): return (new_tape,), lambda res: res[0] @just_pauli_x - @qml.qnode(device) + @qml.qnode(qml.device("default.qubit", wires=2)) def my_circuit(x): qml.RX(x, wires=0) qml.SWAP(wires=(0, 1)) diff --git a/tests/drawer/test_draw_mpl.py b/tests/drawer/test_draw_mpl.py index 91afca5f297..c5ed45534d0 100644 --- a/tests/drawer/test_draw_mpl.py +++ b/tests/drawer/test_draw_mpl.py @@ -81,13 +81,11 @@ def test_fig_argument(): class TestLevelExpansionStrategy: - @pytest.fixture( - params=[qml.device("default.qubit.legacy", wires=3), qml.devices.DefaultQubit()], - ) - def transforms_circuit(self, request): + @pytest.fixture + def transforms_circuit(self): @qml.transforms.merge_rotations @qml.transforms.cancel_inverses - @qml.qnode(request.param, diff_method="parameter-shift") + @qml.qnode(qml.device("default.qubit"), diff_method="parameter-shift") def circ(weights, order): qml.RandomLayers(weights, wires=(0, 1)) qml.Permute(order, wires=(0, 1, 2)) @@ -125,15 +123,11 @@ def test_equivalent_levels(self, transforms_circuit, levels, expected_metadata): plt.close("all") - @pytest.mark.parametrize( - "device", - [qml.device("default.qubit.legacy", wires=3), qml.devices.DefaultQubit(wires=3)], - ) @pytest.mark.parametrize( "strategy, initial_strategy, n_lines", [("gradient", "device", 3), ("device", "gradient", 13)], ) - def test_expansion_strategy(self, device, strategy, initial_strategy, n_lines): + def test_expansion_strategy(self, strategy, initial_strategy, n_lines): """Test that the expansion strategy keyword controls what operations are drawn.""" with pytest.warns( @@ -141,7 +135,7 @@ def test_expansion_strategy(self, device, strategy, initial_strategy, n_lines): match="'expansion_strategy' attribute is deprecated", ): - @qml.qnode(device, expansion_strategy=initial_strategy) + @qml.qnode(qml.device("default.qubit"), expansion_strategy=initial_strategy) def circuit(): qml.Permute([2, 0, 1], wires=(0, 1, 2)) return qml.expval(qml.PauliZ(0)) @@ -516,11 +510,7 @@ def circuit(): assert ax.texts[-1].get_text() == "X†" -@pytest.mark.parametrize( - "device", - [qml.device("default.qubit.legacy", wires=2), qml.device("default.qubit", wires=2)], -) -def test_applied_transforms(device): +def test_applied_transforms(): """Test that any transforms applied to the qnode are included in the output.""" @qml.transform @@ -529,7 +519,7 @@ def just_pauli_x(_): return (new_tape,), lambda res: res[0] @just_pauli_x - @qml.qnode(device) + @qml.qnode(qml.device("default.qubit", wires=2)) def my_circuit(): qml.SWAP(wires=(0, 1)) qml.CNOT(wires=(0, 1)) diff --git a/tests/drawer/test_tape_mpl.py b/tests/drawer/test_tape_mpl.py index c8568c6293f..60d1b997004 100644 --- a/tests/drawer/test_tape_mpl.py +++ b/tests/drawer/test_tape_mpl.py @@ -844,7 +844,7 @@ def test_single_measure_multiple_conds(self): qml.cond(m0, qml.PauliY)(0) tape = qml.tape.QuantumScript.from_queue(q) - _, ax = qml.drawer.tape_mpl(tape) + _, ax = qml.drawer.tape_mpl(tape, style="black_white") assert len(ax.patches) == 5 # three for measure, two for boxes @@ -857,11 +857,11 @@ def test_single_measure_multiple_conds(self): # probably not a good way to test this, but the best I can figure out assert pe1._gc == { - "linewidth": 5 * plt.rcParams["lines.linewidth"], + "linewidth": 5 * 1.5, # hardcoded value to black_white linewidth "foreground": "black", # lines.color for black white style } assert pe2._gc == { - "linewidth": 3 * plt.rcParams["lines.linewidth"], + "linewidth": 3 * 1.5, # hardcoded value to black_white linewidth "foreground": "white", # figure.facecolor for black white sytle } plt.close() @@ -875,7 +875,7 @@ def test_combo_measurement(self): qml.cond(m0 & m1, qml.PauliY)(0) tape = qml.tape.QuantumScript.from_queue(q) - _, ax = qml.drawer.tape_mpl(tape) + _, ax = qml.drawer.tape_mpl(tape, style="black_white") assert len(ax.patches) == 7 # three for 2 measurements, one for box [_, _, cwire1, cwire2, eraser] = ax.lines @@ -891,18 +891,18 @@ def test_combo_measurement(self): # probably not a good way to test this, but the best I can figure out assert pe1._gc == { - "linewidth": 5 * plt.rcParams["lines.linewidth"], + "linewidth": 5 * 1.5, # hardcoded value to black_white linewidth "foreground": "black", # lines.color for black white style } assert pe2._gc == { - "linewidth": 3 * plt.rcParams["lines.linewidth"], + "linewidth": 3 * 1.5, # hardcoded value to black_white linewidth "foreground": "white", # figure.facecolor for black white sytle } assert eraser.get_xdata() == (1.8, 2) assert eraser.get_ydata() == (2, 2) - assert eraser.get_color() == plt.rcParams["figure.facecolor"] - assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"] + assert eraser.get_color() == "white" # hardcoded value to black_white color + assert eraser.get_linewidth() == 3 * 1.5 # hardcoded value to black_white linewidth plt.close() @@ -918,7 +918,7 @@ def test_combo_measurement_non_terminal(self): qml.cond(m1, qml.T)(1) tape = qml.tape.QuantumScript.from_queue(q) - _, ax = qml.drawer.tape_mpl(tape) + _, ax = qml.drawer.tape_mpl(tape, style="black_white") [_, _, cwire1, cwire2, eraser] = ax.lines @@ -933,18 +933,18 @@ def test_combo_measurement_non_terminal(self): # probably not a good way to test this, but the best I can figure out assert pe1._gc == { - "linewidth": 5 * plt.rcParams["lines.linewidth"], + "linewidth": 5 * 1.5, # hardcoded value to black_white linewidth "foreground": "black", # lines.color for black white style } assert pe2._gc == { - "linewidth": 3 * plt.rcParams["lines.linewidth"], + "linewidth": 3 * 1.5, # hardcoded value to black_white linewidth "foreground": "white", # figure.facecolor for black white sytle } assert eraser.get_xdata() == (1.8, 2.2) assert eraser.get_ydata() == (2, 2) - assert eraser.get_color() == plt.rcParams["figure.facecolor"] - assert eraser.get_linewidth() == 3 * plt.rcParams["lines.linewidth"] + assert eraser.get_color() == "white" # hardcoded value to black_white color + assert eraser.get_linewidth() == 3 * 1.5 # hardcoded value to black_white linewidth plt.close() @@ -954,7 +954,7 @@ def test_single_mcm_measure(self): with qml.queuing.AnnotatedQueue() as q: m0 = qml.measure(0) qml.expval(m0) - _, ax = tape_mpl(qml.tape.QuantumScript.from_queue(q)) + _, ax = tape_mpl(qml.tape.QuantumScript.from_queue(q), style="black_white") assert len(ax.patches) == 6 # two measurement boxes assert ax.patches[3].get_x() == 1 - 0.75 / 2 + 0.2 # 1 - box_length/2 + pad @@ -974,11 +974,11 @@ def test_single_mcm_measure(self): # probably not a good way to test this, but the best I can figure out assert pe1._gc == { - "linewidth": 5 * plt.rcParams["lines.linewidth"], + "linewidth": 5 * 1.5, # hardcoded value to black_white linewidth "foreground": "black", # lines.color for black white style } assert pe2._gc == { - "linewidth": 3 * plt.rcParams["lines.linewidth"], + "linewidth": 3 * 1.5, # hardcoded value to black_white linewidth "foreground": "white", # figure.facecolor for black white sytle } diff --git a/tests/gpu/test_gpu_torch.py b/tests/gpu/test_gpu_torch.py index cd461140e17..8c76734e683 100644 --- a/tests/gpu/test_gpu_torch.py +++ b/tests/gpu/test_gpu_torch.py @@ -33,7 +33,7 @@ class TestTorchDevice: def test_device_to_cuda(self): """Checks device executes with cuda is input data is cuda""" - dev = qml.device("default.qubit.torch", wires=1) + dev = qml.device("default.qubit", wires=1) x = torch.tensor(0.1, requires_grad=True, device=torch.device("cuda")) @@ -53,7 +53,7 @@ def test_device_to_cuda(self): def test_mixed_devices(self): """Asserts works with both cuda and cpu input data""" - dev = qml.device("default.qubit.torch", wires=1) + dev = qml.device("default.qubit", wires=1) x = torch.tensor(0.1, requires_grad=True, device=torch.device("cuda")) y = torch.tensor(0.2, requires_grad=True, device=torch.device("cpu")) @@ -77,7 +77,7 @@ def test_mixed_devices(self): def test_matrix_input(self): """Test goes to GPU for matrix valued inputs.""" - dev = qml.device("default.qubit.torch", wires=1) + dev = qml.device("default.qubit", wires=1) U = torch.eye(2, requires_grad=False, device=torch.device("cuda")) @@ -93,7 +93,7 @@ def test_matrix_input(self): def test_resets(self): """Asserts reverts to cpu after execution on gpu""" - dev = qml.device("default.qubit.torch", wires=1) + dev = qml.device("default.qubit", wires=1) x = torch.tensor(0.1, requires_grad=True, device=torch.device("cuda")) y = torch.tensor(0.2, requires_grad=True, device=torch.device("cpu")) @@ -143,7 +143,7 @@ def test_different_devices_creation_and_parameters_warn(self, init_device, par_d PennyLane device creation differs from the Torch device of gate parameters. """ - dev = qml.device("default.qubit.torch", wires=1, torch_device=init_device) + dev = qml.device("default.qubit", wires=1, torch_device=init_device) p = torch.tensor(0.543, dtype=torch.float64, device=par_device) @@ -180,7 +180,7 @@ def circuit_cuda(inputs): @pytest.mark.skipif(not torch_cuda.is_available(), reason="no cuda support") -class TestqnnTorchLayer: +class TestQnnTorchLayer: def test_torch_device_cuda_if_tensors_on_cuda(self): """Test that if any tensor passed to operators is on the GPU then CUDA is set internally as a device option for 'default.qubit.torch'.""" diff --git a/tests/gradients/core/test_hadamard_gradient.py b/tests/gradients/core/test_hadamard_gradient.py index a455824e1df..682f685b99f 100644 --- a/tests/gradients/core/test_hadamard_gradient.py +++ b/tests/gradients/core/test_hadamard_gradient.py @@ -875,61 +875,21 @@ def circuit(weights): with pytest.raises(qml.QuantumFunctionError, match="No trainable parameters."): qml.gradients.hadamard_grad(circuit)(weights) - @pytest.mark.autograd - def test_no_trainable_params_qnode_autograd_legacy_opmath(self): + @pytest.mark.parametrize( + "interface", + [ + pytest.param("jax", marks=pytest.mark.jax), + pytest.param("autograd", marks=pytest.mark.autograd), + pytest.param("torch", marks=pytest.mark.torch), + pytest.param("tf", marks=pytest.mark.tf), + ], + ) + def test_no_trainable_params_qnode_legacy_opmath(self, interface): """Test that the correct ouput and warning is generated in the absence of any trainable parameters""" - dev = qml.device("default.qubit.autograd", wires=2) - - @qml.qnode(dev, interface="autograd") - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - weights = [0.1, 0.2] - with pytest.raises(qml.QuantumFunctionError, match="No trainable parameters."): - qml.gradients.hadamard_grad(circuit)(weights) + dev = qml.device(f"default.qubit.{interface}", wires=2) - @pytest.mark.torch - def test_no_trainable_params_qnode_torch_legacy_opmath(self): - """Test that the correct ouput and warning is generated in the absence of any trainable - parameters""" - dev = qml.device("default.qubit.torch", wires=2) - - @qml.qnode(dev, interface="torch") - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - weights = [0.1, 0.2] - with pytest.raises(qml.QuantumFunctionError, match="No trainable parameters."): - qml.gradients.hadamard_grad(circuit)(weights) - - @pytest.mark.tf - def test_no_trainable_params_qnode_tf_legacy_opmath(self): - """Test that the correct ouput and warning is generated in the absence of any trainable - parameters""" - dev = qml.device("default.qubit.tf", wires=2) - - @qml.qnode(dev, interface="tf") - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - weights = [0.1, 0.2] - with pytest.raises(qml.QuantumFunctionError, match="No trainable parameters."): - qml.gradients.hadamard_grad(circuit)(weights) - - @pytest.mark.jax - def test_no_trainable_params_qnode_jax_legacy_opmath(self): - """Test that the correct ouput and warning is generated in the absence of any trainable - parameters""" - dev = qml.device("default.qubit.jax", wires=2) - - @qml.qnode(dev, interface="jax") + @qml.qnode(dev, interface=interface) def circuit(weights): qml.RX(weights[0], wires=0) qml.RY(weights[1], wires=0) diff --git a/tests/gradients/core/test_metric_tensor.py b/tests/gradients/core/test_metric_tensor.py index 01b43bc8177..364516d6079 100644 --- a/tests/gradients/core/test_metric_tensor.py +++ b/tests/gradients/core/test_metric_tensor.py @@ -14,7 +14,9 @@ """ Unit tests for the metric tensor transform. """ -# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods,not-callable +import importlib + +# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods,not-callable,too-many-statements import pytest from scipy.linalg import block_diag @@ -538,189 +540,31 @@ def layer3_diag(params): G_expected = block_diag(G1, G3, G2) assert qml.math.allclose(G, G_expected, atol=tol, rtol=0) - @pytest.mark.autograd - @pytest.mark.parametrize("interface", ["autograd"]) - def test_argnum_metric_tensor_autograd(self, tol, interface): - """Test that argnum successfully reduces the number of tapes and gives - the desired outcome.""" - dev = qml.device("default.qubit.autograd", wires=3) - - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - qml.CNOT(wires=[0, 1]) - qml.RZ(weights[2], wires=1) - qml.RZ(weights[3], wires=0) - - weights = np.array([0.1, 0.2, 0.3, 0.5], requires_grad=True) - - with qml.tape.QuantumTape() as tape: - circuit(weights) - - tapes, proc_fn = qml.metric_tensor(tape) - res = qml.execute(tapes, dev, None) - mt = proc_fn(res) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt013 = proc_fn(res) - assert isinstance(mt013, np.ndarray) - - assert len(tapes) == 6 - assert mt.shape == mt013.shape - assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0) - assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt23 = proc_fn(res) - assert isinstance(mt23, np.ndarray) - - assert len(tapes) == 1 - assert mt.shape == mt23.shape - assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=0) - res = qml.execute(tapes, dev, None, interface=interface) - mt0 = proc_fn(res) - assert isinstance(mt0, np.ndarray) - - assert len(tapes) == 1 - assert mt.shape == mt0.shape - assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0) - - @pytest.mark.tf - @pytest.mark.parametrize("interface", ["tf"]) - def test_argnum_metric_tensor_tf(self, tol, interface): - """Test that argnum successfully reduces the number of tapes and gives - the desired outcome.""" - import tensorflow as tf - - dev = qml.device("default.qubit.tf", wires=3) - - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - qml.CNOT(wires=[0, 1]) - qml.RZ(weights[2], wires=1) - qml.RZ(weights[3], wires=0) - - weights = tf.Variable([0.1, 0.2, 0.3, 0.5]) - - with qml.tape.QuantumTape() as tape: - circuit(weights) - - tapes, proc_fn = qml.metric_tensor(tape) - res = qml.execute(tapes, dev, None) - mt = proc_fn(res) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt013 = proc_fn(res) - assert isinstance(mt013, tf.Tensor) - - assert len(tapes) == 6 - assert mt.shape == mt013.shape - assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0) - assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt23 = proc_fn(res) - assert isinstance(mt23, tf.Tensor) - - assert len(tapes) == 1 - assert mt.shape == mt23.shape - assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=0) - res = qml.execute(tapes, dev, None, interface=interface) - mt0 = proc_fn(res) - assert isinstance(mt0, tf.Tensor) - - assert len(tapes) == 1 - assert mt.shape == mt0.shape - assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0) - - @pytest.mark.torch - @pytest.mark.parametrize("interface", ["torch"]) - def test_argnum_metric_tensor_torch(self, tol, interface): + @pytest.mark.parametrize( + "interface,array_cls", + [ + pytest.param("jax", "array", marks=pytest.mark.jax), + pytest.param("autograd", "array", marks=pytest.mark.autograd), + pytest.param("tf", "Variable", marks=pytest.mark.tf), + pytest.param("torch", "Tensor", marks=pytest.mark.torch), + ], + ) + def test_argnum_metric_tensor_interfaces(self, tol, interface, array_cls): """Test that argnum successfully reduces the number of tapes and gives the desired outcome.""" - import torch - - dev = qml.device("default.qubit.torch", wires=3) - - def circuit(weights): - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=0) - qml.CNOT(wires=[0, 1]) - qml.RZ(weights[2], wires=1) - qml.RZ(weights[3], wires=0) - - weights = torch.tensor([0.1, 0.2, 0.3, 0.5], requires_grad=True) - - with qml.tape.QuantumTape() as tape: - circuit(weights) - - tapes, proc_fn = qml.metric_tensor(tape) - res = qml.execute(tapes, dev, None) - mt = proc_fn(res) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt013 = proc_fn(res) - assert isinstance(mt013, torch.Tensor) - - assert len(tapes) == 6 - assert mt.shape == mt013.shape - assert qml.math.allclose(mt[:2, :2], mt013[:2, :2], atol=tol, rtol=0) - assert qml.math.allclose(mt[3, 3], mt013[3, 3], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt013[:, 2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3)) - res = qml.execute(tapes, dev, None, interface=interface) - mt23 = proc_fn(res) - assert isinstance(mt23, torch.Tensor) - - assert len(tapes) == 1 - assert mt.shape == mt23.shape - assert qml.math.allclose(mt[2:, 2:], mt23[2:, 2:], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:2, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt23[:, :2], atol=tol, rtol=0) - - tapes, proc_fn = qml.metric_tensor(tape, argnum=0) - res = qml.execute(tapes, dev, None, interface=interface) - mt0 = proc_fn(res) - assert isinstance(mt0, torch.Tensor) - - assert len(tapes) == 1 - assert mt.shape == mt0.shape - assert qml.math.allclose(mt[0, 0], mt0[0, 0], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[1:, :], atol=tol, rtol=0) - assert qml.math.allclose(0, mt0[:, 1:], atol=tol, rtol=0) + if interface == "tf": + interface_name = "tensorflow" + elif interface == "jax": + interface_name = "jax.numpy" + elif interface == "autograd": + interface_name = "numpy" + else: + interface_name = interface - @pytest.mark.jax - @pytest.mark.parametrize("interface", ["jax"]) - def test_argnum_metric_tensor_jax(self, tol, interface): - """Test that argnum successfully reduces the number of tapes and gives - the desired outcome.""" - import jax + mod = importlib.import_module(interface_name) + type_ = type(getattr(mod, array_cls)([])) if interface != "tf" else getattr(mod, "Tensor") - dev = qml.device("default.qubit.jax", wires=3) + dev = qml.device("default.qubit", wires=3) def circuit(weights): qml.RX(weights[0], wires=0) @@ -729,7 +573,7 @@ def circuit(weights): qml.RZ(weights[2], wires=1) qml.RZ(weights[3], wires=0) - weights = jax.numpy.array([0.1, 0.2, 0.3, 0.5]) + weights = getattr(mod, array_cls)([0.1, 0.2, 0.3, 0.5]) with qml.tape.QuantumTape() as tape: circuit(weights) @@ -741,7 +585,7 @@ def circuit(weights): tapes, proc_fn = qml.metric_tensor(tape, argnum=(0, 1, 3)) res = qml.execute(tapes, dev, None, interface=interface) mt013 = proc_fn(res) - assert isinstance(mt013, jax.numpy.ndarray) + assert isinstance(mt013, type_) assert len(tapes) == 6 assert mt.shape == mt013.shape @@ -753,7 +597,7 @@ def circuit(weights): tapes, proc_fn = qml.metric_tensor(tape, argnum=(2, 3)) res = qml.execute(tapes, dev, None, interface=interface) mt23 = proc_fn(res) - assert isinstance(mt23, jax.numpy.ndarray) + assert isinstance(mt23, type_) assert len(tapes) == 1 assert mt.shape == mt23.shape @@ -764,7 +608,7 @@ def circuit(weights): tapes, proc_fn = qml.metric_tensor(tape, argnum=0) res = qml.execute(tapes, dev, None, interface=interface) mt0 = proc_fn(res) - assert isinstance(mt0, jax.numpy.ndarray) + assert isinstance(mt0, type_) assert len(tapes) == 1 assert mt.shape == mt0.shape diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py deleted file mode 100644 index 2a6ee306508..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_autograd_default_qubit_2.py +++ /dev/null @@ -1,839 +0,0 @@ -# Copyright 2018-2023 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. -"""Autograd specific tests for execute and default qubit 2.""" -import autograd -import pytest -from param_shift_dev import ParamShiftDerivativesDevice - -import pennylane as qml -from pennylane import execute -from pennylane import numpy as np -from pennylane.devices import DefaultQubit -from pennylane.gradients import param_shift -from pennylane.measurements import Shots - -pytestmark = pytest.mark.autograd - - -# pylint: disable=too-few-public-methods -class TestCaching: - """Tests for caching behaviour""" - - @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum - when computing Hessians.""" - dev = DefaultQubit() - params = np.arange(1, num_params + 1) / 10 - - N = len(params) - - def cost(x, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - - for i in range(2, num_params): - qml.RZ(x[i], wires=[i % 2]) - - qml.CNOT(wires=[0, 1]) - qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], dev, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 - )[0] - - # No caching: number of executions is not ideal - with qml.Tracker(dev) as tracker: - hess1 = qml.jacobian(qml.grad(cost))(params, cache=False) - - if num_params == 2: - # compare to theoretical result - x, y, *_ = params - expected = np.array( - [ - [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], - ] - ) - assert np.allclose(expected, hess1) - - expected_runs = 1 # forward pass - - # Jacobian of an involutory observable: - # ------------------------------------ - # - # 2 * N execs: evaluate the analytic derivative of - # 1 execs: Get , the expectation value of the tape with unshifted parameters. - num_shifted_evals = 2 * N - runs_for_jacobian = num_shifted_evals + 1 - expected_runs += runs_for_jacobian - - # Each tape used to compute the Jacobian is then shifted again - expected_runs += runs_for_jacobian * num_shifted_evals - assert tracker.totals["executions"] == expected_runs - - # Use caching: number of executions is ideal - - with qml.Tracker(dev) as tracker2: - hess2 = qml.jacobian(qml.grad(cost))(params, cache=True) - assert np.allclose(hess1, hess2) - - expected_runs_ideal = 1 # forward pass - expected_runs_ideal += 2 * N # Jacobian - expected_runs_ideal += N + 1 # Hessian diagonal - expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert tracker2.totals["executions"] == expected_runs_ideal - assert expected_runs_ideal < expected_runs - - def test_single_backward_pass_batch(self): - """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift - is requested for multiple tapes.""" - - dev = qml.device("default.qubit") - - def f(x): - tape1 = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.probs(wires=0)]) - tape2 = qml.tape.QuantumScript([qml.RY(x, 0)], [qml.probs(wires=0)]) - - results = qml.execute([tape1, tape2], dev, gradient_fn=qml.gradients.param_shift) - return results[0] + results[1] - - x = qml.numpy.array(0.1) - with dev.tracker: - out = qml.jacobian(f)(x) - - assert dev.tracker.totals["batches"] == 2 - assert dev.tracker.history["simulations"] == [1, 1, 1, 1, 1, 1] - expected = [-2 * np.cos(x / 2) * np.sin(x / 2), 2 * np.sin(x / 2) * np.cos(x / 2)] - assert qml.math.allclose(out, expected) - - -# add tests for lightning 2 when possible -# set rng for device when possible -test_matrix = [ - ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), - ({"gradient_fn": param_shift}, Shots((100000, 100000)), DefaultQubit(seed=42)), - ({"gradient_fn": param_shift}, Shots(None), DefaultQubit()), - ({"gradient_fn": "backprop"}, Shots(None), DefaultQubit()), - ( - {"gradient_fn": "adjoint", "grad_on_execution": True, "device_vjp": False}, - Shots(None), - DefaultQubit(), - ), - ( - { - "gradient_fn": "adjoint", - "grad_on_execution": False, - "device_vjp": False, - }, - Shots(None), - DefaultQubit(), - ), - ({"gradient_fn": "adjoint", "device_vjp": True}, Shots(None), DefaultQubit()), - ( - {"gradient_fn": "device", "device_vjp": False}, - Shots((100000, 100000)), - ParamShiftDerivativesDevice(seed=904747894), - ), - ( - {"gradient_fn": "device", "device_vjp": True}, - Shots((100000, 100000)), - ParamShiftDerivativesDevice(seed=10490244), - ), -] - - -def atol_for_shots(shots): - """Return higher tolerance if finite shots.""" - return 1e-2 if shots else 1e-6 - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -class TestAutogradExecuteIntegration: - """Test the autograd interface execute function - integrates well for both forward and backward execution""" - - def test_execution(self, execute_kwargs, shots, device): - """Test execution""" - - def cost(a, b): - ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] - tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - - ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] - tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - - return execute([tape1, tape2], device, **execute_kwargs) - - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=False) - with device.tracker: - res = cost(a, b) - - if execute_kwargs.get("grad_on_execution", False): - assert device.tracker.totals["execute_and_derivative_batches"] == 1 - else: - assert device.tracker.totals["batches"] == 1 - assert device.tracker.totals["executions"] == 2 # different wires so different hashes - - assert len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - assert qml.math.allclose(res[0], np.cos(a) * np.cos(b), atol=atol_for_shots(shots)) - assert qml.math.allclose(res[1], np.cos(a) * np.cos(b), atol=atol_for_shots(shots)) - - def test_scalar_jacobian(self, execute_kwargs, shots, device): - """Test scalar jacobian calculation""" - a = np.array(0.1, requires_grad=True) - - def cost(a): - tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - if shots.has_partitioned_shots: - res = qml.jacobian(lambda x: qml.math.hstack(cost(x)))(a) - else: - res = qml.jacobian(cost)(a) - assert res.shape == () # pylint: disable=no-member - - # compare to standard tape jacobian - tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) - tape.trainable_params = [0] - tapes, fn = param_shift(tape) - expected = fn(device.execute(tapes)) - - assert expected.shape == () - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res, -np.sin(a), atol=atol_for_shots(shots)) - - def test_jacobian(self, execute_kwargs, shots, device): - """Test jacobian calculation""" - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - def cost(a, b): - ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - res = cost(a, b) - expected = [np.cos(a), -np.cos(a) * np.sin(b)] - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = qml.jacobian(cost)(a, b) - assert isinstance(res, tuple) and len(res) == 2 - if shots.has_partitioned_shots: - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) - for _r, _e in zip(res, expected): - assert np.allclose(_r[:2], _e, atol=atol_for_shots(shots)) - assert np.allclose(_r[2:], _e, atol=atol_for_shots(shots)) - else: - assert res[0].shape == (2,) - assert res[1].shape == (2,) - - expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) - for _r, _e in zip(res, expected): - assert np.allclose(_r, _e, atol=atol_for_shots(shots)) - - @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") - def test_tape_no_parameters(self, execute_kwargs, shots, device): - """Test that a tape with no parameters is correctly - ignored during the gradient computation""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5, requires_grad=False), wires=0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape4 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5, requires_grad=False), 0)], - [qml.probs(wires=[0, 1])], - shots=shots, - ) - res = qml.execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) - if shots.has_partitioned_shots: - res = tuple(i for r in res for i in r) - return sum(autograd.numpy.hstack(res)) - - params = np.array([0.1, 0.2], requires_grad=True) - x, y = params - - res = cost(params) - expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) - if shots.has_partitioned_shots: - expected = shots.num_copies * expected - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - grad = qml.grad(cost)(params) - expected = np.array([-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)]) - if shots.has_partitioned_shots: - expected = shots.num_copies * expected - assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - - @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") - def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): - """Test that tapes wit different can be executed and differentiated.""" - - if execute_kwargs["gradient_fn"] == "backprop": - pytest.xfail("backprop is not compatible with something about this situation.") - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - shots=shots, - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5, requires_grad=False), 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - res = execute([tape1, tape2, tape3], device, **execute_kwargs) - if shots.has_partitioned_shots: - res = tuple(i for r in res for i in r) - return autograd.numpy.hstack(res) - - params = np.array([0.1, 0.2], requires_grad=True) - x, y = params - - res = cost(params) - assert isinstance(res, np.ndarray) - if not shots: - assert res.shape == (4,) - - if shots.has_partitioned_shots: - for i in (0, 1): - assert np.allclose(res[2 * i], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[2 * i + 1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[4 + i], np.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[6 + i], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) - else: - assert np.allclose(res[0], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[2], np.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[3], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) - - if shots.has_partitioned_shots: - pytest.xfail("autograd jacobians do not work with ragged results and shot vectors.") - # TODO: autograd jacobians with ragged results and shot vectors - jac = qml.jacobian(cost)(params) - assert isinstance(jac, np.ndarray) - if not shots.has_partitioned_shots: - assert jac.shape == (4, 2) # pylint: disable=no-member - - d1 = -np.sin(x) * np.cos(y) - d2 = -np.cos(x) * np.sin(y) - - if shots.has_partitioned_shots: - assert np.allclose(jac[1], 0, atol=atol_for_shots(shots)) - assert np.allclose(jac[3:4], 0, atol=atol_for_shots(shots)) - - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[2, 0], d1, atol=atol_for_shots(shots)) - - assert np.allclose(jac[6, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[7, 0], d1, atol=atol_for_shots(shots)) - - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[6, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[7, 1], d2, atol=atol_for_shots(shots)) - - else: - assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) - - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - - def test_reusing_quantum_tape(self, execute_kwargs, shots, device): - """Test re-using a quantum tape by passing new parameters""" - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - tape = qml.tape.QuantumScript( - [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], - ) - assert tape.trainable_params == [0, 1] - - def cost(a, b): - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return autograd.numpy.hstack(execute([new_tape], device, **execute_kwargs)[0]) - - jac_fn = qml.jacobian(cost) - jac = jac_fn(a, b) - - a = np.array(0.54, requires_grad=True) - b = np.array(0.8, requires_grad=True) - - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - res2 = cost(2 * a, b) - expected = [np.cos(2 * a), -np.cos(2 * a) * np.sin(b)] - assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = qml.jacobian(lambda a, b: cost(2 * a, b)) - jac = jac_fn(a, b) - expected = ( - [-2 * np.sin(2 * a), 2 * np.sin(2 * a) * np.sin(b)], - [0, -np.cos(2 * a) * np.cos(b)], - ) - assert isinstance(jac, tuple) and len(jac) == 2 - for _j, _e in zip(jac, expected): - assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - - def test_classical_processing(self, execute_kwargs, device, shots): - """Test classical processing within the quantum tape""" - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=False) - c = np.array(0.3, requires_grad=True) - - def cost(a, b, c): - ops = [ - qml.RY(a * c, wires=0), - qml.RZ(b, wires=0), - qml.RX(c + c**2 + np.sin(a), wires=0), - ] - - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - if shots.has_partitioned_shots: - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0]) - return execute([tape], device, **execute_kwargs)[0] - - res = qml.jacobian(cost)(a, b, c) - - # Only two arguments are trainable - assert isinstance(res, tuple) and len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - # I tried getting analytic results for this circuit but I kept being wrong and am giving up - - def test_no_trainable_parameters(self, execute_kwargs, shots, device): - """Test evaluation and Jacobian if there are no trainable parameters""" - a = np.array(0.1, requires_grad=False) - b = np.array(0.2, requires_grad=False) - - def cost(a, b): - ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - res = cost(a, b) - assert res.shape == (2 * shots.num_copies,) if shots else (2,) - - with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - res = qml.jacobian(cost)(a, b) - assert len(res) == 0 - - def loss(a, b): - return np.sum(cost(a, b)) - - with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - res = qml.grad(loss)(a, b) - - assert np.allclose(res, 0) - - def test_matrix_parameter(self, execute_kwargs, device, shots): - """Test that the autograd interface works correctly - with a matrix parameter""" - U = np.array([[0, 1], [1, 0]], requires_grad=False) - a = np.array(0.1, requires_grad=True) - - def cost(a, U): - ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) - return execute([tape], device, **execute_kwargs)[0] - - res = cost(a, U) - assert np.allclose(res, -np.cos(a), atol=atol_for_shots(shots)) - - jac_fn = qml.jacobian(cost) - jac = jac_fn(a, U) - assert isinstance(jac, np.ndarray) - assert np.allclose(jac, np.sin(a), atol=atol_for_shots(shots), rtol=0) - - def test_differentiable_expand(self, execute_kwargs, device, shots): - """Test that operation and nested tapes expansion - is differentiable""" - - class U3(qml.U3): - """Dummy operator.""" - - def decomposition(self): - theta, phi, lam = self.data - wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] - - def cost_fn(a, p): - tape = qml.tape.QuantumScript( - [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] - ) - gradient_fn = execute_kwargs["gradient_fn"] - - if gradient_fn is None: - _gradient_method = None - elif isinstance(gradient_fn, str): - _gradient_method = gradient_fn - else: - _gradient_method = "gradient-transform" - config = qml.devices.ExecutionConfig( - interface="autograd", - gradient_method=_gradient_method, - grad_on_execution=execute_kwargs.get("grad_on_execution", None), - ) - program, _ = device.preprocess(execution_config=config) - return execute([tape], device, **execute_kwargs, transform_program=program)[0] - - a = np.array(0.1, requires_grad=False) - p = np.array([0.1, 0.2, 0.3], requires_grad=True) - - res = cost_fn(a, p) - expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( - np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = qml.jacobian(cost_fn) - res = jac_fn(a, p) - expected = np.array( - [ - np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), - np.cos(p[1]) * np.cos(p[2]) * np.sin(a) - - np.sin(p[1]) - * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), - np.sin(a) - * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), - ] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_probability_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob outputs""" - - def cost(x, y): - ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.probs(wires=0), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = np.array(0.543, requires_grad=True) - y = np.array(-0.654, requires_grad=True) - - res = cost(x, y) - expected = np.array( - [ - [ - np.cos(x / 2) ** 2, - np.sin(x / 2) ** 2, - (1 + np.cos(x) * np.cos(y)) / 2, - (1 - np.cos(x) * np.cos(y)) / 2, - ], - ] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = qml.jacobian(cost) - res = jac_fn(x, y) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - expected = ( - np.array( - [ - [ - -np.sin(x) / 2, - np.sin(x) / 2, - -np.sin(x) * np.cos(y) / 2, - np.sin(x) * np.cos(y) / 2, - ], - ] - ), - np.array( - [ - [0, 0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], - ] - ), - ) - - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - def test_ragged_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob and expval outputs""" - - def cost(x, y): - ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = np.array(0.543, requires_grad=True) - y = np.array(-0.654, requires_grad=True) - - res = cost(x, y) - expected = np.array( - [np.cos(x), (1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = qml.jacobian(cost) - res = jac_fn(x, y) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (3,) - assert res[1].shape == (3,) - - expected = ( - np.array([-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]), - np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]), - ) - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - -class TestHigherOrderDerivatives: - """Test that the autograd execute function can be differentiated""" - - @pytest.mark.parametrize( - "params", - [ - np.array([0.543, -0.654], requires_grad=True), - np.array([0, -0.654], requires_grad=True), - np.array([-2.0, 0], requires_grad=True), - ], - ) - def test_parameter_shift_hessian(self, params, tol): - """Tests that the output of the parameter-shift transform - can be differentiated using autograd, yielding second derivatives.""" - dev = DefaultQubit() - - def cost_fn(x): - ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params - expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = qml.grad(cost_fn)(params) - expected = np.array( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = qml.jacobian(qml.grad(cost_fn))(params) - expected = np.array( - [ - [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_max_diff(self, tol): - """Test that setting the max_diff parameter blocks higher-order - derivatives""" - dev = DefaultQubit() - params = np.array([0.543, -0.654], requires_grad=True) - - def cost_fn(x): - ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params - expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = qml.grad(cost_fn)(params) - expected = np.array( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - with pytest.warns(UserWarning, match="Output seems independent"): - res = qml.jacobian(qml.grad(cost_fn))(params) - - expected = np.zeros([2, 2]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -@pytest.mark.usefixtures("use_legacy_and_new_opmath") -class TestHamiltonianWorkflows: - """Test that tapes ending with expectations - of Hamiltonians provide correct results and gradients""" - - @pytest.fixture - def cost_fn(self, execute_kwargs, shots, device): - """Cost function for gradient tests""" - - def _cost_fn(weights, coeffs1, coeffs2): - obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] - H1 = qml.Hamiltonian(coeffs1, obs1) - if qml.operation.active_new_opmath(): - H1 = qml.pauli.pauli_sentence(H1).operation() - - obs2 = [qml.PauliZ(0)] - H2 = qml.Hamiltonian(coeffs2, obs2) - if qml.operation.active_new_opmath(): - H2 = qml.pauli.pauli_sentence(H2).operation() - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H1) - qml.expval(H2) - - tape = qml.tape.QuantumScript.from_queue(q, shots=shots) - return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - - return _cost_fn - - @staticmethod - def cost_fn_expected(weights, coeffs1, coeffs2): - """Analytic value of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] - - @staticmethod - def cost_fn_jacobian(weights, coeffs1, coeffs2): - """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return np.array( - [ - [ - -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), - b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), - np.cos(x), - np.cos(x) * np.sin(y), - -(np.sin(x) * np.sin(y)), - 0, - ], - [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], - ] - ) - - def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with no trainable parameters.""" - - if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): - pytest.skip("adjoint differentiation does not suppport hamiltonians.") - - coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False) - coeffs2 = np.array([0.7], requires_grad=False) - weights = np.array([0.4, 0.5], requires_grad=True) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - if shots.has_partitioned_shots: - assert np.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with trainable parameters.""" - if execute_kwargs["gradient_fn"] == "adjoint": - pytest.skip("trainable hamiltonians not supported with adjoint") - if qml.operation.active_new_opmath(): - pytest.skip("parameter shift derivatives do not yet support sums.") - - coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True) - coeffs2 = np.array([0.7], requires_grad=True) - weights = np.array([0.4, 0.5], requires_grad=True) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2)) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - pytest.xfail( - "multiple hamiltonians with shot vectors does not seem to be differentiable." - ) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py deleted file mode 100644 index 87201a27618..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_execute_default_qubit_2.py +++ /dev/null @@ -1,229 +0,0 @@ -# Copyright 2018-2023 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. -"""Tests for exeuction with default qubit 2 independent of any interface.""" -import pytest - -import pennylane as qml -from pennylane import numpy as np -from pennylane.devices import DefaultQubit -from pennylane.workflow.execution import _preprocess_expand_fn - - -class TestPreprocessExpandFn: - """Tests the _preprocess_expand_fn helper function.""" - - def test_provided_is_callable(self): - """Test that if the expand_fn is not "device", it is simply returned.""" - - dev = DefaultQubit() - - def f(tape): - return tape - - out = _preprocess_expand_fn(f, dev, 10) - assert out is f - - def test_new_device_blank_expand_fn(self): - """Test that the expand_fn is blank if is new device.""" - - dev = DefaultQubit() - - out = _preprocess_expand_fn("device", dev, 10) - - x = [1] - assert out(x) is x - - -class TestBatchTransformHelper: - """Unit tests for the _batch_transform helper function.""" - - def test_warns_if_requested_off(self): - """Test that a warning is raised if the the batch transform is requested to not be used.""" - - # pylint: disable=too-few-public-methods - class CustomOp(qml.operation.Operator): - """Dummy operator.""" - - def decomposition(self): - return [qml.PauliX(self.wires[0])] - - dev = DefaultQubit() - - qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))]) - - with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"): - program, _ = dev.preprocess() - with pytest.warns( - qml.PennyLaneDeprecationWarning, - match="The device_batch_transform argument is deprecated", - ): - qml.execute( - (qs, qs), device=dev, device_batch_transform=False, transform_program=program - ) - - def test_split_and_expand_performed(self): - """Test that preprocess returns the correct tapes when splitting and expanding - is needed.""" - - class NoMatOp(qml.operation.Operation): - """Dummy operation for expanding circuit.""" - - # pylint: disable=missing-function-docstring - num_wires = 1 - - # pylint: disable=arguments-renamed, invalid-overridden-method - @property - def has_matrix(self): - return False - - def decomposition(self): - return [qml.PauliX(self.wires), qml.PauliY(self.wires)] - - ops = [qml.Hadamard(0), NoMatOp(1), qml.RX([np.pi, np.pi / 2], wires=1)] - # Need to specify grouping type to transform tape - measurements = [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(1))] - tapes = [ - qml.tape.QuantumScript(ops=ops, measurements=[measurements[0]]), - qml.tape.QuantumScript(ops=ops, measurements=[measurements[1]]), - ] - - dev = DefaultQubit() - config = qml.devices.ExecutionConfig(gradient_method="adjoint") - - program, new_config = dev.preprocess(config) - res_tapes, batch_fn = program(tapes) - expected_ops = [ - [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi, wires=1)], - [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi / 2, wires=1)], - ] - - assert len(res_tapes) == 4 - for i, t in enumerate(res_tapes): - for op, expected_op in zip(t.operations, expected_ops[i % 2]): - qml.assert_equal(op, expected_op) - assert len(t.measurements) == 1 - if i < 2: - qml.assert_equal(t.measurements[0], measurements[0]) - else: - qml.assert_equal(t.measurements[0], measurements[1]) - - input = ([[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]]) - assert np.array_equal(batch_fn(input), np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])) - - assert new_config.grad_on_execution - assert new_config.use_device_gradient - - -def test_warning_if_not_device_batch_transform(): - """Test that a warning is raised if the users requests to not run device batch transform.""" - - # pylint: disable=too-few-public-methods - class CustomOp(qml.operation.Operator): - """Dummy operator.""" - - def decomposition(self): - return [qml.PauliX(self.wires[0])] - - dev = DefaultQubit() - - qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))]) - - with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"): - program, _ = dev.preprocess() - with pytest.warns( - qml.PennyLaneDeprecationWarning, - match="The device_batch_transform argument is deprecated", - ): - results = qml.execute( - [qs], dev, device_batch_transform=False, transform_program=program - ) - - assert len(results) == 1 - assert qml.math.allclose(results[0], -1) - - -@pytest.mark.parametrize("gradient_fn", (None, "backprop", qml.gradients.param_shift)) -def test_caching(gradient_fn): - """Test that cache execute returns the cached result if the same script is executed - multiple times, both in multiple times in a batch and in separate batches.""" - dev = DefaultQubit() - - qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) - - cache = {} - - with qml.Tracker(dev) as tracker: - results = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn) - results2 = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn) - - assert len(cache) == 1 - assert cache[qs.hash] == -1.0 - - assert list(results) == [-1.0, -1.0] - assert list(results2) == [-1.0, -1.0] - - assert tracker.totals["batches"] == 1 - assert tracker.totals["executions"] == 1 - assert cache[qs.hash] == -1.0 - - -class TestExecuteDeprecations: - """Class to test deprecation warnings in qml.execute. Warnings should be raised even if the default value is used.""" - - @pytest.mark.parametrize("expand_fn", (None, lambda qs: qs, "device")) - def test_expand_fn_is_deprecated(self, expand_fn): - """Test that expand_fn argument of qml.execute is deprecated.""" - dev = DefaultQubit() - qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) - - with pytest.warns( - qml.PennyLaneDeprecationWarning, match="The expand_fn argument is deprecated" - ): - # None is a value used for expand_fn in the QNode - qml.execute([qs], dev, expand_fn=expand_fn) - - def test_max_expansion_is_deprecated(self): - """Test that max_expansion argument of qml.execute is deprecated.""" - dev = DefaultQubit() - qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) - - with pytest.warns( - qml.PennyLaneDeprecationWarning, match="The max_expansion argument is deprecated" - ): - qml.execute([qs], dev, max_expansion=10) - - @pytest.mark.parametrize("override_shots", (False, 10)) - def test_override_shots_is_deprecated(self, override_shots): - """Test that override_shots argument of qml.execute is deprecated.""" - dev = DefaultQubit() - qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) - - with pytest.warns( - qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" - ): - qml.execute([qs], dev, override_shots=override_shots) - - @pytest.mark.parametrize("device_batch_transform", (False, True)) - def test_device_batch_transform_is_deprecated(self, device_batch_transform): - """Test that device_batch_transform argument of qml.execute is deprecated.""" - # Need to use legacy device, otherwise another warning would be raised due to new Device interface - dev = qml.device("default.qubit.legacy", wires=1) - - qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) - - with pytest.warns( - qml.PennyLaneDeprecationWarning, - match="The device_batch_transform argument is deprecated", - ): - qml.execute([qs], dev, device_batch_transform=device_batch_transform) diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py deleted file mode 100644 index 519c0daa028..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_jax_default_qubit_2.py +++ /dev/null @@ -1,817 +0,0 @@ -# Copyright 2018-2023 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. -"""Jax specific tests for execute and default qubit 2.""" -import numpy as np -import pytest -from param_shift_dev import ParamShiftDerivativesDevice - -import pennylane as qml -from pennylane import execute -from pennylane.devices import DefaultQubit -from pennylane.gradients import param_shift -from pennylane.measurements import Shots - -jax = pytest.importorskip("jax") -jnp = pytest.importorskip("jax.numpy") -jax.config.update("jax_enable_x64", True) - -pytestmark = pytest.mark.jax - - -def test_jit_execution(): - """Test that qml.execute can be directly jitted.""" - dev = qml.device("default.qubit") - - tape = qml.tape.QuantumScript( - [qml.RX(jax.numpy.array(0.1), 0)], [qml.expval(qml.s_prod(2.0, qml.PauliZ(0)))] - ) - - out = jax.jit(qml.execute, static_argnames=("device", "gradient_fn"))( - (tape,), device=dev, gradient_fn=qml.gradients.param_shift - ) - expected = 2.0 * jax.numpy.cos(jax.numpy.array(0.1)) - assert qml.math.allclose(out[0], expected) - - -# pylint: disable=too-few-public-methods -class TestCaching: - """Tests for caching behaviour""" - - @pytest.mark.skip("caching is not implemented for jax") - @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum - when computing Hessians.""" - device = DefaultQubit() - params = jnp.arange(1, num_params + 1) / 10 - - N = len(params) - - def cost(x, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - - for i in range(2, num_params): - qml.RZ(x[i], wires=[i % 2]) - - qml.CNOT(wires=[0, 1]) - qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], device, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 - )[0] - - # No caching: number of executions is not ideal - with qml.Tracker(device) as tracker: - hess1 = jax.jacobian(jax.grad(cost))(params, cache=False) - - if num_params == 2: - # compare to theoretical result - x, y, *_ = params - expected = jnp.array( - [ - [2 * jnp.cos(2 * x) * jnp.sin(y) ** 2, jnp.sin(2 * x) * jnp.sin(2 * y)], - [jnp.sin(2 * x) * jnp.sin(2 * y), -2 * jnp.cos(x) ** 2 * jnp.cos(2 * y)], - ] - ) - assert np.allclose(expected, hess1) - - expected_runs = 1 # forward pass - - # Jacobian of an involutory observable: - # ------------------------------------ - # - # 2 * N execs: evaluate the analytic derivative of - # 1 execs: Get , the expectation value of the tape with unshifted parameters. - num_shifted_evals = 2 * N - runs_for_jacobian = num_shifted_evals + 1 - expected_runs += runs_for_jacobian - - # Each tape used to compute the Jacobian is then shifted again - expected_runs += runs_for_jacobian * num_shifted_evals - assert tracker.totals["executions"] == expected_runs - - # Use caching: number of executions is ideal - - with qml.Tracker(device) as tracker2: - hess2 = jax.jacobian(jax.grad(cost))(params, cache=True) - assert np.allclose(hess1, hess2) - - expected_runs_ideal = 1 # forward pass - expected_runs_ideal += 2 * N # Jacobian - expected_runs_ideal += N + 1 # Hessian diagonal - expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert tracker2.totals["executions"] == expected_runs_ideal - assert expected_runs_ideal < expected_runs - - -# add tests for lightning 2 when possible -# set rng for device when possible -no_shots = Shots(None) -shots_2_10k = Shots((10000, 10000)) -dev_def = DefaultQubit() -dev_ps = ParamShiftDerivativesDevice(seed=54353453) -test_matrix = [ - ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), # 0 - ({"gradient_fn": param_shift}, no_shots, dev_def), # 1 - ({"gradient_fn": "backprop"}, no_shots, dev_def), # 2 - ({"gradient_fn": "adjoint"}, no_shots, dev_def), # 3 - ({"gradient_fn": "adjoint", "device_vjp": True}, no_shots, dev_def), # 4 - ({"gradient_fn": "device"}, shots_2_10k, dev_ps), # 5 -] - - -def atol_for_shots(shots): - """Return higher tolerance if finite shots.""" - return 3e-2 if shots else 1e-6 - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -class TestJaxExecuteIntegration: - """Test the jax interface execute function - integrates well for both forward and backward execution""" - - def test_execution(self, execute_kwargs, shots, device): - """Test execution""" - - def cost(a, b): - ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] - tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - - ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] - tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - - return execute([tape1, tape2], device, **execute_kwargs) - - a = jnp.array(0.1) - b = np.array(0.2) - with device.tracker: - res = cost(a, b) - - if execute_kwargs.get("gradient_fn", None) == "adjoint": - assert device.tracker.totals.get("execute_and_derivative_batches", 0) == 0 - else: - assert device.tracker.totals["batches"] == 1 - assert device.tracker.totals["executions"] == 2 # different wires so different hashes - - assert len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - assert qml.math.allclose(res[0], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots)) - assert qml.math.allclose(res[1], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots)) - - def test_scalar_jacobian(self, execute_kwargs, shots, device): - """Test scalar jacobian calculation""" - a = jnp.array(0.1) - - def cost(a): - tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = jax.jacobian(cost)(a) - if not shots.has_partitioned_shots: - assert res.shape == () # pylint: disable=no-member - - # compare to standard tape jacobian - tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) - tape.trainable_params = [0] - tapes, fn = param_shift(tape) - expected = fn(device.execute(tapes)) - - assert expected.shape == () - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res, -jnp.sin(a), atol=atol_for_shots(shots)) - - def test_jacobian(self, execute_kwargs, shots, device): - """Test jacobian calculation""" - - a = jnp.array(0.1) - b = jnp.array(0.2) - - def cost(a, b): - ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = cost(a, b) - expected = [jnp.cos(a), -jnp.cos(a) * jnp.sin(b)] - if shots.has_partitioned_shots: - assert np.allclose(res[0], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected, atol=2 * atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - g = jax.jacobian(cost, argnums=[0, 1])(a, b) - assert isinstance(g, tuple) and len(g) == 2 - - expected = ([-jnp.sin(a), jnp.sin(a) * jnp.sin(b)], [0, -jnp.cos(a) * jnp.cos(b)]) - - if shots.has_partitioned_shots: - for i in (0, 1): - assert np.allclose(g[i][0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[i][1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[i][0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[i][1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(g[0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(g[1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0) - - def test_tape_no_parameters(self, execute_kwargs, shots, device): - """Test that a tape with no parameters is correctly - ignored during the gradient computation""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(jnp.array(0.5), wires=0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape4 = qml.tape.QuantumScript( - [qml.RY(jnp.array(0.5), 0)], [qml.probs(wires=[0, 1])], shots=shots - ) - res = execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) - res = jax.tree_util.tree_leaves(res) - out = sum(jnp.hstack(res)) - if shots.has_partitioned_shots: - out = out / shots.num_copies - return out - - params = jnp.array([0.1, 0.2]) - - x, y = params - - res = cost(params) - expected = 2 + jnp.cos(0.5) + jnp.cos(x) * jnp.cos(y) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - grad = jax.grad(cost)(params) - expected = [-jnp.cos(y) * jnp.sin(x), -jnp.cos(x) * jnp.sin(y)] - assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - - # pylint: disable=too-many-statements - def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): - """Test that tapes wit different can be executed and differentiated.""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - shots=shots, - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), 0)], [qml.expval(qml.PauliZ(0))], shots=shots - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - res = qml.execute([tape1, tape2, tape3], device, **execute_kwargs) - leaves = jax.tree_util.tree_leaves(res) - return jnp.hstack(leaves) - - params = jnp.array([0.1, 0.2]) - x, y = params - - res = cost(params) - assert isinstance(res, jax.Array) - assert res.shape == (4 * shots.num_copies,) if shots.has_partitioned_shots else (4,) - - if shots.has_partitioned_shots: - assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[2], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[3], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[4], jnp.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[5], jnp.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[6], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[7], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - else: - assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[2], jnp.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[3], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) - - jac = jax.jacobian(cost)(params) - assert isinstance(jac, jnp.ndarray) - assert ( - jac.shape == (8, 2) if shots.has_partitioned_shots else (4, 2) - ) # pylint: disable=no-member - - if shots.has_partitioned_shots: - assert np.allclose(jac[1], 0, atol=atol_for_shots(shots)) - assert np.allclose(jac[3:5], 0, atol=atol_for_shots(shots)) - else: - assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) - - d1 = -jnp.sin(x) * jnp.cos(y) - if shots.has_partitioned_shots: - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[2, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[6, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[7, 0], d1, atol=atol_for_shots(shots)) - else: - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - - d2 = -jnp.cos(x) * jnp.sin(y) - if shots.has_partitioned_shots: - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[2, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[6, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[7, 1], d2, atol=atol_for_shots(shots)) - else: - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - - def test_reusing_quantum_tape(self, execute_kwargs, shots, device): - """Test re-using a quantum tape by passing new parameters""" - if execute_kwargs["gradient_fn"] == param_shift: - pytest.skip("Basic QNode execution wipes out trainable params with param-shift") - - a = jnp.array(0.1) - b = jnp.array(0.2) - - tape = qml.tape.QuantumScript( - [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], - shots=shots, - ) - assert tape.trainable_params == [0, 1] - - def cost(a, b): - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return jnp.hstack(execute([new_tape], device, **execute_kwargs)[0]) - - jac_fn = jax.jacobian(cost, argnums=[0, 1]) - jac = jac_fn(a, b) - - a = jnp.array(0.54) - b = jnp.array(0.8) - - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - res2 = cost(2 * a, b) - expected = [jnp.cos(2 * a), -jnp.cos(2 * a) * jnp.sin(b)] - if shots.has_partitioned_shots: - assert np.allclose(res2[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res2[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(lambda a, b: cost(2 * a, b), argnums=[0, 1]) - jac = jac_fn(a, b) - expected = ( - [-2 * jnp.sin(2 * a), 2 * jnp.sin(2 * a) * jnp.sin(b)], - [0, -jnp.cos(2 * a) * jnp.cos(b)], - ) - assert isinstance(jac, tuple) and len(jac) == 2 - if shots.has_partitioned_shots: - for offset in (0, 2): - assert np.allclose(jac[0][0 + offset], expected[0][0], atol=atol_for_shots(shots)) - assert np.allclose(jac[0][1 + offset], expected[0][1], atol=atol_for_shots(shots)) - assert np.allclose(jac[1][0 + offset], expected[1][0], atol=atol_for_shots(shots)) - assert np.allclose(jac[1][1 + offset], expected[1][1], atol=atol_for_shots(shots)) - else: - for _j, _e in zip(jac, expected): - assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - - def test_classical_processing(self, execute_kwargs, shots, device): - """Test classical processing within the quantum tape""" - a = jnp.array(0.1) - b = jnp.array(0.2) - c = jnp.array(0.3) - - def cost(a, b, c): - ops = [ - qml.RY(a * c, wires=0), - qml.RZ(b, wires=0), - qml.RX(c + c**2 + jnp.sin(a), wires=0), - ] - - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = jax.jacobian(cost, argnums=[0, 2])(a, b, c) - - # Only two arguments are trainable - assert isinstance(res, tuple) and len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - # I tried getting analytic results for this circuit but I kept being wrong and am giving up - - def test_matrix_parameter(self, execute_kwargs, device, shots): - """Test that the jax interface works correctly - with a matrix parameter""" - U = jnp.array([[0, 1], [1, 0]]) - a = jnp.array(0.1) - - def cost(a, U): - ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = cost(a, U) - assert np.allclose(res, -jnp.cos(a), atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(cost) - jac = jac_fn(a, U) - if not shots.has_partitioned_shots: - assert isinstance(jac, jnp.ndarray) - assert np.allclose(jac, jnp.sin(a), atol=atol_for_shots(shots), rtol=0) - - def test_differentiable_expand(self, execute_kwargs, device, shots): - """Test that operation and nested tapes expansion - is differentiable""" - - class U3(qml.U3): - """Dummy operator.""" - - def decomposition(self): - theta, phi, lam = self.data - wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] - - def cost_fn(a, p): - tape = qml.tape.QuantumScript( - [qml.RX(a, wires=0), U3(*p, wires=0)], - [qml.expval(qml.PauliX(0))], - shots=shots, - ) - gradient_fn = execute_kwargs["gradient_fn"] - if gradient_fn is None: - _gradient_method = None - elif isinstance(gradient_fn, str): - _gradient_method = gradient_fn - else: - _gradient_method = "gradient-transform" - conf = qml.devices.ExecutionConfig( - interface="autograd", - gradient_method=_gradient_method, - grad_on_execution=execute_kwargs.get("grad_on_execution", None), - ) - program, _ = device.preprocess(execution_config=conf) - return execute([tape], device, **execute_kwargs, transform_program=program)[0] - - a = jnp.array(0.1) - p = jnp.array([0.1, 0.2, 0.3]) - - res = cost_fn(a, p) - expected = jnp.cos(a) * jnp.cos(p[1]) * jnp.sin(p[0]) + jnp.sin(a) * ( - jnp.cos(p[2]) * jnp.sin(p[1]) + jnp.cos(p[0]) * jnp.cos(p[1]) * jnp.sin(p[2]) - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(cost_fn, argnums=[1]) - res = jac_fn(a, p) - expected = jnp.array( - [ - jnp.cos(p[1]) - * (jnp.cos(a) * jnp.cos(p[0]) - jnp.sin(a) * jnp.sin(p[0]) * jnp.sin(p[2])), - jnp.cos(p[1]) * jnp.cos(p[2]) * jnp.sin(a) - - jnp.sin(p[1]) - * (jnp.cos(a) * jnp.sin(p[0]) + jnp.cos(p[0]) * jnp.sin(a) * jnp.sin(p[2])), - jnp.sin(a) - * (jnp.cos(p[0]) * jnp.cos(p[1]) * jnp.cos(p[2]) - jnp.sin(p[1]) * jnp.sin(p[2])), - ] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_probability_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob outputs""" - - def cost(x, y): - ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.probs(wires=0), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return jnp.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = jnp.array(0.543) - y = jnp.array(-0.654) - - res = cost(x, y) - expected = jnp.array( - [ - [ - jnp.cos(x / 2) ** 2, - jnp.sin(x / 2) ** 2, - (1 + jnp.cos(x) * jnp.cos(y)) / 2, - (1 - jnp.cos(x) * jnp.cos(y)) / 2, - ], - ] - ) - if shots.has_partitioned_shots: - assert np.allclose(res[:, 0:2].flatten(), expected, atol=atol_for_shots(shots)) - assert np.allclose(res[:, 2:].flatten(), expected, atol=atol_for_shots(shots)) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(cost, argnums=[0, 1]) - res = jac_fn(x, y) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (2, 4) if shots.has_partitioned_shots else (4,) - assert res[1].shape == (2, 4) if shots.has_partitioned_shots else (4,) - - expected = ( - jnp.array( - [ - [ - -jnp.sin(x) / 2, - jnp.sin(x) / 2, - -jnp.sin(x) * jnp.cos(y) / 2, - jnp.sin(x) * jnp.cos(y) / 2, - ], - ] - ), - jnp.array( - [ - [0, 0, -jnp.cos(x) * jnp.sin(y) / 2, jnp.cos(x) * jnp.sin(y) / 2], - ] - ), - ) - - if shots.has_partitioned_shots: - assert np.allclose(res[0][:, 0:2].flatten(), expected[0], atol=atol_for_shots(shots)) - assert np.allclose(res[0][:, 2:].flatten(), expected[0], atol=atol_for_shots(shots)) - assert np.allclose(res[1][:, :2].flatten(), expected[1], atol=atol_for_shots(shots)) - assert np.allclose(res[1][:, 2:].flatten(), expected[1], atol=atol_for_shots(shots)) - else: - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - def test_ragged_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob and expval outputs""" - - def cost(x, y): - ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - res = qml.execute([tape], device, **execute_kwargs)[0] - return jnp.hstack(jax.tree_util.tree_leaves(res)) - - x = jnp.array(0.543) - y = jnp.array(-0.654) - - res = cost(x, y) - expected = jnp.array( - [jnp.cos(x), (1 + jnp.cos(x) * jnp.cos(y)) / 2, (1 - jnp.cos(x) * jnp.cos(y)) / 2] - ) - if shots.has_partitioned_shots: - assert np.allclose(res[:3], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[3:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac_fn = jax.jacobian(cost, argnums=[0, 1]) - res = jac_fn(x, y) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,) - assert res[1].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,) - - expected = ( - jnp.array([-jnp.sin(x), -jnp.sin(x) * jnp.cos(y) / 2, jnp.sin(x) * jnp.cos(y) / 2]), - jnp.array([0, -jnp.cos(x) * jnp.sin(y) / 2, jnp.cos(x) * jnp.sin(y) / 2]), - ) - if shots.has_partitioned_shots: - assert np.allclose(res[0][:3], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[0][3:], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1][:3], expected[1], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1][3:], expected[1], atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - -class TestHigherOrderDerivatives: - """Test that the jax execute function can be differentiated""" - - @pytest.mark.parametrize( - "params", - [ - jnp.array([0.543, -0.654]), - jnp.array([0, -0.654]), - jnp.array([-2.0, 0]), - ], - ) - def test_parameter_shift_hessian(self, params, tol): - """Tests that the output of the parameter-shift transform - can be differentiated using jax, yielding second derivatives.""" - dev = DefaultQubit() - - def cost_fn(x): - ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params - expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.grad(cost_fn)(params) - expected = jnp.array( - [-jnp.cos(x) * jnp.cos(2 * y) * jnp.sin(x), -jnp.cos(x) ** 2 * jnp.sin(2 * y)] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.jacobian(jax.grad(cost_fn))(params) - expected = jnp.array( - [ - [-jnp.cos(2 * x) * jnp.cos(2 * y), jnp.sin(2 * x) * jnp.sin(2 * y)], - [jnp.sin(2 * x) * jnp.sin(2 * y), -2 * jnp.cos(x) ** 2 * jnp.cos(2 * y)], - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_max_diff(self, tol): - """Test that setting the max_diff parameter blocks higher-order - derivatives""" - dev = DefaultQubit() - params = jnp.array([0.543, -0.654]) - - def cost_fn(x): - ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params - expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.grad(cost_fn)(params) - expected = jnp.array( - [-jnp.cos(x) * jnp.cos(2 * y) * jnp.sin(x), -jnp.cos(x) ** 2 * jnp.sin(2 * y)] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.jacobian(jax.grad(cost_fn))(params) - expected = jnp.zeros([2, 2]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -@pytest.mark.usefixtures("use_legacy_and_new_opmath") -class TestHamiltonianWorkflows: - """Test that tapes ending with expectations - of Hamiltonians provide correct results and gradients""" - - @pytest.fixture - def cost_fn(self, execute_kwargs, shots, device): - """Cost function for gradient tests""" - - def _cost_fn(weights, coeffs1, coeffs2): - obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] - H1 = qml.Hamiltonian(coeffs1, obs1) - if qml.operation.active_new_opmath(): - H1 = qml.pauli.pauli_sentence(H1).operation() - - obs2 = [qml.PauliZ(0)] - H2 = qml.Hamiltonian(coeffs2, obs2) - if qml.operation.active_new_opmath(): - H2 = qml.pauli.pauli_sentence(H2).operation() - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H1) - qml.expval(H2) - - tape = qml.tape.QuantumScript.from_queue(q, shots=shots) - res = execute([tape], device, **execute_kwargs)[0] - if shots.has_partitioned_shots: - return jnp.hstack(res[0] + res[1]) - return jnp.hstack(res) - - return _cost_fn - - @staticmethod - def cost_fn_expected(weights, coeffs1, coeffs2): - """Analytic value of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return [-c * jnp.sin(x) * jnp.sin(y) + jnp.cos(x) * (a + b * jnp.sin(y)), d * jnp.cos(x)] - - @staticmethod - def cost_fn_jacobian(weights, coeffs1, coeffs2): - """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return jnp.array( - [ - [ - -c * jnp.cos(x) * jnp.sin(y) - jnp.sin(x) * (a + b * jnp.sin(y)), - b * jnp.cos(x) * jnp.cos(y) - c * jnp.cos(y) * jnp.sin(x), - jnp.cos(x), - jnp.cos(x) * jnp.sin(y), - -(jnp.sin(x) * jnp.sin(y)), - 0, - ], - [-d * jnp.sin(x), 0, 0, 0, 0, jnp.cos(x)], - ] - ) - - def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with no trainable parameters.""" - - if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): - pytest.skip("adjoint differentiation does not suppport hamiltonians.") - - coeffs1 = jnp.array([0.1, 0.2, 0.3]) - coeffs2 = jnp.array([0.7]) - weights = jnp.array([0.4, 0.5]) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = jax.jacobian(cost_fn)(weights, coeffs1, coeffs2) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - if shots.has_partitioned_shots: - assert np.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with trainable parameters.""" - if execute_kwargs["gradient_fn"] == "adjoint": - pytest.skip("trainable hamiltonians not supported with adjoint") - if qml.operation.active_new_opmath(): - pytest.skip("parameter shift derivatives do not yet support sums.") - - coeffs1 = jnp.array([0.1, 0.2, 0.3]) - coeffs2 = jnp.array([0.7]) - weights = jnp.array([0.4, 0.5]) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = jnp.hstack(jax.jacobian(cost_fn, argnums=[0, 1, 2])(weights, coeffs1, coeffs2)) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - pytest.xfail( - "multiple hamiltonians with shot vectors does not seem to be differentiable." - ) - else: - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py deleted file mode 100644 index 11a54575cfa..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_default_qubit_2.py +++ /dev/null @@ -1,844 +0,0 @@ -# Copyright 2018-2023 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. -"""Tensorflow specific tests for execute and default qubit 2.""" -import numpy as np -import pytest - -import pennylane as qml -from pennylane import execute -from pennylane.devices import DefaultQubit -from pennylane.gradients import param_shift - -pytestmark = pytest.mark.tf -tf = pytest.importorskip("tensorflow") - - -# pylint: disable=too-few-public-methods -class TestCaching: - """Tests for caching behaviour""" - - @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum - when computing Hessians.""" - dev = DefaultQubit() - params = tf.Variable(tf.range(1, num_params + 1) / 10) - - N = num_params - - def cost(x, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - - for i in range(2, num_params): - qml.RZ(x[i], wires=[i % 2]) - - qml.CNOT(wires=[0, 1]) - qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], dev, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 - )[0] - - # No caching: number of executions is not ideal - with qml.Tracker(dev) as tracker: - with tf.GradientTape() as jac_tape: - with tf.GradientTape() as grad_tape: - res = cost(params, cache=False) - grad = grad_tape.gradient(res, params) - hess1 = jac_tape.jacobian(grad, params) - - if num_params == 2: - # compare to theoretical result - x, y, *_ = params - expected = tf.convert_to_tensor( - [ - [2 * tf.cos(2 * x) * tf.sin(y) ** 2, tf.sin(2 * x) * tf.sin(2 * y)], - [tf.sin(2 * x) * tf.sin(2 * y), -2 * tf.cos(x) ** 2 * tf.cos(2 * y)], - ] - ) - assert np.allclose(expected, hess1) - - expected_runs = 1 # forward pass - - # Jacobian of an involutory observable: - # ------------------------------------ - # - # 2 * N execs: evaluate the analytic derivative of - # 1 execs: Get , the expectation value of the tape with unshifted parameters. - num_shifted_evals = 2 * N - runs_for_jacobian = num_shifted_evals + 1 - expected_runs += runs_for_jacobian - - # Each tape used to compute the Jacobian is then shifted again - expected_runs += runs_for_jacobian * num_shifted_evals - assert tracker.totals["executions"] == expected_runs - - # Use caching: number of executions is ideal - - with qml.Tracker(dev) as tracker2: - with tf.GradientTape() as jac_tape: - with tf.GradientTape() as grad_tape: - res = cost(params, cache=True) - grad = grad_tape.gradient(res, params) - hess2 = jac_tape.jacobian(grad, params) - assert np.allclose(hess1, hess2) - - expected_runs_ideal = 1 # forward pass - expected_runs_ideal += 2 * N # Jacobian - expected_runs_ideal += N + 1 # Hessian diagonal - expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert tracker2.totals["executions"] == expected_runs_ideal - assert expected_runs_ideal < expected_runs - - -# add tests for lightning 2 when possible -# set rng for device when possible -test_matrix = [ - ({"gradient_fn": param_shift, "interface": "tensorflow"}, 100000, DefaultQubit(seed=42)), # 0 - ({"gradient_fn": param_shift, "interface": "tensorflow"}, None, DefaultQubit()), # 1 - ({"gradient_fn": "backprop", "interface": "tensorflow"}, None, DefaultQubit()), # 2 - ({"gradient_fn": "adjoint", "interface": "tensorflow"}, None, DefaultQubit()), # 3 - ({"gradient_fn": param_shift, "interface": "tf-autograph"}, 100000, DefaultQubit(seed=42)), # 4 - ({"gradient_fn": param_shift, "interface": "tf-autograph"}, None, DefaultQubit()), # 5 - ({"gradient_fn": "backprop", "interface": "tf-autograph"}, None, DefaultQubit()), # 6 - ({"gradient_fn": "adjoint", "interface": "tf-autograph"}, None, DefaultQubit()), # 7 - ({"gradient_fn": "adjoint", "interface": "tf", "device_vjp": True}, None, DefaultQubit()), # 8 -] - - -def atol_for_shots(shots): - """Return higher tolerance if finite shots.""" - return 1e-2 if shots else 1e-6 - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -class TestTensorflowExecuteIntegration: - """Test the tensorflow interface execute function - integrates well for both forward and backward execution""" - - def test_execution(self, execute_kwargs, shots, device): - """Test execution""" - - def cost(a, b): - ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] - tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - - ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] - tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - - return execute([tape1, tape2], device, **execute_kwargs) - - a = tf.Variable(0.1, dtype="float64") - b = tf.constant(0.2, dtype="float64") - with device.tracker: - res = cost(a, b) - - if execute_kwargs.get("gradient_fn", None) == "adjoint" and not execute_kwargs.get( - "device_vjp", False - ): - assert device.tracker.totals["execute_and_derivative_batches"] == 1 - else: - assert device.tracker.totals["batches"] == 1 - assert device.tracker.totals["executions"] == 2 # different wires so different hashes - - assert len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () - - assert qml.math.allclose(res[0], tf.cos(a) * tf.cos(b), atol=atol_for_shots(shots)) - assert qml.math.allclose(res[1], tf.cos(a) * tf.cos(b), atol=atol_for_shots(shots)) - - def test_scalar_jacobian(self, execute_kwargs, shots, device): - """Test scalar jacobian calculation""" - a = tf.Variable(0.1, dtype=tf.float64) - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a): - tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost(a) - res = tape.jacobian(cost_res, a, experimental_use_pfor=not device_vjp) - assert res.shape == () # pylint: disable=no-member - - # compare to standard tape jacobian - tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) - tape.trainable_params = [0] - tapes, fn = param_shift(tape) - expected = fn(device.execute(tapes)) - - assert expected.shape == () - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res, -tf.sin(a), atol=atol_for_shots(shots)) - - def test_jacobian(self, execute_kwargs, shots, device): - """Test jacobian calculation""" - a = tf.Variable(0.1) - b = tf.Variable(0.2) - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a, b): - ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - with tf.GradientTape(persistent=device_vjp) as tape: - res = cost(a, b) - expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) - assert isinstance(jac, list) and len(jac) == 2 - assert jac[0].shape == (2,) - assert jac[1].shape == (2,) - - expected = ([-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)]) - for _r, _e in zip(jac, expected): - assert np.allclose(_r, _e, atol=atol_for_shots(shots)) - - def test_tape_no_parameters(self, execute_kwargs, shots, device): - """Test that a tape with no parameters is correctly - ignored during the gradient computation""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(tf.constant(0.5), wires=0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape4 = qml.tape.QuantumScript( - [qml.RY(tf.constant(0.5), 0)], - [qml.probs(wires=[0, 1])], - shots=shots, - ) - return tf.reduce_sum( - qml.math.hstack( - execute([tape1, tape2, tape3, tape4], device, **execute_kwargs), - like="tensorflow", - ) - ) - - params = tf.Variable([0.1, 0.2]) - x, y = params - - with tf.GradientTape() as tape: - res = cost(params) - expected = 2 + tf.cos(0.5) + tf.cos(x) * tf.cos(y) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - if ( - execute_kwargs.get("interface", "") == "tf-autograph" - and execute_kwargs.get("gradient_fn", "") == "adjoint" - ): - with pytest.raises(NotImplementedError): - tape.gradient(res, params) - return - - grad = tape.gradient(res, params) - expected = [-tf.cos(y) * tf.sin(x), -tf.cos(x) * tf.sin(y)] - assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - - def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): - """Test that tapes wit different can be executed and differentiated.""" - - if ( - execute_kwargs["gradient_fn"] == "adjoint" - and execute_kwargs["interface"] == "tf-autograph" - ): - pytest.skip("Cannot compute the jacobian with adjoint-differentation and tf-autograph") - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - shots=shots, - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(tf.constant(0.5, dtype=tf.float64), 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - return qml.math.hstack( - execute([tape1, tape2, tape3], device, **execute_kwargs), like="tensorflow" - ) - - params = tf.Variable([0.1, 0.2], dtype=tf.float64) - x, y = params - - with tf.GradientTape(persistent=device_vjp) as tape: - res = cost(params) - - assert isinstance(res, tf.Tensor) - assert res.shape == (4,) - - assert np.allclose(res[0], tf.cos(x) * tf.cos(y), atol=atol_for_shots(shots)) - assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) - assert np.allclose(res[2], tf.cos(0.5), atol=atol_for_shots(shots)) - assert np.allclose(res[3], tf.cos(x) * tf.cos(y), atol=atol_for_shots(shots)) - - jac = tape.jacobian(res, params, experimental_use_pfor=not device_vjp) - assert isinstance(jac, tf.Tensor) - assert jac.shape == (4, 2) # pylint: disable=no-member - - assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) - - d1 = -tf.sin(x) * tf.cos(y) - assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - - d2 = -tf.cos(x) * tf.sin(y) - assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) - assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - - def test_reusing_quantum_tape(self, execute_kwargs, shots, device): - """Test re-using a quantum tape by passing new parameters""" - a = tf.Variable(0.1) - b = tf.Variable(0.2) - - tape = qml.tape.QuantumScript( - [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], - ) - assert tape.trainable_params == [0, 1] - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a, b): - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return qml.math.hstack( - execute([new_tape], device, **execute_kwargs)[0], like="tensorflow" - ) - - with tf.GradientTape(persistent=device_vjp) as t: - res = cost(a, b) - - jac = t.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) - a = tf.Variable(0.54, dtype=tf.float64) - b = tf.Variable(0.8, dtype=tf.float64) - - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - - with tf.GradientTape(persistent=device_vjp): - res2 = cost(2 * a, b) - - expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] - assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - - with tf.GradientTape(persistent=device_vjp) as t: - res = cost(2 * a, b) - - jac = t.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) - expected = ( - [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], - [0, -tf.cos(2 * a) * tf.cos(b)], - ) - assert isinstance(jac, list) and len(jac) == 2 - for _j, _e in zip(jac, expected): - assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - - def test_classical_processing(self, execute_kwargs, device, shots): - """Test classical processing within the quantum tape""" - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.constant(0.2, dtype=tf.float64) - c = tf.Variable(0.3, dtype=tf.float64) - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a, b, c): - ops = [ - qml.RY(a * c, wires=0), - qml.RZ(b, wires=0), - qml.RX(c + c**2 + tf.sin(a), wires=0), - ] - - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost(a, b, c) - - res = tape.jacobian(cost_res, [a, c], experimental_use_pfor=not device_vjp) - - # Only two arguments are trainable - assert isinstance(res, list) and len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () - - # I tried getting analytic results for this circuit but I kept being wrong and am giving up - - def test_no_trainable_parameters(self, execute_kwargs, shots, device): - """Test evaluation and Jacobian if there are no trainable parameters""" - a = tf.constant(0.1) - b = tf.constant(0.2) - - def cost(a, b): - ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - with tf.GradientTape() as tape: - cost_res = cost(a, b) - - assert cost_res.shape == (2,) - - res = tape.jacobian(cost_res, [a, b]) - assert len(res) == 2 - assert all(r is None for r in res) - - def loss(a, b): - return tf.reduce_sum(cost(a, b)) - - with tf.GradientTape() as tape: - loss_res = loss(a, b) - - res = tape.gradient(loss_res, [a, b]) - assert all(r is None for r in res) - - def test_matrix_parameter(self, execute_kwargs, device, shots): - """Test that the tensorflow interface works correctly - with a matrix parameter""" - U = tf.constant([[0, 1], [1, 0]], dtype=tf.complex128) - a = tf.Variable(0.1) - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(a, U): - ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) - return execute([tape], device, **execute_kwargs)[0] - - with tf.GradientTape(persistent=device_vjp) as tape: - res = cost(a, U) - - assert np.allclose(res, -tf.cos(a), atol=atol_for_shots(shots)) - - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) - assert isinstance(jac, tf.Tensor) - assert np.allclose(jac, tf.sin(a), atol=atol_for_shots(shots), rtol=0) - - def test_differentiable_expand(self, execute_kwargs, device, shots): - """Test that operation and nested tapes expansion - is differentiable""" - - device_vjp = execute_kwargs.get("device_vjp", False) - - class U3(qml.U3): - """Dummy operator.""" - - def decomposition(self): - theta, phi, lam = self.data - wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] - - def cost_fn(a, p): - tape = qml.tape.QuantumScript( - [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] - ) - gradient_fn = execute_kwargs["gradient_fn"] - if gradient_fn is None: - _gradient_method = None - elif isinstance(gradient_fn, str): - _gradient_method = gradient_fn - else: - _gradient_method = "gradient-transform" - config = qml.devices.ExecutionConfig( - interface="autograd", - gradient_method=_gradient_method, - grad_on_execution=execute_kwargs.get("grad_on_execution", None), - ) - program, _ = device.preprocess(execution_config=config) - return execute([tape], device, **execute_kwargs, transform_program=program)[0] - - a = tf.constant(0.1) - p = tf.Variable([0.1, 0.2, 0.3]) - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost_fn(a, p) - - expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( - tf.cos(p[2]) * tf.sin(p[1]) + tf.cos(p[0]) * tf.cos(p[1]) * tf.sin(p[2]) - ) - assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - - res = tape.jacobian(cost_res, p, experimental_use_pfor=not device_vjp) - expected = tf.convert_to_tensor( - [ - tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), - tf.cos(p[1]) * tf.cos(p[2]) * tf.sin(a) - - tf.sin(p[1]) - * (tf.cos(a) * tf.sin(p[0]) + tf.cos(p[0]) * tf.sin(a) * tf.sin(p[2])), - tf.sin(a) - * (tf.cos(p[0]) * tf.cos(p[1]) * tf.cos(p[2]) - tf.sin(p[1]) * tf.sin(p[2])), - ] - ) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_probability_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob outputs""" - - def cost(x, y): - ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.probs(wires=0), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - device_vjp = execute_kwargs.get("device_vjp", False) - - x = tf.Variable(0.543, dtype=tf.float64) - y = tf.Variable(-0.654, dtype=tf.float64) - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost(x, y) - - expected = tf.convert_to_tensor( - [ - [ - tf.cos(x / 2) ** 2, - tf.sin(x / 2) ** 2, - (1 + tf.cos(x) * tf.cos(y)) / 2, - (1 - tf.cos(x) * tf.cos(y)) / 2, - ], - ] - ) - assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - - if ( - execute_kwargs.get("interface", "") == "tf-autograph" - and execute_kwargs.get("gradient_fn", "") == "adjoint" - ): - with pytest.raises(tf.errors.UnimplementedError): - tape.jacobian(cost_res, [x, y]) - return - res = tape.jacobian(cost_res, [x, y], experimental_use_pfor=not device_vjp) - assert isinstance(res, list) and len(res) == 2 - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - expected = ( - tf.convert_to_tensor( - [ - [ - -tf.sin(x) / 2, - tf.sin(x) / 2, - -tf.sin(x) * tf.cos(y) / 2, - tf.sin(x) * tf.cos(y) / 2, - ], - ] - ), - tf.convert_to_tensor( - [ - [0, 0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], - ] - ), - ) - - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - def test_ragged_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob and expval outputs""" - - device_vjp = execute_kwargs.get("device_vjp", False) - - def cost(x, y): - ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - x = tf.Variable(0.543, dtype=tf.float64) - y = tf.Variable(-0.654, dtype=tf.float64) - - with tf.GradientTape(persistent=device_vjp) as tape: - cost_res = cost(x, y) - - expected = tf.convert_to_tensor( - [tf.cos(x), (1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2] - ) - assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - - if ( - execute_kwargs.get("interface", "") == "tf-autograph" - and execute_kwargs.get("gradient_fn", "") == "adjoint" - ): - with pytest.raises(tf.errors.UnimplementedError): - tape.jacobian(cost_res, [x, y]) - return - res = tape.jacobian(cost_res, [x, y], experimental_use_pfor=not device_vjp) - assert isinstance(res, list) and len(res) == 2 - assert res[0].shape == (3,) - assert res[1].shape == (3,) - - expected = ( - tf.convert_to_tensor( - [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.sin(x) * tf.cos(y) / 2] - ), - tf.convert_to_tensor([0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2]), - ) - assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - -class TestHigherOrderDerivatives: - """Test that the tensorflow execute function can be differentiated""" - - @pytest.mark.parametrize( - "params", - [ - tf.Variable([0.543, -0.654], dtype=tf.float64), - tf.Variable([0, -0.654], dtype=tf.float64), - tf.Variable([-2.0, 0], dtype=tf.float64), - ], - ) - def test_parameter_shift_hessian(self, params, tol): - """Tests that the output of the parameter-shift transform - can be differentiated using tensorflow, yielding second derivatives.""" - dev = DefaultQubit() - - def cost_fn(x): - ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) - return result[0] + result[1][0] - - with tf.GradientTape() as jac_tape: - with tf.GradientTape() as grad_tape: - res = cost_fn(params) - grad = grad_tape.gradient(res, params) - hess = jac_tape.jacobian(grad, params) - - x, y = params - expected = 0.5 * (3 + tf.cos(x) ** 2 * tf.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - expected = tf.convert_to_tensor( - [-tf.cos(x) * tf.cos(2 * y) * tf.sin(x), -tf.cos(x) ** 2 * tf.sin(2 * y)] - ) - assert np.allclose(grad, expected, atol=tol, rtol=0) - - expected = tf.convert_to_tensor( - [ - [-tf.cos(2 * x) * tf.cos(2 * y), tf.sin(2 * x) * tf.sin(2 * y)], - [tf.sin(2 * x) * tf.sin(2 * y), -2 * tf.cos(x) ** 2 * tf.cos(2 * y)], - ] - ) - assert np.allclose(hess, expected, atol=tol, rtol=0) - - def test_max_diff(self, tol): - """Test that setting the max_diff parameter blocks higher-order - derivatives""" - dev = DefaultQubit() - params = tf.Variable([0.543, -0.654]) - - def cost_fn(x): - ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) - return result[0] + result[1][0] - - with tf.GradientTape() as jac_tape: - with tf.GradientTape() as grad_tape: - res = cost_fn(params) - grad = grad_tape.gradient(res, params) - hess = jac_tape.gradient(grad, params) - - x, y = params - expected = 0.5 * (3 + tf.cos(x) ** 2 * tf.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - expected = tf.convert_to_tensor( - [-tf.cos(x) * tf.cos(2 * y) * tf.sin(x), -tf.cos(x) ** 2 * tf.sin(2 * y)] - ) - assert np.allclose(grad, expected, atol=tol, rtol=0) - assert hess is None - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -@pytest.mark.usefixtures("use_legacy_and_new_opmath") -class TestHamiltonianWorkflows: - """Test that tapes ending with expectations - of Hamiltonians provide correct results and gradients""" - - @pytest.fixture - def cost_fn(self, execute_kwargs, shots, device): - """Cost function for gradient tests""" - - def _cost_fn(weights, coeffs1, coeffs2): - obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] - H1 = qml.Hamiltonian(coeffs1, obs1) - if qml.operation.active_new_opmath(): - H1 = qml.pauli.pauli_sentence(H1).operation() - - obs2 = [qml.PauliZ(0)] - H2 = qml.Hamiltonian(coeffs2, obs2) - if qml.operation.active_new_opmath(): - H2 = qml.pauli.pauli_sentence(H2).operation() - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H1) - qml.expval(H2) - - tape = qml.tape.QuantumScript.from_queue(q, shots=shots) - return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - - return _cost_fn - - @staticmethod - def cost_fn_expected(weights, coeffs1, coeffs2): - """Analytic value of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return [-c * tf.sin(x) * tf.sin(y) + tf.cos(x) * (a + b * tf.sin(y)), d * tf.cos(x)] - - @staticmethod - def cost_fn_jacobian(weights, coeffs1, coeffs2): - """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1 - d = coeffs2[0] - x, y = weights - return tf.convert_to_tensor( - [ - [ - -c * tf.cos(x) * tf.sin(y) - tf.sin(x) * (a + b * tf.sin(y)), - b * tf.cos(x) * tf.cos(y) - c * tf.cos(y) * tf.sin(x), - tf.cos(x), - tf.cos(x) * tf.sin(y), - -(tf.sin(x) * tf.sin(y)), - 0, - ], - [-d * tf.sin(x), 0, 0, 0, 0, tf.cos(x)], - ] - ) - - def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with no trainable parameters.""" - - if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): - pytest.skip("adjoint differentiation does not suppport hamiltonians.") - - device_vjp = execute_kwargs.get("device_vjp", False) - - coeffs1 = tf.constant([0.1, 0.2, 0.3], dtype=tf.float64) - coeffs2 = tf.constant([0.7], dtype=tf.float64) - weights = tf.Variable([0.4, 0.5], dtype=tf.float64) - - with tf.GradientTape(persistent=device_vjp) as tape: - res = cost_fn(weights, coeffs1, coeffs2) - - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac = tape.jacobian(res, [weights], experimental_use_pfor=not device_vjp) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0) - - def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, shots): - """Test hamiltonian with trainable parameters.""" - if execute_kwargs["gradient_fn"] == "adjoint": - pytest.skip("trainable hamiltonians not supported with adjoint") - if qml.operation.active_new_opmath(): - pytest.skip("parameter shift derivatives do not yet support sums.") - - coeffs1 = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) - coeffs2 = tf.Variable([0.7], dtype=tf.float64) - weights = tf.Variable([0.4, 0.5], dtype=tf.float64) - - with tf.GradientTape() as tape: - res = cost_fn(weights, coeffs1, coeffs2) - - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - jac = qml.math.hstack(tape.jacobian(res, [weights, coeffs1, coeffs2]), like="tensorflow") - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0) - - -@pytest.mark.parametrize("diff_method", ("adjoint", "parameter-shift")) -def test_device_returns_float32(diff_method): - """Test that if the device returns float32, the derivative succeeds.""" - - def _to_float32(results): - if isinstance(results, (list, tuple)): - return tuple(_to_float32(r) for r in results) - return np.array(results, dtype=np.float32) - - class Float32Dev(qml.devices.DefaultQubit): - def execute(self, circuits, execution_config=qml.devices.DefaultExecutionConfig): - results = super().execute(circuits, execution_config) - return _to_float32(results) - - dev = Float32Dev() - - @qml.qnode(dev, diff_method=diff_method) - def circuit(x): - qml.RX(tf.cos(x), wires=0) - return qml.expval(qml.Z(0)) - - x = tf.Variable(0.1, dtype=tf.float64) - - with tf.GradientTape() as tape: - y = circuit(x) - - assert qml.math.allclose(y, np.cos(np.cos(0.1))) - - g = tape.gradient(y, x) - expected_g = np.sin(np.cos(0.1)) * np.sin(0.1) - assert qml.math.allclose(g, expected_g) diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py b/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py deleted file mode 100644 index 3cdcf5eae30..00000000000 --- a/tests/interfaces/default_qubit_2_integration/test_torch_default_qubit_2.py +++ /dev/null @@ -1,849 +0,0 @@ -# Copyright 2018-2023 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. -"""Torch specific tests for execute and default qubit 2.""" -import numpy as np -import pytest -from param_shift_dev import ParamShiftDerivativesDevice - -import pennylane as qml -from pennylane import execute -from pennylane.devices import DefaultQubit -from pennylane.gradients import param_shift -from pennylane.measurements import Shots - -torch = pytest.importorskip("torch") - -pytestmark = pytest.mark.torch - - -@pytest.fixture(autouse=True) -def run_before_and_after_tests(): - torch.set_default_dtype(torch.float64) - yield - torch.set_default_dtype(torch.float32) - - -# pylint: disable=too-few-public-methods -class TestCaching: - """Tests for caching behaviour""" - - @pytest.mark.skip("caching is not implemented for torch") - @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum - when computing Hessians.""" - dev = DefaultQubit() - params = torch.arange(1, num_params + 1, requires_grad=True, dtype=torch.float64) - - N = len(params) - - def get_cost_tape(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - - for i in range(2, num_params): - qml.RZ(x[i], wires=[i % 2]) - - qml.CNOT(wires=[0, 1]) - qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - - return qml.tape.QuantumScript.from_queue(q) - - def cost_no_cache(x): - return qml.execute( - [get_cost_tape(x)], - dev, - gradient_fn=qml.gradients.param_shift, - cache=False, - max_diff=2, - )[0] - - def cost_cache(x): - return qml.execute( - [get_cost_tape(x)], - dev, - gradient_fn=qml.gradients.param_shift, - cache=True, - max_diff=2, - )[0] - - # No caching: number of executions is not ideal - with qml.Tracker(dev) as tracker: - hess1 = torch.autograd.functional.hessian(cost_no_cache, params) - - if num_params == 2: - # compare to theoretical result - x, y, *_ = params.clone().detach() - expected = torch.tensor( - [ - [2 * torch.cos(2 * x) * torch.sin(y) ** 2, torch.sin(2 * x) * torch.sin(2 * y)], - [ - torch.sin(2 * x) * torch.sin(2 * y), - -2 * torch.cos(x) ** 2 * torch.cos(2 * y), - ], - ] - ) - assert torch.allclose(expected, hess1) - - expected_runs = 1 # forward pass - - # Jacobian of an involutory observable: - # ------------------------------------ - # - # 2 * N execs: evaluate the analytic derivative of - # 1 execs: Get , the expectation value of the tape with unshifted parameters. - num_shifted_evals = 2 * N - runs_for_jacobian = num_shifted_evals + 1 - expected_runs += runs_for_jacobian - - # Each tape used to compute the Jacobian is then shifted again - expected_runs += runs_for_jacobian * num_shifted_evals - assert tracker.totals["executions"] == expected_runs - - # Use caching: number of executions is ideal - - with qml.Tracker(dev) as tracker2: - hess2 = torch.autograd.functional.hessian(cost_cache, params) - assert torch.allclose(hess1, hess2) - - expected_runs_ideal = 1 # forward pass - expected_runs_ideal += 2 * N # Jacobian - expected_runs_ideal += N + 1 # Hessian diagonal - expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert tracker2.totals["executions"] == expected_runs_ideal - assert expected_runs_ideal < expected_runs - - -# add tests for lightning 2 when possible -# set rng for device when possible -test_matrix = [ - ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), - ({"gradient_fn": param_shift}, Shots((100000, 100000)), DefaultQubit(seed=42)), - ({"gradient_fn": param_shift}, Shots(None), DefaultQubit()), - ({"gradient_fn": "backprop"}, Shots(None), DefaultQubit()), - ( - {"gradient_fn": "adjoint", "grad_on_execution": True, "device_vjp": False}, - Shots(None), - DefaultQubit(), - ), - ( - { - "gradient_fn": "adjoint", - "grad_on_execution": False, - "device_vjp": False, - }, - Shots(None), - DefaultQubit(), - ), - ({"gradient_fn": "adjoint", "device_vjp": True}, Shots(None), DefaultQubit()), - ( - {"gradient_fn": "device", "device_vjp": False}, - Shots((100000, 100000)), - ParamShiftDerivativesDevice(), - ), - ( - {"gradient_fn": "device", "device_vjp": True}, - Shots((100000, 100000)), - ParamShiftDerivativesDevice(), - ), -] - - -def atol_for_shots(shots): - """Return higher tolerance if finite shots.""" - return 1e-2 if shots else 1e-6 - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -class TestTorchExecuteIntegration: - """Test the torch interface execute function - integrates well for both forward and backward execution""" - - def test_execution(self, execute_kwargs, shots, device): - """Test execution""" - - def cost(a, b): - ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] - tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - - ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] - tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - - return execute([tape1, tape2], device, **execute_kwargs) - - a = torch.tensor(0.1, requires_grad=True) - b = torch.tensor(0.2, requires_grad=False) - - with device.tracker: - res = cost(a, b) - - if execute_kwargs.get("grad_on_execution", False): - assert device.tracker.totals["execute_and_derivative_batches"] == 1 - else: - assert device.tracker.totals["batches"] == 1 - assert device.tracker.totals["executions"] == 2 # different wires so different hashes - - assert len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - exp = torch.cos(a) * torch.cos(b) - if shots.has_partitioned_shots: - for shot in range(2): - for wire in range(2): - assert qml.math.allclose(res[shot][wire], exp, atol=atol_for_shots(shots)) - else: - for wire in range(2): - assert qml.math.allclose(res[wire], exp, atol=atol_for_shots(shots)) - - def test_scalar_jacobian(self, execute_kwargs, shots, device): - """Test scalar jacobian calculation""" - a = torch.tensor(0.1, requires_grad=True) - - def cost(a): - tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - res = torch.autograd.functional.jacobian(cost, a) - if not shots.has_partitioned_shots: - assert res.shape == () # pylint: disable=no-member - - # compare to standard tape jacobian - tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) - tape.trainable_params = [0] - tapes, fn = param_shift(tape) - expected = fn(device.execute(tapes)) - - assert expected.shape == () - if shots.has_partitioned_shots: - for i in range(shots.num_copies): - assert torch.allclose(res[i], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[i], -torch.sin(a), atol=atol_for_shots(shots)) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res, -torch.sin(a), atol=atol_for_shots(shots)) - - def test_jacobian(self, execute_kwargs, shots, device): - """Test jacobian calculation""" - a = torch.tensor(0.1, requires_grad=True) - b = torch.tensor(0.2, requires_grad=True) - - def cost(a, b): - ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - [res] = execute([tape], device, **execute_kwargs) - if shots.has_partitioned_shots: - return torch.hstack(res[0] + res[1]) - return torch.hstack(res) - - res = cost(a, b) - expected = torch.tensor([torch.cos(a), -torch.cos(a) * torch.sin(b)]) - if shots.has_partitioned_shots: - assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(cost, (a, b)) - assert isinstance(res, tuple) and len(res) == 2 - - expected = ( - torch.tensor([-torch.sin(a), torch.sin(a) * torch.sin(b)]), - torch.tensor([0, -torch.cos(a) * torch.cos(b)]), - ) - if shots.has_partitioned_shots: - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - for _r, _e in zip(res, expected): - assert torch.allclose(_r[:2], _e, atol=atol_for_shots(shots)) - assert torch.allclose(_r[2:], _e, atol=atol_for_shots(shots)) - - else: - assert res[0].shape == (2,) - assert res[1].shape == (2,) - - for _r, _e in zip(res, expected): - assert torch.allclose(_r, _e, atol=atol_for_shots(shots)) - - def test_tape_no_parameters(self, execute_kwargs, shots, device): - """Test that a tape with no parameters is correctly - ignored during the gradient computation""" - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), wires=0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape4 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), 0)], - [qml.probs(wires=[0, 1])], - shots=shots, - ) - res = execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) - if shots.has_partitioned_shots: - res = [qml.math.asarray(ri, like="torch") for r in res for ri in r] - else: - res = [qml.math.asarray(r, like="torch") for r in res] - return sum(torch.hstack(res)) - - params = torch.tensor([0.1, 0.2], requires_grad=True) - x, y = params.clone().detach() - - res = cost(params) - expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) - - if shots.has_partitioned_shots: - assert torch.allclose(res, 2 * expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res.backward() - expected = torch.tensor([-torch.cos(y) * torch.sin(x), -torch.cos(x) * torch.sin(y)]) - if shots.has_partitioned_shots: - assert torch.allclose(params.grad, 2 * expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(params.grad, expected, atol=atol_for_shots(shots), rtol=0) - - @pytest.mark.skip("torch cannot reuse tensors in various computations") - def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): - """Test that tapes wit different can be executed and differentiated.""" - - if execute_kwargs["gradient_fn"] == "backprop": - pytest.xfail("backprop is not compatible with something about this situation.") - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - shots=shots, - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - shots=shots, - ) - res = execute([tape1, tape2, tape3], device, **execute_kwargs) - return torch.hstack([qml.math.asarray(r, like="torch") for r in res]) - - params = torch.tensor([0.1, 0.2], requires_grad=True) - x, y = params.clone().detach() - - res = cost(params) - assert isinstance(res, torch.Tensor) - assert res.shape == (4,) - - assert torch.allclose(res[0], torch.cos(x) * torch.cos(y), atol=atol_for_shots(shots)) - assert torch.allclose(res[1], torch.tensor(1.0), atol=atol_for_shots(shots)) - assert torch.allclose(res[2], torch.cos(torch.tensor(0.5)), atol=atol_for_shots(shots)) - assert torch.allclose(res[3], torch.cos(x) * torch.cos(y), atol=atol_for_shots(shots)) - - jac = torch.autograd.functional.jacobian(cost, params) - assert isinstance(jac, torch.Tensor) - assert jac.shape == (4, 2) # pylint: disable=no-member - - assert torch.allclose(jac[1:3], torch.tensor(0.0), atol=atol_for_shots(shots)) - - d1 = -torch.sin(x) * torch.cos(y) - assert torch.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) # fails for torch - assert torch.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - - d2 = -torch.cos(x) * torch.sin(y) - assert torch.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) # fails for torch - assert torch.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - - def test_reusing_quantum_tape(self, execute_kwargs, shots, device): - """Test re-using a quantum tape by passing new parameters""" - a = torch.tensor(0.1, requires_grad=True) - b = torch.tensor(0.2, requires_grad=True) - - tape = qml.tape.QuantumScript( - [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], - ) - assert tape.trainable_params == [0, 1] - - def cost(a, b): - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return torch.hstack(execute([new_tape], device, **execute_kwargs)[0]) - - jac = torch.autograd.functional.jacobian(cost, (a, b)) - - a = torch.tensor(0.54, requires_grad=True) - b = torch.tensor(0.8, requires_grad=True) - - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - res2 = cost(2 * a, b) - expected = torch.tensor([torch.cos(2 * a), -torch.cos(2 * a) * torch.sin(b)]) - assert torch.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - - jac = torch.autograd.functional.jacobian(lambda a, b: cost(2 * a, b), (a, b)) - expected = ( - torch.tensor([-2 * torch.sin(2 * a), 2 * torch.sin(2 * a) * torch.sin(b)]), - torch.tensor([0, -torch.cos(2 * a) * torch.cos(b)]), - ) - assert isinstance(jac, tuple) and len(jac) == 2 - for _j, _e in zip(jac, expected): - assert torch.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - - def test_classical_processing(self, execute_kwargs, device, shots): - """Test classical processing within the quantum tape""" - a = torch.tensor(0.1, requires_grad=True) - b = torch.tensor(0.2, requires_grad=False) - c = torch.tensor(0.3, requires_grad=True) - - def cost(a, b, c): - ops = [ - qml.RY(a * c, wires=0), - qml.RZ(b, wires=0), - qml.RX(c + c**2 + torch.sin(a), wires=0), - ] - - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) - return execute([tape], device, **execute_kwargs)[0] - - # PyTorch docs suggest a lambda for cost functions with some non-trainable args - # See for more: https://pytorch.org/docs/stable/autograd.html#functional-higher-level-api - res = torch.autograd.functional.jacobian(lambda _a, _c: cost(_a, b, _c), (a, c)) - - # Only two arguments are trainable - assert isinstance(res, tuple) and len(res) == 2 - if not shots.has_partitioned_shots: - assert res[0].shape == () - assert res[1].shape == () - - # I tried getting analytic results for this circuit but I kept being wrong and am giving up - - @pytest.mark.skip("torch handles gradients and jacobians differently") - def test_no_trainable_parameters(self, execute_kwargs, shots, device): - """Test evaluation and Jacobian if there are no trainable parameters""" - a = torch.tensor(0.1, requires_grad=False) - b = torch.tensor(0.2, requires_grad=False) - - def cost(a, b): - ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] - tape = qml.tape.QuantumScript(ops, m, shots=shots) - return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - - res = cost(a, b) - assert res.shape == (2,) - - res = torch.autograd.functional.jacobian(cost, (a, b)) - assert len(res) == 0 - - def loss(a, b): - return torch.sum(cost(a, b)) - - res = loss(a, b) - res.backward() - - assert torch.allclose(torch.tensor([a.grad, b.grad]), 0) - - def test_matrix_parameter(self, execute_kwargs, device, shots): - """Test that the torch interface works correctly - with a matrix parameter""" - U = torch.tensor([[0, 1], [1, 0]], requires_grad=False, dtype=torch.float64) - a = torch.tensor(0.1, requires_grad=True) - - def cost(a, U): - ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] - tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) - return execute([tape], device, **execute_kwargs)[0] - - res = cost(a, U) - assert torch.allclose(res, -torch.cos(a), atol=atol_for_shots(shots)) - - jac = torch.autograd.functional.jacobian(lambda y: cost(y, U), a) - assert isinstance(jac, torch.Tensor) - assert torch.allclose(jac, torch.sin(a), atol=atol_for_shots(shots), rtol=0) - - def test_differentiable_expand(self, execute_kwargs, device, shots): - """Test that operation and nested tapes expansion - is differentiable""" - - class U3(qml.U3): - """Dummy operator.""" - - def decomposition(self): - theta, phi, lam = self.data - wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] - - def cost_fn(a, p): - tape = qml.tape.QuantumScript( - [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] - ) - gradient_fn = execute_kwargs["gradient_fn"] - if gradient_fn is None: - _gradient_method = None - elif isinstance(gradient_fn, str): - _gradient_method = gradient_fn - else: - _gradient_method = "gradient-transform" - config = qml.devices.ExecutionConfig( - interface="autograd", - gradient_method=_gradient_method, - grad_on_execution=execute_kwargs.get("grad_on_execution", None), - ) - program, _ = device.preprocess(execution_config=config) - return execute([tape], device, **execute_kwargs, transform_program=program)[0] - - a = torch.tensor(0.1, requires_grad=False) - p = torch.tensor([0.1, 0.2, 0.3], requires_grad=True) - - res = cost_fn(a, p) - expected = torch.cos(a) * torch.cos(p[1]) * torch.sin(p[0]) + torch.sin(a) * ( - torch.cos(p[2]) * torch.sin(p[1]) + torch.cos(p[0]) * torch.cos(p[1]) * torch.sin(p[2]) - ) - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(lambda _p: cost_fn(a, _p), p) - expected = torch.tensor( - [ - torch.cos(p[1]) - * ( - torch.cos(a) * torch.cos(p[0]) - - torch.sin(a) * torch.sin(p[0]) * torch.sin(p[2]) - ), - torch.cos(p[1]) * torch.cos(p[2]) * torch.sin(a) - - torch.sin(p[1]) - * ( - torch.cos(a) * torch.sin(p[0]) - + torch.cos(p[0]) * torch.sin(a) * torch.sin(p[2]) - ), - torch.sin(a) - * ( - torch.cos(p[0]) * torch.cos(p[1]) * torch.cos(p[2]) - - torch.sin(p[1]) * torch.sin(p[2]) - ), - ] - ) - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_probability_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob outputs""" - - def cost(x, y): - ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.probs(wires=0), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = torch.tensor(0.543, requires_grad=True) - y = torch.tensor(-0.654, requires_grad=True) - - res = cost(x, y) - expected = torch.tensor( - [ - [ - torch.cos(x / 2) ** 2, - torch.sin(x / 2) ** 2, - (1 + torch.cos(x) * torch.cos(y)) / 2, - (1 - torch.cos(x) * torch.cos(y)) / 2, - ], - ] - ) - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(cost, (x, y)) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (4,) - assert res[1].shape == (4,) - - expected = ( - torch.tensor( - [ - [ - -torch.sin(x) / 2, - torch.sin(x) / 2, - -torch.sin(x) * torch.cos(y) / 2, - torch.sin(x) * torch.cos(y) / 2, - ], - ] - ), - torch.tensor( - [ - [0, 0, -torch.cos(x) * torch.sin(y) / 2, torch.cos(x) * torch.sin(y) / 2], - ] - ), - ) - - assert torch.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - def test_ragged_differentiation(self, execute_kwargs, device, shots): - """Tests correct output shape and evaluation for a tape - with prob and expval outputs""" - - def cost(x, y): - ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] - m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] - tape = qml.tape.QuantumScript(ops, m) - return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - - x = torch.tensor(0.543, requires_grad=True) - y = torch.tensor(-0.654, requires_grad=True) - - res = cost(x, y) - expected = torch.tensor( - [ - torch.cos(x), - (1 + torch.cos(x) * torch.cos(y)) / 2, - (1 - torch.cos(x) * torch.cos(y)) / 2, - ] - ) - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(cost, (x, y)) - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (3,) - assert res[1].shape == (3,) - - expected = ( - torch.tensor( - [-torch.sin(x), -torch.sin(x) * torch.cos(y) / 2, torch.sin(x) * torch.cos(y) / 2] - ), - torch.tensor([0, -torch.cos(x) * torch.sin(y) / 2, torch.cos(x) * torch.sin(y) / 2]), - ) - assert torch.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - - -class TestHigherOrderDerivatives: - """Test that the torch execute function can be differentiated""" - - @pytest.mark.parametrize( - "params", - [ - torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64), - torch.tensor([0, -0.654], requires_grad=True, dtype=torch.float64), - torch.tensor([-2.0, 0], requires_grad=True, dtype=torch.float64), - ], - ) - def test_parameter_shift_hessian(self, params, tol): - """Tests that the output of the parameter-shift transform - can be differentiated using torch, yielding second derivatives.""" - dev = DefaultQubit() - - def cost_fn(x): - ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params.clone().detach() - expected = 0.5 * (3 + torch.cos(x) ** 2 * torch.cos(2 * y)) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - res.backward() - expected = torch.tensor( - [-torch.cos(x) * torch.cos(2 * y) * torch.sin(x), -torch.cos(x) ** 2 * torch.sin(2 * y)] - ) - assert torch.allclose(params.grad, expected, atol=tol, rtol=0) - - res = torch.autograd.functional.hessian(cost_fn, params) - expected = torch.tensor( - [ - [-torch.cos(2 * x) * torch.cos(2 * y), torch.sin(2 * x) * torch.sin(2 * y)], - [torch.sin(2 * x) * torch.sin(2 * y), -2 * torch.cos(x) ** 2 * torch.cos(2 * y)], - ] - ) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - def test_max_diff(self, tol): - """Test that setting the max_diff parameter blocks higher-order - derivatives""" - dev = DefaultQubit() - params = torch.tensor([0.543, -0.654], requires_grad=True) - - def cost_fn(x): - ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] - tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - - ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] - tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - - result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) - return result[0] + result[1][0] - - res = cost_fn(params) - x, y = params.clone().detach() - expected = 0.5 * (3 + torch.cos(x) ** 2 * torch.cos(2 * y)) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - res.backward() - expected = torch.tensor( - [-torch.cos(x) * torch.cos(2 * y) * torch.sin(x), -torch.cos(x) ** 2 * torch.sin(2 * y)] - ) - assert torch.allclose(params.grad, expected, atol=tol, rtol=0) - - res = torch.autograd.functional.hessian(cost_fn, params) - expected = torch.zeros([2, 2]) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) -@pytest.mark.usefixtures("use_legacy_and_new_opmath") -class TestHamiltonianWorkflows: - """Test that tapes ending with expectations - of Hamiltonians provide correct results and gradients""" - - @pytest.fixture - def cost_fn(self, execute_kwargs, shots, device): - """Cost function for gradient tests""" - - def _cost_fn(weights, coeffs1, coeffs2): - obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] - H1 = qml.Hamiltonian(coeffs1, obs1) - if qml.operation.active_new_opmath(): - H1 = qml.pauli.pauli_sentence(H1).operation() - - obs2 = [qml.PauliZ(0)] - H2 = qml.Hamiltonian(coeffs2, obs2) - if qml.operation.active_new_opmath(): - H2 = qml.pauli.pauli_sentence(H2).operation() - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(weights[0], wires=0) - qml.RY(weights[1], wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H1) - qml.expval(H2) - - tape = qml.tape.QuantumScript.from_queue(q, shots=shots) - res = execute([tape], device, **execute_kwargs)[0] - if shots.has_partitioned_shots: - return torch.hstack(res[0] + res[1]) - return torch.hstack(res) - - return _cost_fn - - @staticmethod - def cost_fn_expected(weights, coeffs1, coeffs2): - """Analytic value of cost_fn above""" - a, b, c = coeffs1.clone().detach() - d = coeffs2[0].clone().detach() - x, y = weights.clone().detach() - return torch.tensor( - [ - -c * torch.sin(x) * torch.sin(y) + torch.cos(x) * (a + b * torch.sin(y)), - d * torch.cos(x), - ] - ) - - @staticmethod - def cost_fn_jacobian(weights, coeffs1, coeffs2): - """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1.clone().detach() - d = coeffs2[0].clone().detach() - x, y = weights.clone().detach() - return torch.tensor( - [ - [ - -c * torch.cos(x) * torch.sin(y) - torch.sin(x) * (a + b * torch.sin(y)), - b * torch.cos(x) * torch.cos(y) - c * torch.cos(y) * torch.sin(x), - torch.cos(x), - torch.cos(x) * torch.sin(y), - -(torch.sin(x) * torch.sin(y)), - 0, - ], - [-d * torch.sin(x), 0, 0, 0, 0, torch.cos(x)], - ] - ) - - def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with no trainable parameters.""" - - if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): - pytest.skip("adjoint differentiation does not suppport hamiltonians.") - - coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False) - coeffs2 = torch.tensor([0.7], requires_grad=False) - weights = torch.tensor([0.4, 0.5], requires_grad=True) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.autograd.functional.jacobian(lambda w: cost_fn(w, coeffs1, coeffs2), weights) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - if shots.has_partitioned_shots: - assert torch.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): - """Test hamiltonian with trainable parameters.""" - if execute_kwargs["gradient_fn"] == "adjoint": - pytest.skip("trainable hamiltonians not supported with adjoint") - if qml.operation.active_new_opmath(): - pytest.skip("parameter shift derivatives do not yet support sums.") - - coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True) - coeffs2 = torch.tensor([0.7], requires_grad=True) - weights = torch.tensor([0.4, 0.5], requires_grad=True) - - res = cost_fn(weights, coeffs1, coeffs2) - expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - - res = torch.hstack(torch.autograd.functional.jacobian(cost_fn, (weights, coeffs1, coeffs2))) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - if shots.has_partitioned_shots: - pytest.xfail( - "multiple hamiltonians with shot vectors does not seem to be differentiable." - ) - else: - assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/legacy_devices_integration/test_autograd_legacy.py b/tests/interfaces/legacy_devices_integration/test_autograd_legacy.py new file mode 100644 index 00000000000..311c2f32b21 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_autograd_legacy.py @@ -0,0 +1,1552 @@ +# Copyright 2018-2021 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. +"""Unit tests for the autograd interface""" +# pylint: disable=protected-access,too-few-public-methods +import sys + +import autograd +import pytest + +import pennylane as qml +from pennylane import numpy as np +from pennylane.devices import DefaultQubitLegacy +from pennylane.gradients import finite_diff, param_shift +from pennylane.operation import AnyWires, Observable + +pytestmark = pytest.mark.autograd + + +class TestAutogradExecuteUnitTests: + """Unit tests for autograd execution""" + + def test_import_error(self, mocker): + """Test that an exception is caught on import error""" + + mock = mocker.patch.object(autograd.extend, "defvjp") + mock.side_effect = ImportError() + + try: + del sys.modules["pennylane.workflow.interfaces.autograd"] + except KeyError: + pass + + dev = qml.device("default.qubit.legacy", wires=2, shots=None) + + with qml.queuing.AnnotatedQueue() as q: + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + with pytest.raises( + qml.QuantumFunctionError, + match="autograd not found. Please install the latest version " + "of autograd to enable the 'autograd' interface", + ): + qml.execute([tape], dev, gradient_fn=param_shift, interface="autograd") + + def test_jacobian_options(self, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = np.array([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute( + [tape], + device, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + )[0] + + qml.jacobian(cost)(a, device=dev) + + for args in spy.call_args_list: + assert args[1]["shifts"] == [(np.pi / 4,)] * 2 + + def test_incorrect_grad_on_execution(self): + """Test that an error is raised if a gradient transform + is used with grad_on_execution=True""" + a = np.array([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, gradient_fn=param_shift, grad_on_execution=True)[0] + + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + qml.jacobian(cost)(a, device=dev) + + def test_unknown_interface(self): + """Test that an error is raised if the interface is unknown""" + a = np.array([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, gradient_fn=param_shift, interface="None")[0] + + with pytest.raises(ValueError, match="interface must be in"): + cost(a, device=dev) + + def test_grad_on_execution(self, mocker): + """Test that grad on execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(dev, "execute_and_gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + )[0] + + a = np.array([0.1, 0.2], requires_grad=True) + cost(a) + + # adjoint method only performs a single device execution, but gets both result and gradient + assert dev.num_executions == 1 + spy.assert_called() + + def test_no_gradients_on_execution(self, mocker): + """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + a = np.array([0.1, 0.2], requires_grad=True) + cost(a) + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + qml.jacobian(cost)(a) + spy_gradients.assert_called() + + +class TestBatchTransformExecution: + """Tests to ensure batch transforms can be correctly executed + via qml.execute and batch_transform""" + + def test_no_batch_transform(self, mocker): + """Test that batch transforms can be disabled and enabled""" + dev = qml.device("default.qubit.legacy", wires=2, shots=100000) + + H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0) + x = 0.6 + y = 0.2 + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=0) + qml.RY(y, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(H) + + tape = qml.tape.QuantumScript.from_queue(q) + spy = mocker.spy(dev, "batch_transform") + + if not qml.operation.active_new_opmath(): + with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"): + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + _ = qml.execute([tape], dev, None, device_batch_transform=False) + else: + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + res = qml.execute([tape], dev, None, device_batch_transform=False) + assert np.allclose(res[0], np.cos(y), atol=0.1) + + spy.assert_not_called() + + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + res = qml.execute([tape], dev, None, device_batch_transform=True) + spy.assert_called() + + assert qml.math.shape(res[0]) == () + assert np.allclose(res[0], np.cos(y), rtol=0.05) + + def test_batch_transform_dynamic_shots(self): + """Tests that the batch transform considers the number of shots for the execution, not those + statically on the device.""" + dev = qml.device("default.qubit.legacy", wires=1) + H = 2.0 * qml.PauliZ(0) + qscript = qml.tape.QuantumScript(measurements=[qml.expval(H)]) + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + res = qml.execute([qscript], dev, interface=None, override_shots=10) + assert res == (2.0,) + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cachesize): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, gradient_fn=param_shift, cachesize=cachesize)[0] + + params = np.array([0.1, 0.2]) + qml.jacobian(cost)(params, cachesize=2) + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache)[0] + + custom_cache = {} + params = np.array([0.1, 0.2]) + qml.jacobian(cost)(params, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_caching_param_shift(self, tol): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum.""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache)[0] + + # Without caching, 5 jacobians should still be performed + params = np.array([0.1, 0.2]) + qml.jacobian(cost)(params, cache=None) + assert dev.num_executions == 5 + + # With caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + jac_fn = qml.jacobian(cost) + grad1 = jac_fn(params, cache=True) + assert dev.num_executions == 5 + + # Check that calling the cost function again + # continues to evaluate the device (that is, the cache + # is emptied between calls) + grad2 = jac_fn(params, cache=True) + assert dev.num_executions == 10 + assert np.allclose(grad1, grad2, atol=tol, rtol=0) + + # Check that calling the cost function again + # with different parameters produces a different Jacobian + grad2 = jac_fn(2 * params, cache=True) + assert dev.num_executions == 15 + assert not np.allclose(grad1, grad2, atol=tol, rtol=0) + + @pytest.mark.parametrize("num_params", [2, 3]) + def test_caching_param_shift_hessian(self, num_params, tol): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum + when computing Hessians.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = np.arange(1, num_params + 1) / 10 + + N = len(params) + + def cost(x, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + + for i in range(2, num_params): + qml.RZ(x[i], wires=[i % 2]) + + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache, max_diff=2)[0] + + # No caching: number of executions is not ideal + hess1 = qml.jacobian(qml.grad(cost))(params, cache=False) + + if num_params == 2: + # compare to theoretical result + x, y, *_ = params + expected = np.array( + [ + [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(expected, hess1, atol=tol, rtol=0) + + expected_runs = 1 # forward pass + + # Jacobian of an involutory observable: + # ------------------------------------ + # + # 2 * N execs: evaluate the analytic derivative of + # 1 execs: Get , the expectation value of the tape with unshifted parameters. + num_shifted_evals = 2 * N + runs_for_jacobian = num_shifted_evals + 1 + expected_runs += runs_for_jacobian + + # Each tape used to compute the Jacobian is then shifted again + expected_runs += runs_for_jacobian * num_shifted_evals + assert dev.num_executions == expected_runs + + # Use caching: number of executions is ideal + dev._num_executions = 0 + hess2 = qml.jacobian(qml.grad(cost))(params, cache=True) + assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + expected_runs_ideal = 1 # forward pass + expected_runs_ideal += 2 * N # Jacobian + expected_runs_ideal += N + 1 # Hessian diagonal + expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal + assert dev.num_executions == expected_runs_ideal + assert expected_runs_ideal < expected_runs + + def test_caching_adjoint_no_grad_on_execution(self): + """Test that caching reduces the number of adjoint evaluations + when the grads is not on execution.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = np.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack( + qml.execute( + [tape], + dev, + gradient_fn="device", + cache=cache, + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + ) + + # no caching, but jac for each batch still stored. + qml.jacobian(cost)(params, cache=None) + assert dev.num_executions == 2 + + # With caching, only 2 evaluations are required. One + # for the forward pass, and one for the backward pass. + dev._num_executions = 0 + jac_fn = qml.jacobian(cost) + jac_fn(params, cache=True) + assert dev.num_executions == 2 + + def test_single_backward_pass_batch(self): + """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift + is requested for multiple tapes.""" + + dev = qml.device("default.qubit.legacy", wires=2) + + def f(x): + tape1 = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.probs(wires=0)]) + tape2 = qml.tape.QuantumScript([qml.RY(x, 0)], [qml.probs(wires=0)]) + + results = qml.execute([tape1, tape2], dev, gradient_fn=qml.gradients.param_shift) + return results[0] + results[1] + + x = qml.numpy.array(0.1) + with dev.tracker: + out = qml.jacobian(f)(x) + + assert dev.tracker.totals["batches"] == 2 + assert dev.tracker.history["batch_len"] == [2, 4] + expected = [-2 * np.cos(x / 2) * np.sin(x / 2), 2 * np.sin(x / 2) * np.cos(x / 2)] + assert qml.math.allclose(out, expected) + + def test_single_backward_pass_split_hamiltonian(self): + """Tests that the backward pass is one single batch, not a bunch of batches, when parameter + shift derivatives are requested for a tape that the device split into batches.""" + + dev = qml.device("default.qubit.legacy", wires=2, shots=50000) + + H = qml.Hamiltonian([1, 1], [qml.PauliY(0), qml.PauliZ(0)], grouping_type="qwc") + + def f(x): + tape = qml.tape.QuantumScript([qml.RX(x, wires=0)], [qml.expval(H)]) + return qml.execute([tape], dev, gradient_fn=qml.gradients.param_shift)[0] + + x = qml.numpy.array(0.1) + with dev.tracker: + out = qml.grad(f)(x) + + assert dev.tracker.totals["batches"] == 2 + assert dev.tracker.history["batch_len"] == [2, 4] + + assert qml.math.allclose(out, -np.cos(x) - np.sin(x), atol=0.05) + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + }, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestAutogradExecuteIntegration: + """Test the autograd interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, execute_kwargs): + """Test execution""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, b): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + return qml.execute([tape1, tape2], dev, **execute_kwargs) + + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=False) + res = cost(a, b) + + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_scalar_jacobian(self, execute_kwargs, tol): + """Test scalar jacobian calculation""" + a = np.array(0.1, requires_grad=True) + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], dev, **execute_kwargs)[0] + + res = qml.jacobian(cost)(a) + assert res.shape == () + + # compare to standard tape jacobian + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(dev.batch_execute(tapes)) + + assert expected.shape == () + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_jacobian(self, execute_kwargs, tol): + """Test jacobian calculation""" + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + def cost(a, b, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost(a, b, device=dev) + expected = [np.cos(a), -np.cos(a) * np.sin(b)] + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.jacobian(cost)(a, b, device=dev) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (2,) + assert res[1].shape == (2,) + + expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) + assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected)) + + @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") + def test_tape_no_parameters(self, execute_kwargs, tol): + """Test that a tape with no parameters is correctly + ignored during the gradient computation""" + + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(params): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.expval(qml.PauliX(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(np.array(0.5, requires_grad=False), wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + with qml.queuing.AnnotatedQueue() as q3: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape3 = qml.tape.QuantumScript.from_queue(q3) + with qml.queuing.AnnotatedQueue() as q4: + qml.RY(np.array(0.5, requires_grad=False), wires=0) + qml.probs(wires=[0, 1]) + + tape4 = qml.tape.QuantumScript.from_queue(q4) + return sum( + autograd.numpy.hstack( + qml.execute([tape1, tape2, tape3, tape4], dev, **execute_kwargs) + ) + ) + + params = np.array([0.1, 0.2], requires_grad=True) + x, y = params + + res = cost(params) + expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) + assert np.allclose(res, expected, atol=tol, rtol=0) + + grad = qml.grad(cost)(params) + expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)] + assert np.allclose(grad, expected, atol=tol, rtol=0) + + @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") + def test_tapes_with_different_return_size(self, execute_kwargs): + """Test that tapes wit different can be executed and differentiated.""" + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(params): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(np.array(0.5, requires_grad=False), wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + with qml.queuing.AnnotatedQueue() as q3: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape3 = qml.tape.QuantumScript.from_queue(q3) + return autograd.numpy.hstack(qml.execute([tape1, tape2, tape3], dev, **execute_kwargs)) + + params = np.array([0.1, 0.2], requires_grad=True) + + res = cost(params) + assert isinstance(res, np.ndarray) + assert res.shape == (4,) + + jac = qml.jacobian(cost)(params) + assert isinstance(jac, np.ndarray) + assert jac.shape == (4, 2) + + def test_reusing_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + assert tape.trainable_params == [0, 1] + + def cost(a, b): + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return autograd.numpy.hstack(qml.execute([new_tape], dev, **execute_kwargs)[0]) + + jac_fn = qml.jacobian(cost) + jac = jac_fn(a, b) + + a = np.array(0.54, requires_grad=True) + b = np.array(0.8, requires_grad=True) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = [np.cos(2 * a), -np.cos(2 * a) * np.sin(b)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac_fn = qml.jacobian(lambda a, b: cost(2 * a, b)) + jac = jac_fn(a, b) + expected = ( + [-2 * np.sin(2 * a), 2 * np.sin(2 * a) * np.sin(b)], + [0, -np.cos(2 * a) * np.cos(b)], + ) + assert isinstance(jac, tuple) and len(jac) == 2 + assert all(np.allclose(_j, _e, atol=tol, rtol=0) for _j, _e in zip(jac, expected)) + + def test_classical_processing(self, execute_kwargs): + """Test classical processing within the quantum tape""" + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=False) + c = np.array(0.3, requires_grad=True) + + def cost(a, b, c, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a * c, wires=0) + qml.RZ(b, wires=0) + qml.RX(c + c**2 + np.sin(a), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = qml.jacobian(cost)(a, b, c, device=dev) + + # Only two arguments are trainable + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_no_trainable_parameters(self, execute_kwargs): + """Test evaluation and Jacobian if there are no trainable parameters""" + a = np.array(0.1, requires_grad=False) + b = np.array(0.2, requires_grad=False) + + def cost(a, b, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + + dev = qml.device("default.qubit.legacy", wires=2) + res = cost(a, b, device=dev) + assert res.shape == (2,) + + with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): + res = qml.jacobian(cost)(a, b, device=dev) + assert len(res) == 0 + + def loss(a, b): + return np.sum(cost(a, b, device=dev)) + + with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): + res = qml.grad(loss)(a, b) + + assert np.allclose(res, 0) + + def test_matrix_parameter(self, execute_kwargs, tol): + """Test that the autograd interface works correctly + with a matrix parameter""" + U = np.array([[0, 1], [1, 0]], requires_grad=False) + a = np.array(0.1, requires_grad=True) + + def cost(a, U, device): + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = cost(a, U, device=dev) + assert isinstance(res, np.ndarray) + assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + + jac_fn = qml.jacobian(cost) + jac = jac_fn(a, U, device=dev) + assert isinstance(jac, np.ndarray) + assert np.allclose(jac, np.sin(a), atol=tol, rtol=0) + + def test_differentiable_expand(self, execute_kwargs, tol): + """Test that operation and nested tapes expansion + is differentiable""" + + class U3(qml.U3): + def decomposition(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + def cost_fn(a, p, device): + with qml.queuing.AnnotatedQueue() as q_tape: + qml.RX(a, wires=0) + U3(*p, wires=0) + qml.expval(qml.PauliX(0)) + + tape = qml.tape.QuantumScript.from_queue(q_tape) + tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name)) + return qml.execute([tape], device, **execute_kwargs)[0] + + a = np.array(0.1, requires_grad=False) + p = np.array([0.1, 0.2, 0.3], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + res = cost_fn(a, p, device=dev) + expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( + np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = qml.jacobian(cost_fn) + res = jac_fn(a, p, device=dev) + expected = np.array( + [ + np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), + np.cos(p[1]) * np.cos(p[2]) * np.sin(a) + - np.sin(p[1]) + * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), + np.sin(a) + * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_probability_differentiation(self, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob outputs""" + + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + def cost(x, y, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + + dev = qml.device("default.qubit.legacy", wires=2) + x = np.array(0.543, requires_grad=True) + y = np.array(-0.654, requires_grad=True) + + res = cost(x, y, device=dev) + expected = np.array( + [ + [ + np.cos(x / 2) ** 2, + np.sin(x / 2) ** 2, + (1 + np.cos(x) * np.cos(y)) / 2, + (1 - np.cos(x) * np.cos(y)) / 2, + ], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = qml.jacobian(cost) + res = jac_fn(x, y, device=dev) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (4,) + assert res[1].shape == (4,) + + expected = ( + np.array( + [ + [ + -np.sin(x) / 2, + np.sin(x) / 2, + -np.sin(x) * np.cos(y) / 2, + np.sin(x) * np.cos(y) / 2, + ], + ] + ), + np.array( + [ + [0, 0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], + ] + ), + ) + + assert np.allclose(res[0], expected[0], atol=tol, rtol=0) + assert np.allclose(res[1], expected[1], atol=tol, rtol=0) + + def test_ragged_differentiation(self, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + def cost(x, y, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + + dev = qml.device("default.qubit.legacy", wires=2) + x = np.array(0.543, requires_grad=True) + y = np.array(-0.654, requires_grad=True) + + res = cost(x, y, device=dev) + expected = np.array( + [np.cos(x), (1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = qml.jacobian(cost) + res = jac_fn(x, y, device=dev) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (3,) + assert res[1].shape == (3,) + + expected = ( + np.array([-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]), + np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]), + ) + assert np.allclose(res[0], expected[0], atol=tol, rtol=0) + assert np.allclose(res[1], expected[1], atol=tol, rtol=0) + + def test_sampling(self, execute_kwargs): + """Test sampling works as expected""" + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + + def cost(device): + with qml.queuing.AnnotatedQueue() as q: + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.sample(qml.PauliZ(0)) + qml.sample(qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute([tape], device, **execute_kwargs)[0] + + shots = 10 + dev = qml.device("default.qubit.legacy", wires=2, shots=shots) + res = cost(device=dev) + assert isinstance(res, tuple) + assert len(res) == 2 + assert res[0].shape == (shots,) + assert res[1].shape == (shots,) + + +class TestHigherOrderDerivatives: + """Test that the autograd execute function can be differentiated""" + + @pytest.mark.parametrize( + "params", + [ + np.array([0.543, -0.654], requires_grad=True), + np.array([0, -0.654], requires_grad=True), + np.array([-2.0, 0], requires_grad=True), + ], + ) + def test_parameter_shift_hessian(self, params, tol): + """Tests that the output of the parameter-shift transform + can be differentiated using autograd, yielding second derivatives.""" + dev = qml.device("default.qubit.autograd", wires=2) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[0], wires=0) + qml.RY(x[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = qml.execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) + return result[0] + result[1][0] + + res = cost_fn(params) + x, y = params + expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.grad(cost_fn)(params) + expected = np.array( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.jacobian(qml.grad(cost_fn))(params) + expected = np.array( + [ + [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_adjoint_hessian_one_param(self, tol): + """Since the adjoint hessian is not a differentiable transform, + higher-order derivatives are not supported.""" + dev = qml.device("default.qubit.autograd", wires=2) + params = np.array([0.543, -0.654], requires_grad=True) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return qml.execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + )[0] + + with pytest.warns(UserWarning, match="Output seems independent"): + res = qml.jacobian(qml.grad(cost_fn))(params) + + assert np.allclose(res, np.zeros([2, 2]), atol=tol, rtol=0) + + def test_adjoint_hessian_multiple_params(self, tol): + """Since the adjoint hessian is not a differentiable transform, + higher-order derivatives are not supported.""" + dev = qml.device("default.qubit.autograd", wires=2) + params = np.array([0.543, -0.654], requires_grad=True) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack( + qml.execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + )[0] + ) + + with pytest.warns(UserWarning, match="Output seems independent"): + res = qml.jacobian(qml.jacobian(cost_fn))(params) + + assert np.allclose(res, np.zeros([2, 2, 2]), atol=tol, rtol=0) + + def test_max_diff(self, tol): + """Test that setting the max_diff parameter blocks higher-order + derivatives""" + dev = qml.device("default.qubit.legacy", wires=2) + params = np.array([0.543, -0.654], requires_grad=True) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[0], wires=0) + qml.RY(x[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = qml.execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) + return result[0] + result[1][0] + + res = cost_fn(params) + x, y = params + expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.grad(cost_fn)(params) + expected = np.array( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + with pytest.warns(UserWarning, match="Output seems independent"): + res = qml.jacobian(qml.grad(cost_fn))(params) + + expected = np.zeros([2, 2]) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +class TestOverridingShots: + """Test overriding shots on execution""" + + def test_changing_shots(self, mocker, tol): + """Test that changing shots works on execution""" + dev = qml.device("default.qubit.legacy", wires=2, shots=None) + a, b = np.array([0.543, -0.654], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + spy = mocker.spy(dev, "sample") + + # execute with device default shots (None) + res = qml.execute([tape], dev, gradient_fn=param_shift) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() + + # execute with shots=100 + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + res = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100) + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = qml.execute([tape], dev, gradient_fn=param_shift) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # same single call from above, no additional calls + + def test_overriding_shots_with_same_value(self, mocker): + """Overriding shots with the same value as the device will have no effect""" + dev = qml.device("default.qubit.legacy", wires=2, shots=123) + a, b = np.array([0.543, -0.654], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + spy = mocker.Mock(wraps=qml.Device.shots.fset) + # pylint:disable=assignment-from-no-return,too-many-function-args + mock_property = qml.Device.shots.setter(spy) + mocker.patch.object(qml.Device, "shots", mock_property) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + qml.execute([tape], dev, gradient_fn=param_shift, override_shots=123) + # overriden shots is the same, no change + spy.assert_not_called() + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100) + # overriden shots is not the same, shots were changed + spy.assert_called() + + # shots were temporarily set to the overriden value + assert spy.call_args_list[0][0] == (dev, 100) + # shots were then returned to the built-in value + assert spy.call_args_list[1][0] == (dev, 123) + + def test_overriding_device_with_shot_vector(self): + """Overriding a device that has a batch of shots set + results in original shots being returned after execution""" + dev = qml.device("default.qubit.legacy", wires=2, shots=[10, (1, 3), 5]) + + assert dev.shots == 18 + assert dev._shot_vector == [(10, 1), (1, 3), (5, 1)] + + a, b = np.array([0.543, -0.654], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + res = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100)[0] + + assert isinstance(res, np.ndarray) + assert res.shape == () + + # device is unchanged + assert dev.shots == 18 + assert dev._shot_vector == [(10, 1), (1, 3), (5, 1)] + + res = qml.execute([tape], dev, gradient_fn=param_shift)[0] + assert len(res) == 5 + + @pytest.mark.xfail(reason="Shots vector must be adapted for new return types.") + def test_gradient_integration(self): + """Test that temporarily setting the shots works + for gradient computations""" + # TODO: Update here when shot vectors are supported + dev = qml.device("default.qubit.legacy", wires=2, shots=None) + a, b = np.array([0.543, -0.654], requires_grad=True) + + def cost_fn(a, b, shots): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + result = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=shots) + return result[0] + + res = qml.jacobian(cost_fn)(a, b, shots=[10000, 10000, 10000]) + assert dev.shots is None + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (3,) + assert res[1].shape == (3,) + + expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] + assert all( + np.allclose(np.mean(r, axis=0), e, atol=0.1, rtol=0) for r, e in zip(res, expected) + ) + + +execute_kwargs_hamiltonian = [ + {"gradient_fn": param_shift}, + {"gradient_fn": finite_diff}, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +class TestHamiltonianWorkflows: + """Test that tapes ending with expectations + of Hamiltonians provide correct results and gradients""" + + @pytest.fixture + def cost_fn(self, execute_kwargs): + """Cost function for gradient tests""" + + def _cost_fn(weights, coeffs1, coeffs2, dev=None): + obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] + H1 = qml.Hamiltonian(coeffs1, obs1) + + obs2 = [qml.PauliZ(0)] + H2 = qml.Hamiltonian(coeffs2, obs2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(weights[0], wires=0) + qml.RY(weights[1], wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(H1) + qml.expval(H2) + + tape = qml.tape.QuantumScript.from_queue(q) + return autograd.numpy.hstack(qml.execute([tape], dev, **execute_kwargs)[0]) + + return _cost_fn + + @staticmethod + def cost_fn_expected(weights, coeffs1, coeffs2): + """Analytic value of cost_fn above""" + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + + @staticmethod + def cost_fn_jacobian(weights, coeffs1, coeffs2): + """Analytic jacobian of cost_fn above""" + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return np.array( + [ + [ + -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), + b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), + np.cos(x), + np.cos(x) * np.sin(y), + -(np.sin(x) * np.sin(y)), + 0, + ], + [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + ] + ) + + def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): + """Test hamiltonian with no trainable parameters.""" + # pylint: disable=unused-argument + coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False) + coeffs2 = np.array([0.7], requires_grad=False) + weights = np.array([0.4, 0.5], requires_grad=True) + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): + """Test hamiltonian with trainable parameters.""" + # pylint: disable=unused-argument + coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True) + coeffs2 = np.array([0.7], requires_grad=True) + weights = np.array([0.4, 0.5], requires_grad=True) + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev)) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +class TestCustomJacobian: + """Test for custom Jacobian.""" + + def test_custom_jacobians(self): + """Test custom Jacobian device methood""" + + class CustomJacobianDevice(DefaultQubitLegacy): + @classmethod + def capabilities(cls): + capabilities = super().capabilities() + capabilities["provides_jacobian"] = True + return capabilities + + def jacobian(self, tape): + # pylint: disable=unused-argument + return np.array([1.0, 2.0, 3.0, 4.0]) + + dev = CustomJacobianDevice(wires=2) + + @qml.qnode(dev, diff_method="device") + def circuit(v): + qml.RX(v, wires=0) + return qml.probs(wires=[0, 1]) + + d_circuit = qml.jacobian(circuit, argnum=0) + + params = np.array(1.0, requires_grad=True) + + d_out = d_circuit(params) + assert np.allclose(d_out, np.array([1.0, 2.0, 3.0, 4.0])) + + def test_custom_jacobians_param_shift(self): + """Test computing the gradient using the parameter-shift + rule with a device that provides a jacobian""" + + class MyQubit(DefaultQubitLegacy): + @classmethod + def capabilities(cls): + capabilities = super().capabilities().copy() + capabilities.update( + provides_jacobian=True, + ) + return capabilities + + def jacobian(self, *args, **kwargs): + raise NotImplementedError() + + dev = MyQubit(wires=2) + + @qml.qnode(dev, diff_method="parameter-shift", grad_on_execution=False) + def qnode(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) + + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + def cost(a, b): + return autograd.numpy.hstack(qnode(a, b)) + + res = qml.jacobian(cost)(a, b) + expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) + + assert np.allclose(res[0], expected[0]) + assert np.allclose(res[1], expected[1]) + + +class SpecialObject: + """SpecialObject + A special object that conveniently encapsulates the return value of + a special observable supported by a special device and which supports + multiplication with scalars and addition. + """ + + def __init__(self, val): + self.val = val + + def __mul__(self, other): + new = SpecialObject(self.val) + new *= other + return new + + def __imul__(self, other): + self.val *= other + return self + + def __rmul__(self, other): + return self * other + + def __iadd__(self, other): + self.val += other.val if isinstance(other, self.__class__) else other + return self + + def __add__(self, other): + new = SpecialObject(self.val) + new += other.val if isinstance(other, self.__class__) else other + return new + + def __radd__(self, other): + return self + other + + +class SpecialObservable(Observable): + """SpecialObservable""" + + num_wires = AnyWires + num_params = 0 + par_domain = None + + def diagonalizing_gates(self): + """Diagonalizing gates""" + return [] + + +class DeviceSupportingSpecialObservable(DefaultQubitLegacy): + name = "Device supporting SpecialObservable" + short_name = "default.qubit.specialobservable" + observables = DefaultQubitLegacy.observables.union({"SpecialObservable"}) + + @staticmethod + def _asarray(arr, dtype=None): + # pylint: disable=unused-argument + return arr + + @classmethod + def capabilities(cls): + capabilities = super().capabilities().copy() + capabilities.update( + provides_jacobian=True, + ) + return capabilities + + def expval(self, observable, **kwargs): + if self.analytic and isinstance(observable, SpecialObservable): + val = super().expval(qml.PauliZ(wires=0), **kwargs) + return np.array(SpecialObject(val)) + + return super().expval(observable, **kwargs) + + def jacobian(self, tape): + # we actually let pennylane do the work of computing the + # jacobian for us but return it as a device jacobian + gradient_tapes, fn = qml.gradients.param_shift(tape) + tape_jacobian = fn(qml.execute(gradient_tapes, self, None)) + return tape_jacobian + + +@pytest.mark.autograd +class TestObservableWithObjectReturnType: + """Unit tests for qnode returning a custom object""" + + def test_custom_return_type(self): + """Test custom return values for a qnode""" + + dev = DeviceSupportingSpecialObservable(wires=1, shots=None) + + # force diff_method='parameter-shift' because otherwise + # PennyLane swaps out dev for default.qubit.autograd + @qml.qnode(dev, diff_method="parameter-shift") + def qnode(x): + qml.RY(x, wires=0) + return qml.expval(SpecialObservable(wires=0)) + + @qml.qnode(dev, diff_method="parameter-shift") + def reference_qnode(x): + qml.RY(x, wires=0) + return qml.expval(qml.PauliZ(wires=0)) + + out = qnode(0.2) + assert isinstance(out, np.ndarray) + assert isinstance(out.item(), SpecialObject) + assert np.isclose(out.item().val, reference_qnode(0.2)) + + def test_jacobian_with_custom_return_type(self): + """Test differentiation of a QNode on a device supporting a + special observable that returns an object rather than a number.""" + + dev = DeviceSupportingSpecialObservable(wires=1, shots=None) + + # force diff_method='parameter-shift' because otherwise + # PennyLane swaps out dev for default.qubit.autograd + @qml.qnode(dev, diff_method="parameter-shift") + def qnode(x): + qml.RY(x, wires=0) + return qml.expval(SpecialObservable(wires=0)) + + @qml.qnode(dev, diff_method="parameter-shift") + def reference_qnode(x): + qml.RY(x, wires=0) + return qml.expval(qml.PauliZ(wires=0)) + + reference_jac = (qml.jacobian(reference_qnode)(np.array(0.2, requires_grad=True)),) + + assert np.isclose( + reference_jac, + qml.jacobian(qnode)(np.array(0.2, requires_grad=True)).item().val, + ) + + # now check that also the device jacobian works with a custom return type + @qml.qnode(dev, diff_method="device") + def device_gradient_qnode(x): + qml.RY(x, wires=0) + return qml.expval(SpecialObservable(wires=0)) + + assert np.isclose( + reference_jac, + qml.jacobian(device_gradient_qnode)(np.array(0.2, requires_grad=True)).item().val, + ) diff --git a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py similarity index 69% rename from tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py index a0d1e0f2e66..818d7d58ff5 100644 --- a/tests/interfaces/default_qubit_2_integration/test_autograd_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2020 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. @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the autograd interface with a QNode""" -# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg, use-dict-literal, no-name-in-module - +# pylint: disable=too-many-arguments,too-few-public-methods, use-dict-literal, use-implicit-booleaness-not-comparison, +# pylint: disable=unnecessary-lambda-assignment import autograd import autograd.numpy as anp import pytest @@ -21,47 +21,32 @@ import pennylane as qml from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit -from tests.param_shift_dev import ParamShiftDerivativesDevice -# dev, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - [qml.device("default.qubit"), "finite-diff", False, False], - [qml.device("default.qubit"), "parameter-shift", False, False], - [qml.device("default.qubit"), "backprop", True, False], - [qml.device("default.qubit"), "adjoint", True, False], - [qml.device("default.qubit"), "adjoint", False, False], - [qml.device("default.qubit"), "spsa", False, False], - [qml.device("default.qubit"), "hadamard", False, False], - [ParamShiftDerivativesDevice(), "parameter-shift", False, False], - [ParamShiftDerivativesDevice(), "best", False, False], - [ParamShiftDerivativesDevice(), "parameter-shift", True, False], - [ParamShiftDerivativesDevice(), "parameter-shift", False, True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] interface_qubit_device_and_diff_method = [ - ["autograd", DefaultQubit(), "finite-diff", False, False], - ["autograd", DefaultQubit(), "parameter-shift", False, False], - ["autograd", DefaultQubit(), "backprop", True, False], - ["autograd", DefaultQubit(), "adjoint", True, False], - ["autograd", DefaultQubit(), "adjoint", False, False], - ["autograd", DefaultQubit(), "adjoint", True, True], - ["autograd", DefaultQubit(), "adjoint", False, True], - ["autograd", DefaultQubit(), "spsa", False, False], - ["autograd", DefaultQubit(), "hadamard", False, False], - ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", False, True], - ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", True, True], - ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", False, False], - ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", True, False], - ["auto", DefaultQubit(), "finite-diff", False, False], - ["auto", DefaultQubit(), "parameter-shift", False, False], - ["auto", DefaultQubit(), "backprop", True, False], - ["auto", DefaultQubit(), "adjoint", True, False], - ["auto", DefaultQubit(), "adjoint", False, False], - ["auto", DefaultQubit(), "spsa", False, False], - ["auto", DefaultQubit(), "hadamard", False, False], - ["auto", qml.device("lightning.qubit", wires=5), "adjoint", False, False], - ["auto", qml.device("lightning.qubit", wires=5), "adjoint", True, False], + ["autograd", "default.qubit.legacy", "finite-diff", False], + ["autograd", "default.qubit.legacy", "parameter-shift", False], + ["autograd", "default.qubit.legacy", "backprop", True], + ["autograd", "default.qubit.legacy", "adjoint", True], + ["autograd", "default.qubit.legacy", "adjoint", False], + ["autograd", "default.qubit.legacy", "spsa", False], + ["autograd", "default.qubit.legacy", "hadamard", False], + ["auto", "default.qubit.legacy", "finite-diff", False], + ["auto", "default.qubit.legacy", "parameter-shift", False], + ["auto", "default.qubit.legacy", "backprop", True], + ["auto", "default.qubit.legacy", "adjoint", True], + ["auto", "default.qubit.legacy", "adjoint", False], + ["auto", "default.qubit.legacy", "spsa", False], + ["auto", "default.qubit.legacy", "hadamard", False], ] pytestmark = pytest.mark.autograd @@ -72,21 +57,78 @@ @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with Autograd integrates with the PennyLane stack""" - # pylint:disable=unused-argument - def test_execution_no_interface( - self, interface, dev, diff_method, grad_on_execution, device_vjp + # pylint: disable=unused-argument + + def test_nondiff_param_unwrapping( + self, interface, dev_name, diff_method, grad_on_execution, mocker ): + """Test that non-differentiable parameters are correctly unwrapped + to NumPy ndarrays or floats (if 0-dimensional)""" + if diff_method != "parameter-shift": + pytest.skip("Test only supports parameter-shift") + + dev = qml.device("default.qubit.legacy", wires=1) + + @qnode(dev, interface=interface, diff_method=diff_method) + def circuit(x, y): + qml.RX(x[0], wires=0) + qml.Rot(*x[1:], wires=0) + qml.RY(y[0], wires=0) + return qml.expval(qml.PauliZ(0)) + + x = np.array([0.1, 0.2, 0.3, 0.4], requires_grad=False) + y = np.array([0.5], requires_grad=True) + + param_data = [] + + def mock_apply(*args, **kwargs): + for op in args[0]: + param_data.extend(op.data) + + mocker.patch.object(dev, "apply", side_effect=mock_apply) + circuit(x, y) + assert param_data == [0.1, 0.2, 0.3, 0.4, 0.5] + assert not any(isinstance(p, np.tensor) for p in param_data) + + # test the jacobian works correctly + param_data = [] + qml.grad(circuit)(x, y) + assert param_data == [ + 0.1, + 0.2, + 0.3, + 0.4, + 0.5, + 0.1, + 0.2, + 0.3, + 0.4, + 0.5 + np.pi / 2, + 0.1, + 0.2, + 0.3, + 0.4, + 0.5 - np.pi / 2, + ] + assert not any(isinstance(p, np.tensor) for p in param_data) + + def test_execution_no_interface(self, interface, dev_name, diff_method, grad_on_execution): """Test execution works without an interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - @qnode(dev, interface=None) + num_wires = 1 + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface=None, diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -96,24 +138,26 @@ def circuit(a): res = circuit(a) - # without the interface, the QNode simply returns a scalar array or float - assert isinstance(res, (np.ndarray, float)) - assert qml.math.shape(res) == tuple() # pylint: disable=comparison-with-callable + # without the interface, the QNode simply returns a scalar array + assert isinstance(res, np.ndarray) + assert res.shape == tuple() - def test_execution_with_interface( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): + # gradients should cause an error + with pytest.raises(TypeError, match="must be real number, not ArrayBox"): + qml.grad(circuit)(a) + + def test_execution_with_interface(self, interface, dev_name, diff_method, grad_on_execution): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 1 + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface=interface, diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -121,27 +165,31 @@ def circuit(a): a = np.array(0.1, requires_grad=True) assert circuit.interface == interface + # gradients should work grad = qml.grad(circuit)(a) + assert isinstance(grad, float) assert grad.shape == tuple() - def test_jacobian(self, interface, dev, diff_method, grad_on_execution, tol, device_vjp): + def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test jacobian calculation""" + num_wires = 2 kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA + elif diff_method == "hadamard": + num_wires = 3 a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -172,23 +220,24 @@ def cost(x, y): assert res[1].shape == (2,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - def test_jacobian_no_evaluate( - self, interface, dev, diff_method, grad_on_execution, tol, device_vjp - ): + def test_jacobian_no_evaluate(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test jacobian calculation when no prior circuit evaluation has been performed""" + num_wires = 2 kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + elif diff_method == "hadamard": + num_wires = 3 a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -214,31 +263,36 @@ def cost(x, y): assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - def test_jacobian_options(self, interface, dev, diff_method, grad_on_execution, device_vjp): + def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execution): """Test setting jacobian options""" - if diff_method != "finite-diff": - pytest.skip("Test only supports finite diff.") + wires = [0] + if diff_method in ["backprop", "adjoint"]: + pytest.skip("Test does not support backprop or adjoint method") + elif diff_method == "finite-diff": + kwargs = {"h": 1e-8, "approx_order": 2} + elif diff_method == "parameter-shift": + kwargs = {"shifts": [(0.1,), (0.2,)]} + elif diff_method == "hadamard": + wires = [0, "aux"] + kwargs = {"aux_wire": qml.wires.Wires("aux"), "device_wires": wires} + else: + kwargs = {} a = np.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface=interface, - h=1e-8, - order=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + dev = qml.device("default.qubit.legacy", wires=wires) + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) + circuit(a) + qml.jacobian(circuit)(a) - def test_changing_trainability( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -247,13 +301,9 @@ def test_changing_trainability( a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - @qnode( - dev, - interface=interface, - diff_method="parameter-shift", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + dev = qml.device("default.qubit.legacy", wires=2) + + @qnode(dev, interface=interface, diff_method=diff_method) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -290,18 +340,21 @@ def loss(a, b): circuit(a, b) assert circuit.qtape.trainable_params == [1] - def test_classical_processing(self, interface, dev, diff_method, grad_on_execution, device_vjp): + def test_classical_processing(self, interface, dev_name, diff_method, grad_on_execution): """Test classical processing within the quantum tape""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=False) c = np.array(0.3, requires_grad=True) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -319,17 +372,12 @@ def circuit(a, b, c): assert res[0].shape == () assert res[1].shape == () - def test_no_trainable_parameters( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): + def test_no_trainable_parameters(self, interface, dev_name, diff_method, grad_on_execution): """Test evaluation and Jacobian if there are no trainable parameters""" + dev = qml.device(dev_name, wires=2) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -348,34 +396,35 @@ def circuit(a, b): assert len(res) == 2 assert isinstance(res, tuple) - def cost(x, y): + def cost0(x, y): return autograd.numpy.hstack(circuit(x, y)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - assert not qml.jacobian(cost)(a, b) + assert not qml.jacobian(cost0)(a, b) - def cost2(a, b): + def cost1(a, b): return np.sum(circuit(a, b)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - grad = qml.grad(cost2)(a, b) + grad = qml.grad(cost1)(a, b) assert grad == tuple() - def test_matrix_parameter( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_matrix_parameter(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the autograd interface works correctly with a matrix parameter""" U = np.array([[0, 1], [1, 0]], requires_grad=False) a = np.array(0.1, requires_grad=True) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -391,18 +440,13 @@ def circuit(U, a): assert np.allclose(res, np.sin(a), atol=tol, rtol=0) def test_gradient_non_differentiable_exception( - self, interface, dev, diff_method, grad_on_execution, device_vjp + self, interface, dev_name, diff_method, grad_on_execution ): """Test that an exception is raised if non-differentiable data is differentiated""" + dev = qml.device(dev_name, wires=2) - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface=interface, diff_method=diff_method) def circuit(data1): qml.templates.AmplitudeEmbedding(data1, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) @@ -413,27 +457,18 @@ def circuit(data1): with pytest.raises(qml.numpy.NonDifferentiableError, match="is non-differentiable"): grad_fn(data1) - def test_differentiable_expand( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_differentiable_expand(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that operation and nested tape expansion is differentiable""" kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 10 + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=20) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA - # pylint: disable=too-few-public-methods class U3(qml.U3): - """Custom U3.""" - def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -442,6 +477,7 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] + dev = qml.device(dev_name, wires=2) a = np.array(0.1, requires_grad=False) p = np.array([0.1, 0.2, 0.3], requires_grad=True) @@ -455,14 +491,14 @@ def circuit(a, p): expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) ) - # assert isinstance(res, np.ndarray) - # assert res.shape == () + assert isinstance(res, np.ndarray) + assert res.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) res = qml.grad(circuit)(a, p) - # assert isinstance(res, np.ndarray) - # assert len(res) == 3 + assert isinstance(res, np.ndarray) + assert len(res) == 3 expected = np.array( [ @@ -481,9 +517,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and remains differentiable.""" - def test_changing_shots(self): + def test_changing_shots(self, mocker, tol): """Test that changing shots works on execution""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = np.array([0.543, -0.654], requires_grad=True) @qnode(dev, diff_method=qml.gradients.param_shift) @@ -491,21 +527,31 @@ def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - res = circuit(a, b) + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 - res = circuit(a, b, shots=100) - assert res.shape == (100, 2) # pylint: disable=comparison-with-callable + res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # same single call performed above @pytest.mark.xfail(reason="Param shift and shot vectors.") def test_gradient_integration(self): """Test that temporarily setting the shots works for gradient computations""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = np.array([0.543, -0.654], requires_grad=True) @qnode(dev, diff_method=qml.gradients.param_shift) @@ -528,57 +574,62 @@ def cost_fn(a, b): def test_update_diff_method(self, mocker): """Test that temporarily setting the shots updates the diff method""" + dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = np.array([0.543, -0.654], requires_grad=True) spy = mocker.spy(qml, "execute") - @qnode(DefaultQubit()) + @qnode(dev) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - assert cost_fn.gradient_fn == "backprop" # gets restored to default - - cost_fn(a, b, shots=100) + cost_fn(a, b) # since we are using finite shots, parameter-shift will # be chosen - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - # if we use the default shots value of None, backprop can now be used - cost_fn(a, b) - assert cost_fn.gradient_fn == "backprop" + # if we set the shots to None, backprop can now be used + cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert cost_fn.gradient_fn == "backprop" + + cost_fn(a, b) + assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_qubit_device_and_diff_method ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" + # pylint: disable=unused-argument + def test_probability_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Tests correct output shape and evaluation for a tape with a single prob output""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measuring probabilities with adjoint.") - kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -599,23 +650,25 @@ def circuit(x, y): assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) def test_multiple_probability_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measuring probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -664,24 +717,22 @@ def cost(x, y): ) assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) - def test_ragged_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_ragged_differentiation(self, interface, dev_name, diff_method, grad_on_execution, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measuring probabilities with adjoint.") - kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + num_wires = 2 + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + elif diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -717,25 +768,22 @@ def cost(x, y): assert np.allclose(res[1], expected[1], atol=tol, rtol=0) def test_ragged_differentiation_variance( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measuring probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") + dev = qml.device(dev_name, wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -756,12 +804,12 @@ def circuit(x, y): assert isinstance(res, tuple) assert len(res) == 2 - # assert isinstance(res[0], np.ndarray) - # assert res[0].shape == () + assert isinstance(res[0], np.ndarray) + assert res[0].shape == () assert np.allclose(res[0], expected_var, atol=tol, rtol=0) - # assert isinstance(res[1], np.ndarray) - # assert res[1].shape == (2,) + assert isinstance(res[1], np.ndarray) + assert res[1].shape == (2,) assert np.allclose(res[1], expected_probs, atol=tol, rtol=0) def cost(x, y): @@ -778,38 +826,33 @@ def cost(x, y): assert isinstance(jac, tuple) assert len(jac) == 2 - # assert isinstance(jac[0], np.ndarray) - # assert jac[0].shape == (3,) + assert isinstance(jac[0], np.ndarray) + assert jac[0].shape == (3,) assert np.allclose(jac[0], expected[0], atol=tol, rtol=0) - # assert isinstance(jac[1], np.ndarray) - # assert jac[1].shape == (3,) + assert isinstance(jac[1], np.ndarray) + assert jac[1].shape == (3,) assert np.allclose(jac[1], expected[1], atol=tol, rtol=0) - def test_chained_qnodes(self, interface, dev, diff_method, grad_on_execution, device_vjp): + def test_chained_qnodes(self, interface, dev_name, diff_method, grad_on_execution): """Test that the gradient of chained QNodes works without error""" + num_wires = 2 - # pylint: disable=too-few-public-methods - class Template(qml.templates.StronglyEntanglingLayers): - """Custom template.""" + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution - ) + @qnode(dev, interface=interface, diff_method=diff_method) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface=interface, diff_method=diff_method) def circuit2(data, weights): qml.templates.AngleEmbedding(data, wires=[0, 1]) Template(weights, wires=[0, 1]) @@ -833,22 +876,19 @@ def cost(w1, w2): assert len(res) == 2 - def test_chained_gradient_value( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_chained_gradient_value(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the returned gradient value for two chained qubit QNodes is correct.""" - kwargs = dict( - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - + kwargs = dict(interface=interface, diff_method=diff_method) if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - dev1 = qml.device("default.qubit") + num_wires = 3 + + if diff_method == "hadamard": + num_wires = 4 + + dev1 = qml.device(dev_name, wires=num_wires) @qnode(dev1, **kwargs) def circuit1(a, b, c): @@ -859,11 +899,9 @@ def circuit1(a, b, c): qml.CNOT(wires=[1, 2]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(2)) - dev2 = dev + dev2 = qml.device("default.qubit.legacy", wires=num_wires) - @qnode( - dev2, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution - ) + @qnode(dev2, interface=interface, diff_method=diff_method) def circuit2(data, weights): qml.RX(data[0], wires=0) qml.RX(data[1], wires=1) @@ -933,20 +971,19 @@ def cost(a, b, c, weights): # to the first parameter of circuit1. assert circuit1.qtape.trainable_params == [1, 2] - def test_second_derivative( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_second_derivative(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test second derivative calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, max_diff=2, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -976,17 +1013,18 @@ def circuit(x): assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1001,8 +1039,8 @@ def circuit(x): expected_res = np.cos(a) * np.cos(b) - # assert isinstance(res, np.ndarray) - # assert res.shape == () + assert isinstance(res, np.ndarray) + assert res.shape == () assert np.allclose(res, expected_res, atol=tol, rtol=0) grad_fn = qml.grad(circuit) @@ -1030,18 +1068,19 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_unused_parameter( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1070,20 +1109,19 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a vector valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1099,7 +1137,7 @@ def circuit(x): expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] assert isinstance(res, np.ndarray) - assert res.shape == (2,) # pylint: disable=comparison-with-callable + assert res.shape == (2,) assert np.allclose(res, expected_res, atol=tol, rtol=0) jac_fn = qml.jacobian(circuit) @@ -1136,18 +1174,19 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1190,18 +1229,19 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(a, b): @@ -1215,7 +1255,7 @@ def circuit(a, b): expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] assert isinstance(res, np.ndarray) - assert res.shape == (2,) # pylint: disable=comparison-with-callable + assert res.shape == (2,) assert np.allclose(res, expected_res, atol=tol, rtol=0) jac_fn = qml.jacobian(circuit) @@ -1232,12 +1272,8 @@ def circuit(a, b): assert g[1].shape == (2,) assert np.allclose(g[1], expected_g[1], atol=tol, rtol=0) - def jac_fn_a(*args): - return jac_fn(*args)[0] - - def jac_fn_b(*args): - return jac_fn(*args)[1] - + jac_fn_a = lambda *args: jac_fn(*args)[0] + jac_fn_b = lambda *args: jac_fn(*args)[1] hess_a = qml.jacobian(jac_fn_a)(a, b) hess_b = qml.jacobian(jac_fn_b)(a, b) assert isinstance(hess_a, tuple) and len(hess_a) == 2 @@ -1259,17 +1295,18 @@ def jac_fn_b(*args): assert np.allclose(hess[0], exp_hess[0], atol=tol, rtol=0) assert np.allclose(hess[1], exp_hess[1], atol=tol, rtol=0) - def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a ragged QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") + dev = qml.device(dev_name, wires=2) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1283,18 +1320,14 @@ def circuit(x): a, b = x - expected_res = [ - np.cos(a) * np.cos(b), - 0.5 + 0.5 * np.cos(a) * np.cos(b), - 0.5 - 0.5 * np.cos(a) * np.cos(b), - ] + cos_prod = np.cos(a) * np.cos(b) + expected_res = (cos_prod, [0.5 + 0.5 * cos_prod, 0.5 - 0.5 * cos_prod]) + res = circuit(x) + assert all(qml.math.allclose(r, e) for r, e in zip(res, expected_res)) def cost_fn(x): return autograd.numpy.hstack(circuit(x)) - res = cost_fn(x) - assert qml.math.allclose(res, expected_res) - jac_fn = qml.jacobian(cost_fn) hess = qml.jacobian(jac_fn)(x) @@ -1318,21 +1351,18 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_state(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("Lightning does not support state adjoint diff.") + dev = qml.device(dev_name, wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1363,25 +1393,20 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector( - self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_projector(self, state, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the variance of a projector is correctly returned""" - if diff_method == "adjoint": - pytest.skip("adjoint supports either expvals or diagonal measurements.") kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) - - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("Adjoint does not support projectors") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") + dev = qml.device(dev_name, wires=2) P = np.array(state, requires_grad=False) x, y = np.array([0.765, -0.654], requires_grad=True) @@ -1394,8 +1419,8 @@ def circuit(x, y): res = circuit(x, y) expected = 0.25 * np.sin(x / 2) ** 2 * (3 + np.cos(2 * y) + 2 * np.cos(x) * np.sin(y) ** 2) - # assert isinstance(res, np.ndarray) - # assert res.shape == () + assert isinstance(res, np.ndarray) + assert res.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) jac = qml.jacobian(circuit)(x, y) @@ -1419,60 +1444,106 @@ def circuit(x, y): assert np.allclose(jac, expected, atol=tol, rtol=0) - def test_postselection_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") +@pytest.mark.parametrize( + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": 0.05}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +class TestCV: + """Tests for CV integration""" - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + def test_first_order_observable(self, diff_method, kwargs, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + elif diff_method == "hadamard": + pytest.skip("Hadamard gradient does not support variances.") + + r = np.array(0.543, requires_grad=True) + phi = np.array(-0.654, requires_grad=True) + + @qnode(dev, diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + res = circuit(r, phi) + expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = qml.jacobian(circuit)(r, phi) + expected = np.array( + [ + [ + 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, + 2 * np.sinh(2 * r) * np.sin(2 * phi), + ] + ] ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) + assert np.allclose(res, expected, atol=tol, rtol=0) - phi = np.array(1.23, requires_grad=True) - theta = np.array(4.56, requires_grad=True) + def test_second_order_observable(self, diff_method, kwargs, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + elif diff_method == "hadamard": + pytest.skip("Hadamard gradient does not support variances.") - assert np.allclose(circuit(phi, theta), expected_circuit(theta)) + n = np.array(0.12, requires_grad=True) + a = np.array(0.765, requires_grad=True) - gradient = qml.grad(circuit)(phi, theta) - exp_theta_grad = qml.grad(expected_circuit)(theta) - assert np.allclose(gradient, [0.0, exp_theta_grad]) + @qnode(dev, diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) + res = circuit(n, a) + expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) + assert np.allclose(res, expected, atol=tol, rtol=0) -@pytest.mark.parametrize( - "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method -) + # circuit jacobians + res = qml.jacobian(circuit)(n, a) + expected = np.array([[2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]]) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +def test_adjoint_reuse_device_state(mocker): + """Tests that the autograd interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=1) + + @qnode(dev, diff_method="adjoint") + def circ(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + qml.grad(circ, argnum=0)(1.0) + assert circ.device.num_executions == 1 + + spy.assert_called_with(mocker.ANY, use_device_state=True) + + +@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the Autograd interface""" @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev, diff_method, grad_on_execution, max_diff, device_vjp + self, dev_name, diff_method, grad_on_execution, max_diff ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" @@ -1481,22 +1552,20 @@ def test_gradient_expansion_trainable_only( if max_diff == 2 and diff_method == "hadamard": pytest.skip("Max diff > 1 not supported for Hadamard gradient.") - # pylint: disable=too-few-public-methods - class PhaseShift(qml.PhaseShift): - """dummy phase shift.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + dev = qml.device(dev_name, wires=num_wires) + + class PhaseShift(qml.PhaseShift): grad_method = None def decomposition(self): return [qml.RY(3 * self.data[0], wires=self.wires)] - @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - max_diff=max_diff, - device_vjp=device_vjp, - ) + @qnode(dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff) def circuit(x, y): qml.Hadamard(wires=0) PhaseShift(x, wires=0) @@ -1507,11 +1576,11 @@ def circuit(x, y): y = np.array(0.7, requires_grad=False) circuit(x, y) - _ = qml.grad(circuit)(x, y) + qml.grad(circuit)(x, y) @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, max_diff, tol, device_vjp + self, dev_name, diff_method, grad_on_execution, max_diff, tol ): """Test that if there are non-commuting groups and the number of shots is None the first and second order gradients are correctly evaluated""" @@ -1519,16 +1588,15 @@ def test_hamiltonian_expansion_analytic( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, - device_vjp=device_vjp, ) - if diff_method in ["adjoint", "hadamard"]: pytest.skip("The diff method requested does not yet support Hamiltonians") elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 10 tol = TOL_FOR_SPSA + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} + dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1558,11 +1626,7 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if ( - diff_method in ("parameter-shift", "backprop") - and max_diff == 2 - and dev.name != "param_shift.qubit" - ): + if diff_method in ("parameter-shift", "backprop") and max_diff == 2: if diff_method == "backprop": with pytest.warns(UserWarning, match=r"Output seems independent of input."): grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c) @@ -1580,26 +1644,28 @@ def circuit(data, weights, coeffs): @pytest.mark.slow @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, grad_on_execution, max_diff, device_vjp + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, max_diff, mocker ): - """Test that the Hamiltonian is correctly measured if there + """Test that the Hamiltonian is expanded if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.1 + tol = 0.3 if diff_method in ("adjoint", "backprop", "hadamard"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), - "num_directions": 10, - } + gradient_kwargs = dict( + h=H_FOR_SPSA, + sampler_rng=np.random.default_rng(SEED_FOR_SPSA), + num_directions=20, + ) tol = TOL_FOR_SPSA elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} + dev = qml.device(dev_name, wires=3, shots=50000) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1607,7 +1673,6 @@ def test_hamiltonian_finite_shots( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1623,15 +1688,13 @@ def circuit(data, weights, coeffs): c = np.array([-0.6543, 0.24, 0.54], requires_grad=True) # test output - res = circuit(d, w, c, shots=50000) + res = circuit(d, w, c) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) + spy.assert_called() # test gradients - if diff_method in ["finite-diff", "spsa"]: - pytest.skip(f"{diff_method} not compatible") - - grad = qml.grad(circuit)(d, w, c, shots=50000) + grad = qml.grad(circuit)(d, w, c) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1641,11 +1704,12 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2 and dev.name != "param_shift.qubit": - grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c, shots=50000) + if diff_method == "parameter-shift" and max_diff == 2: + with pytest.warns(UserWarning, match=r"Output seems independent of input."): + grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c) assert np.allclose(grad2_c, 0, atol=tol) - grad2_w_c = qml.jacobian(qml.grad(circuit, argnum=1), argnum=2)(d, w, c, shots=50000) + grad2_w_c = qml.jacobian(qml.grad(circuit, argnum=1), argnum=2)(d, w, c) expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ 0, -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), @@ -1659,38 +1723,36 @@ class TestSample: def test_backprop_error(self): """Test that sampling in backpropagation grad_on_execution raises an error""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2) @qnode(dev, diff_method="backprop") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - with pytest.raises( - qml.QuantumFunctionError, match="does not support backprop with requested" - ): - circuit(shots=10) + with pytest.raises(qml.QuantumFunctionError, match="only supported when shots=None"): + circuit(shots=10) # pylint: disable=unexpected-keyword-arg def test_sample_dimension(self): """Test that the sample function outputs samples of the right size""" n_sample = 10 - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample) - @qnode(dev, diff_method=None) + @qnode(dev) def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=n_sample) + res = circuit() assert isinstance(res, tuple) assert len(res) == 2 - assert res[0].shape == (10,) # pylint: disable=comparison-with-callable + assert res[0].shape == (10,) assert isinstance(res[0], np.ndarray) - assert res[1].shape == (10,) # pylint: disable=comparison-with-callable + assert res[1].shape == (10,) assert isinstance(res[1], np.ndarray) def test_sample_combination(self): @@ -1698,7 +1760,7 @@ def test_sample_combination(self): n_sample = 10 - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) @qnode(dev, diff_method="parameter-shift") def circuit(): @@ -1706,29 +1768,29 @@ def circuit(): return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit(shots=n_sample) + result = circuit() assert isinstance(result, tuple) assert len(result) == 3 assert np.array_equal(result[0].shape, (n_sample,)) - assert isinstance(result[1], (float, np.ndarray)) - assert isinstance(result[2], (float, np.ndarray)) + assert isinstance(result[1], np.ndarray) + assert isinstance(result[2], np.ndarray) assert result[0].dtype == np.dtype("float") def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" n_sample = 10 - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) - @qnode(dev, diff_method=None) + @qnode(dev) def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit(shots=n_sample) + result = circuit() assert isinstance(result, np.ndarray) assert np.array_equal(result.shape, (n_sample,)) @@ -1738,44 +1800,44 @@ def test_multi_wire_sample_regular_shape(self): where a rectangular array is expected""" n_sample = 10 - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(dev, diff_method=None) + @qnode(dev) def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit(shots=n_sample) + result = circuit() # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) assert len(result) == 3 - assert result[0].shape == (10,) # pylint: disable=comparison-with-callable + assert result[0].shape == (10,) assert isinstance(result[0], np.ndarray) - assert result[1].shape == (10,) # pylint: disable=comparison-with-callable + assert result[1].shape == (10,) assert isinstance(result[1], np.ndarray) - assert result[2].shape == (10,) # pylint: disable=comparison-with-callable + assert result[2].shape == (10,) assert isinstance(result[2], np.ndarray) -@pytest.mark.parametrize( - "dev,diff_method,grad_on_execution,device_vjp", qubit_device_and_diff_method -) +@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" - def test_grad_single_measurement_param(self, dev, diff_method, grad_on_execution, device_vjp): + # pylint: disable=unused-argument + + def test_grad_single_measurement_param(self, dev_name, diff_method, grad_on_execution): """For one measurement and one param, the gradient is a float.""" + num_wires = 1 - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1785,20 +1847,25 @@ def circuit(a): grad = qml.grad(circuit)(a) - assert isinstance(grad, np.tensor if diff_method == "backprop" else float) + import sys - def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp - ): + python_version = sys.version_info.minor + if diff_method == "backprop" and python_version > 7: + # Since numpy 1.23.0 + assert isinstance(grad, np.ndarray) + else: + assert isinstance(grad, float) + + def test_grad_single_measurement_multiple_param(self, dev_name, diff_method, grad_on_execution): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + num_wires = 1 - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1815,17 +1882,17 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + num_wires = 1 - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1840,18 +1907,21 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1865,18 +1935,21 @@ def circuit(a): assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1896,17 +1969,20 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """For a multi dimensional measurement (probs), check that a single array is returned.""" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1919,17 +1995,20 @@ def circuit(a): assert jac.shape == (4, 2) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """The jacobian of multiple measurements with a single params return an array.""" + num_wires = 2 - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1946,17 +2025,21 @@ def cost(x): assert jac.shape == (5,) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1980,17 +2063,21 @@ def cost(x, y): assert jac[1].shape == (5,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + @qnode(dev, interface="autograd", diff_method=diff_method) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2006,8 +2093,14 @@ def cost(x): assert isinstance(jac, np.ndarray) assert jac.shape == (5, 2) - def test_hessian_expval_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp): + def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_execution): """The hessian of single a measurement with multiple params return a tuple of arrays.""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2015,14 +2108,7 @@ def test_hessian_expval_multiple_params(self, dev, diff_method, grad_on_executio par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2043,24 +2129,22 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (2,) - def test_hessian_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_expval_multiple_param_array(self, dev_name, diff_method, grad_on_execution): """The hessian of single measurement with a multiple params array return a single array.""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2072,8 +2156,9 @@ def circuit(x): assert isinstance(hess, np.ndarray) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp): + def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution): """The hessian of single a measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2083,14 +2168,7 @@ def test_hessian_var_multiple_params(self, dev, diff_method, grad_on_execution, par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2111,25 +2189,18 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (2,) - def test_hessian_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_var_multiple_param_array(self, dev_name, diff_method, grad_on_execution): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") + dev = qml.device(dev_name, wires=2) + params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2141,24 +2212,19 @@ def circuit(x): assert isinstance(hess, np.ndarray) assert hess.shape == (2, 2) - def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_probs_expval_multiple_params(self, dev_name, diff_method, grad_on_execution): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint method does not currently support second-order diff.") par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2183,23 +2249,18 @@ def cost(x, y): assert hess[1].shape == (6,) def test_hessian_expval_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp + self, dev_name, diff_method, grad_on_execution ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint method does not currently support second-order diff.") + dev = qml.device(dev_name, wires=2) + params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2212,12 +2273,11 @@ def cost(x): hess = qml.jacobian(qml.jacobian(cost))(params) assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) # pylint: disable=no-member + assert hess.shape == (3, 2, 2) - def test_hessian_probs_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, grad_on_execution): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2227,14 +2287,7 @@ def test_hessian_probs_var_multiple_params( par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2258,25 +2311,18 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (6,) - def test_hessian_var_multiple_param_array2( - self, dev, diff_method, grad_on_execution, device_vjp - ): + def test_hessian_var_probs_multiple_param_array(self, dev_name, diff_method, grad_on_execution): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") + dev = qml.device(dev_name, wires=2) + params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode( - dev, - interface="autograd", - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2289,14 +2335,15 @@ def cost(x): hess = qml.jacobian(qml.jacobian(cost))(params) assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) # pylint: disable=no-member + assert hess.shape == (3, 2, 2) -def test_no_ops(): +@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) +def test_no_ops(dev_name): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = DefaultQubit() + dev = qml.device(dev_name, wires=1) @qml.qnode(dev, interface="autograd") def circuit(): diff --git a/tests/interfaces/legacy_devices_integration/test_autograd_qnode_shot_vector_legacy.py b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_shot_vector_legacy.py new file mode 100644 index 00000000000..24d1d824ced --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_shot_vector_legacy.py @@ -0,0 +1,658 @@ +# Copyright 2022 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. +"""Integration tests for using the Autograd interface with shot vectors and with a QNode""" +# pylint: disable=too-many-arguments,redefined-outer-name + +import pytest + +import pennylane as qml +from pennylane import numpy as np +from pennylane import qnode + +pytestmark = pytest.mark.autograd + +shots_and_num_copies = [(((5, 2), 1, 10), 4), ((1, 10, (5, 2)), 4)] +shots_and_num_copies_hess = [(((5, 1), 10), 2), ((10, (5, 1)), 2)] + + +kwargs = { + "finite-diff": {"h": 0.05}, + "parameter-shift": {}, + "spsa": {"h": 0.05, "num_directions": 20}, +} + +qubit_device_and_diff_method = [ + ["default.qubit.legacy", "finite-diff"], + ["default.qubit.legacy", "parameter-shift"], + ["default.qubit.legacy", "spsa"], +] + +TOLS = { + "finite-diff": 0.3, + "parameter-shift": 1e-2, + "spsa": 0.3, +} + + +@pytest.fixture +def gradient_kwargs(request): + diff_method = request.node.funcargs["diff_method"] + return kwargs[diff_method] | ( + {"sampler_rng": np.random.default_rng(42)} if diff_method == "spsa" else {} + ) + + +@pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) +class TestReturnWithShotVectors: + """Class to test the shape of the Jacobian/Hessian with different return types and shot vectors.""" + + def test_jac_single_measurement_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For one measurement and one param, the gradient is a float.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = np.array(0.1) + + def cost(a): + return qml.math.stack(circuit(a)) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies,) + + def test_jac_single_measurement_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For one measurement and multiple param, the gradient is a tuple of arrays.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = np.array(0.1) + b = np.array(0.2) + + def cost(a, b): + return qml.math.stack(circuit(a, b)) + + jac = qml.jacobian(cost, argnum=[0, 1])(a, b) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies,) + + def test_jacobian_single_measurement_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For one measurement and multiple param as a single array params, the gradient is an array.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0)) + + a = np.array([0.1, 0.2]) + + def cost(a): + return qml.math.stack(circuit(a)) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 2) + + def test_jacobian_single_measurement_param_probs( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For a multi dimensional measurement (probs), check that a single array is returned with the correct + dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.probs(wires=[0, 1]) + + a = np.array(0.1) + + def cost(a): + return qml.math.stack(circuit(a)) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 4) + + def test_jacobian_single_measurement_probs_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with + the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.probs(wires=[0, 1]) + + a = np.array(0.1) + b = np.array(0.2) + + def cost(a, b): + return qml.math.stack(circuit(a, b)) + + jac = qml.jacobian(cost, argnum=[0, 1])(a, b) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies, 4) + + def test_jacobian_single_measurement_probs_multiple_param_single_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with + the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.probs(wires=[0, 1]) + + a = np.array([0.1, 0.2]) + + def cost(a): + return qml.math.stack(circuit(a)) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 4, 2) + + def test_jacobian_expval_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The gradient of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, max_diff=1, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) + + def cost(x, y): + res = circuit(x, y) + return qml.math.stack([qml.math.stack(r) for r in res]) + + jac = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies, 2) + + def test_jacobian_expval_expval_multiple_params_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) + + a = np.array([0.1, 0.2, 0.3]) + + def cost(a): + res = circuit(a) + return qml.math.stack([qml.math.stack(r) for r in res]) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 2, 3) + + def test_jacobian_var_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) + + def cost(x, y): + res = circuit(x, y) + return qml.math.stack([qml.math.stack(r) for r in res]) + + jac = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies, 2) + + def test_jacobian_var_var_multiple_params_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) + + a = np.array([0.1, 0.2, 0.3]) + + def cost(a): + res = circuit(a) + return qml.math.stack([qml.math.stack(r) for r in res]) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 2, 3) + + def test_jacobian_multiple_measurement_single_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a single params return an array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = np.array(0.1) + + def cost(a): + res = circuit(a) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 5) + + def test_jacobian_multiple_measurement_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + def cost(a, b): + res = circuit(a, b) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + jac = qml.jacobian(cost, argnum=[0, 1])(a, b) + + assert isinstance(jac, tuple) + assert len(jac) == 2 + for j in jac: + assert isinstance(j, np.ndarray) + assert j.shape == (num_copies, 5) + + def test_jacobian_multiple_measurement_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = np.array([0.1, 0.2]) + + def cost(a): + res = circuit(a) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + jac = qml.jacobian(cost)(a) + + assert isinstance(jac, np.ndarray) + assert jac.shape == (num_copies, 5, 2) + + +@pytest.mark.slow +@pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) +class TestReturnShotVectorHessian: + """Class to test the shape of the Hessian with different return types and shot vectors.""" + + def test_hessian_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of a single measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x, y): + def cost2(x, y): + res = circuit(x, y) + return qml.math.stack(res) + + return qml.math.stack(qml.jacobian(cost2, argnum=[0, 1])(x, y)) + + hess = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(hess, tuple) + assert len(hess) == 2 + for h in hess: + assert isinstance(h, np.ndarray) + assert h.shape == (2, num_copies) + + def test_hessian_expval_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = np.array([0.1, 0.2]) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x): + def cost2(x): + res = circuit(x) + return qml.math.stack(res) + + return qml.jacobian(cost2)(x) + + hess = qml.jacobian(cost)(params) + + assert isinstance(hess, np.ndarray) + assert hess.shape == (num_copies, 2, 2) + + def test_hessian_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of a single measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x, y): + def cost2(x, y): + res = circuit(x, y) + return qml.math.stack(res) + + return qml.math.stack(qml.jacobian(cost2, argnum=[0, 1])(x, y)) + + hess = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(hess, tuple) + assert len(hess) == 2 + for h in hess: + assert isinstance(h, np.ndarray) + assert h.shape == (2, num_copies) + + def test_hessian_var_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = np.array([0.1, 0.2]) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x): + def cost2(x): + res = circuit(x) + return qml.math.stack(res) + + return qml.jacobian(cost2)(x) + + hess = qml.jacobian(cost)(params) + + assert isinstance(hess, np.ndarray) + assert hess.shape == (num_copies, 2, 2) + + def test_hessian_probs_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + if diff_method == "spsa": + pytest.skip("SPSA does not support iterated differentiation in Autograd.") + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = np.array(0.1) + par_1 = np.array(0.2) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + def cost(x, y): + def cost2(x, y): + res = circuit(x, y) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + return qml.math.stack(qml.jacobian(cost2, argnum=[0, 1])(x, y)) + + hess = qml.jacobian(cost, argnum=[0, 1])(par_0, par_1) + + assert isinstance(hess, tuple) + assert len(hess) == 2 + for h in hess: + assert isinstance(h, np.ndarray) + assert h.shape == (2, num_copies, 3) + + def test_hessian_expval_probs_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """The hessian of multiple measurements with a multiple param array return a single array.""" + if diff_method == "spsa": + pytest.skip("SPSA does not support iterated differentiation in Autograd.") + + dev = qml.device(dev_name, wires=2, shots=shots) + + params = np.array([0.1, 0.2]) + + @qnode(dev, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + def cost(x): + def cost2(x): + res = circuit(x) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + return qml.jacobian(cost2)(x) + + hess = qml.jacobian(cost)(params) + + assert isinstance(hess, np.ndarray) + assert hess.shape == (num_copies, 3, 2, 2) + + +shots_and_num_copies = [((30000, 28000, 26000), 3), ((30000, (28000, 2)), 3)] + + +@pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) +class TestReturnShotVectorIntegration: + """Tests for the integration of shots with the autograd interface.""" + + def test_single_expectation_value( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """Tests correct output shape and evaluation for a tape + with a single expval output""" + dev = qml.device(dev_name, wires=2, shots=shots) + x = np.array(0.543) + y = np.array(-0.654) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + def cost(x, y): + res = circuit(x, y) + return qml.math.stack(res) + + all_res = qml.jacobian(cost, argnum=[0, 1])(x, y) + + assert isinstance(all_res, tuple) + assert len(all_res) == 2 + + expected = np.array([-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]) + tol = TOLS[diff_method] + + for res, exp in zip(all_res, expected): + assert isinstance(res, np.ndarray) + assert res.shape == (num_copies,) + assert np.allclose(res, exp, atol=tol, rtol=0) + + def test_prob_expectation_values( + self, dev_name, diff_method, gradient_kwargs, shots, num_copies + ): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + dev = qml.device(dev_name, wires=2, shots=shots) + x = np.array(0.543) + y = np.array(-0.654) + + @qnode(dev, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + def cost(x, y): + res = circuit(x, y) + return qml.math.stack([qml.math.hstack(r) for r in res]) + + all_res = qml.jacobian(cost, argnum=[0, 1])(x, y) + + assert isinstance(all_res, tuple) + assert len(all_res) == 2 + + expected = np.array( + [ + [ + -np.sin(x), + -(np.cos(y / 2) ** 2 * np.sin(x)) / 2, + -(np.sin(x) * np.sin(y / 2) ** 2) / 2, + (np.sin(x) * np.sin(y / 2) ** 2) / 2, + (np.cos(y / 2) ** 2 * np.sin(x)) / 2, + ], + [ + 0, + -(np.cos(x / 2) ** 2 * np.sin(y)) / 2, + (np.cos(x / 2) ** 2 * np.sin(y)) / 2, + (np.sin(x / 2) ** 2 * np.sin(y)) / 2, + -(np.sin(x / 2) ** 2 * np.sin(y)) / 2, + ], + ] + ) + + tol = TOLS[diff_method] + + for res, exp in zip(all_res, expected): + assert isinstance(res, np.ndarray) + assert res.shape == (num_copies, 5) + assert np.allclose(res, exp, atol=tol, rtol=0) diff --git a/tests/interfaces/legacy_devices_integration/test_execute_legacy.py b/tests/interfaces/legacy_devices_integration/test_execute_legacy.py new file mode 100644 index 00000000000..e12c5d8ff55 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_execute_legacy.py @@ -0,0 +1,28 @@ +# Copyright 2023 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. +""" +Interface independent tests for qml.execute +""" + +import pytest + +import pennylane as qml + + +def test_old_interface_no_device_jacobian_products(): + """Test that an error is always raised for the old device interface if device jacobian products are requested.""" + dev = qml.device("default.qubit.legacy", wires=2) + tape = qml.tape.QuantumScript([qml.RX(1.0, wires=0)], [qml.expval(qml.PauliZ(0))]) + with pytest.raises(qml.QuantumFunctionError): + qml.execute((tape,), dev, device_vjp=True) diff --git a/tests/interfaces/legacy_devices_integration/test_jax_jit_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_jit_legacy.py new file mode 100644 index 00000000000..c3fd7df5d1a --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_jax_jit_legacy.py @@ -0,0 +1,901 @@ +# Copyright 2018-2021 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. +"""Unit tests for the JAX-JIT interface""" +import numpy as np + +# pylint: disable=protected-access,too-few-public-methods +import pytest + +import pennylane as qml +from pennylane import execute +from pennylane.gradients import param_shift +from pennylane.typing import TensorLike + +pytestmark = pytest.mark.jax + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) + + +class TestJaxExecuteUnitTests: + """Unit tests for jax execution""" + + def test_jacobian_options(self, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + device, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + )[0] + + jax.grad(cost)(a, device=dev) + + for args in spy.call_args_list: + assert args[1]["shifts"] == [(np.pi / 4,)] * 2 + + def test_incorrect_gradients_on_execution(self): + """Test that an error is raised if an gradient transform + is used with grad_on_execution=True""" + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + device, + gradient_fn=param_shift, + grad_on_execution=True, + )[0] + + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + jax.grad(cost)(a, device=dev) + + def test_unknown_interface(self): + """Test that an error is raised if the interface is unknown""" + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + device, + gradient_fn=param_shift, + interface="None", + )[0] + + with pytest.raises(ValueError, match="Unknown interface"): + cost(a, device=dev) + + def test_grad_on_execution(self, mocker): + """Test that grad on execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(dev, "execute_and_gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={ + "method": "adjoint_jacobian", + "use_device_state": True, + }, + )[0] + + a = jax.numpy.array([0.1, 0.2]) + jax.jit(cost)(a) + + # adjoint method only performs a single device execution + # gradients are not requested when we only want the results + assert dev.num_executions == 1 + spy.assert_not_called() + + # when the jacobian is requested, we always calculate it at the same time as the results + jax.grad(jax.jit(cost))(a) + spy.assert_called() + + def test_no_gradients_on_execution(self, mocker): + """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "execute_and_gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + a = jax.numpy.array([0.1, 0.2]) + jax.jit(cost)(a) + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + jax.grad(jax.jit(cost))(a) + spy_gradients.assert_called() + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cachesize): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn=param_shift, + cachesize=cachesize, + )[0] + + params = jax.numpy.array([0.1, 0.2]) + jax.jit(jax.grad(cost), static_argnums=1)(params, cachesize=2) + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn=param_shift, + cache=cache, + )[0] + + custom_cache = {} + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_custom_cache_multiple(self, mocker): + """Test the use of a custom cache object with multiple tapes""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + def cost(a, b, cache): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + res = execute( + [tape1, tape2], + dev, + gradient_fn=param_shift, + cache=cache, + ) + return res[0] + + custom_cache = {} + jax.grad(cost)(a, b, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_caching_param_shift(self, tol): + """Test that, when using parameter-shift transform, + caching produces the optimum number of evaluations.""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn=param_shift, + cache=cache, + )[0] + + # Without caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cache=None) + assert dev.num_executions == 5 + + # With caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + jac_fn = jax.grad(cost) + grad1 = jac_fn(params, cache=True) + assert dev.num_executions == 5 + + # Check that calling the cost function again + # continues to evaluate the device (that is, the cache + # is emptied between calls) + grad2 = jac_fn(params, cache=True) + assert dev.num_executions == 10 + assert np.allclose(grad1, grad2, atol=tol, rtol=0) + + # Check that calling the cost function again + # with different parameters produces a different Jacobian + grad2 = jac_fn(2 * params, cache=True) + assert dev.num_executions == 15 + assert not np.allclose(grad1, grad2, atol=tol, rtol=0) + + def test_caching_adjoint_backward(self): + """Test that caching produces the optimum number of adjoint evaluations + when mode=backward""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + cache=cache, + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + # Without caching, 2 evaluations are required. + # 1 for the forward pass, and one per output dimension + # on the backward pass. + jax.grad(cost)(params, cache=None) + assert dev.num_executions == 2 + + # With caching, also 2 evaluations are required. One + # for the forward pass, and one for the backward pass. + dev._num_executions = 0 + jac_fn = jax.grad(cost) + jac_fn(params, cache=True) + assert dev.num_executions == 2 + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + }, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestJaxExecuteIntegration: + """Test the jax interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, execute_kwargs): + """Test execution""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, b): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return execute([tape1, tape2], dev, **execute_kwargs) + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + res = cost(a, b) + + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_scalar_jacobian(self, execute_kwargs, tol): + """Test scalar jacobian calculation""" + a = jax.numpy.array(0.1) + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, **execute_kwargs)[0] + + res = jax.jit(jax.grad(cost))(a) + assert res.shape == () + + # compare to standard tape jacobian + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(dev.batch_execute(tapes)) + + assert expected.shape == () + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_reusing_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + assert tape.trainable_params == [0, 1] + + def cost(a, b): + # An explicit call to _update() is required here to update the + # trainable parameters in between tape executions. + # This is different from how the autograd interface works. + # Unless the update is issued, the validation check related to the + # number of provided parameters fails in the tape: (len(params) != + # required_length) and the tape produces incorrect results. + tape._update() + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return execute([new_tape], dev, **execute_kwargs)[0] + + jac_fn = jax.jit(jax.grad(cost)) + jac = jac_fn(a, b) + + a = jax.numpy.array(0.54) + b = jax.numpy.array(0.8) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = [np.cos(2 * a)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac_fn = jax.jit(jax.grad(lambda a, b: cost(2 * a, b))) + jac = jac_fn(a, b) + expected = -2 * np.sin(2 * a) + assert np.allclose(jac, expected, atol=tol, rtol=0) + + def test_grad_with_backward_mode(self, execute_kwargs): + """Test jax grad for adjoint diff method in backward mode""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res + + cost = jax.jit(cost) + + results = jax.grad(cost)(params, cache=None) + for r, e in zip(results, expected_results): + assert jax.numpy.allclose(r, e, atol=1e-7) + + def test_classical_processing_single_tape(self, execute_kwargs): + """Test classical processing within the quantum tape for a single tape""" + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + c = jax.numpy.array(0.3) + + def cost(a, b, c, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a * c, wires=0) + qml.RZ(b, wires=0) + qml.RX(c + c**2 + jax.numpy.sin(a), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = jax.jit(jax.grad(cost, argnums=(0, 1, 2)), static_argnums=3)(a, b, c, device=dev) + assert len(res) == 3 + + def test_classical_processing_multiple_tapes(self, execute_kwargs): + """Test classical processing within the quantum tape for multiple + tapes""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.RY(x[0], wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.Hadamard(0) + qml.CRX(2 * x[0] * x[1], wires=[0, 1]) + qml.RX(2 * x[1], wires=[1]) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + result = execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + return result[0] + result[1] - 7 * result[1] + + res = jax.jit(jax.grad(cost_fn))(params) + assert res.shape == (2,) + + def test_multiple_tapes_output(self, execute_kwargs): + """Test the output types for the execution of multiple quantum tapes""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.RY(x[0], wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.Hadamard(0) + qml.CRX(2 * x[0] * x[1], wires=[0, 1]) + qml.RX(2 * x[1], wires=[1]) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + + res = jax.jit(cost_fn)(params) + assert isinstance(res, TensorLike) + assert all(isinstance(r, jax.numpy.ndarray) for r in res) + assert all(r.shape == () for r in res) + + def test_matrix_parameter(self, execute_kwargs, tol): + """Test that the jax interface works correctly + with a matrix parameter""" + a = jax.numpy.array(0.1) + U = jax.numpy.array([[0, 1], [1, 0]]) + + def cost(a, U, device): + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + tape.trainable_params = [0] + return execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = jax.jit(cost, static_argnums=2)(a, U, device=dev) + assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + + jac_fn = jax.grad(cost, argnums=0) + res = jac_fn(a, U, device=dev) + assert np.allclose(res, np.sin(a), atol=tol, rtol=0) + + def test_differentiable_expand(self, execute_kwargs, tol): + """Test that operation and nested tapes expansion + is differentiable""" + + class U3(qml.U3): + def expand(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + def cost_fn(a, p, device): + qscript = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] + ) + qscript = qscript.expand(stop_at=lambda obj: device.supports_operation(obj.name)) + return execute([qscript], device, **execute_kwargs)[0] + + a = jax.numpy.array(0.1) + p = jax.numpy.array([0.1, 0.2, 0.3]) + + dev = qml.device("default.qubit.legacy", wires=1) + res = jax.jit(cost_fn, static_argnums=2)(a, p, device=dev) + expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( + np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = jax.jit(jax.grad(cost_fn, argnums=1), static_argnums=2) + res = jac_fn(a, p, device=dev) + expected = jax.numpy.array( + [ + np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), + np.cos(p[1]) * np.cos(p[2]) * np.sin(a) + - np.sin(p[1]) + * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), + np.sin(a) + * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_independent_expval(self, execute_kwargs): + """Tests computing an expectation value that is independent of trainable + parameters.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jit(jax.grad(cost), static_argnums=1)(params, cache=None) + assert res.shape == (3,) + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestVectorValuedJIT: + """Test vector-valued returns for the JAX-JIT interface.""" + + @pytest.mark.parametrize( + "ret_type, shape, expected_type", + [ + ([qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], (), tuple), + ([qml.probs(wires=[0, 1])], (4,), jax.numpy.ndarray), + ([qml.probs()], (4,), jax.numpy.ndarray), + ], + ) + def test_shapes(self, execute_kwargs, ret_type, shape, expected_type): + """Test the shape of the result of vector-valued QNodes.""" + adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian" + if adjoint: + pytest.skip("The adjoint diff method doesn't support probabilities.") + + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + for r in ret_type: + qml.apply(r) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jit(cost)(params, cache=None) + assert isinstance(res, expected_type) + + if expected_type is tuple: + for r in res: + assert r.shape == shape + else: + assert res.shape == shape + + def test_independent_expval(self, execute_kwargs): + """Tests computing an expectation value that is independent of trainable + parameters.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jit(jax.grad(cost), static_argnums=1)(params, cache=None) + assert res.shape == (3,) + + ret_and_output_dim = [ + ([qml.probs(wires=0)], (2,), jax.numpy.ndarray), + ([qml.state()], (4,), jax.numpy.ndarray), + ([qml.density_matrix(wires=0)], (2, 2), jax.numpy.ndarray), + # Multi measurements + ([qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], (), tuple), + ([qml.var(qml.PauliZ(0)), qml.var(qml.PauliZ(1))], (), tuple), + ([qml.probs(wires=0), qml.probs(wires=1)], (2,), tuple), + ] + + @pytest.mark.parametrize("ret, out_dim, expected_type", ret_and_output_dim) + def test_vector_valued_qnode(self, execute_kwargs, ret, out_dim, expected_type): + """Tests the shape of vector-valued QNode results.""" + + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + grad_meth = ( + execute_kwargs["gradient_kwargs"]["method"] + if "gradient_kwargs" in execute_kwargs + else "" + ) + if "adjoint" in grad_meth and any( + r.return_type + in (qml.measurements.Probability, qml.measurements.State, qml.measurements.Variance) + for r in ret + ): + pytest.skip("Adjoint does not support probs") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + + for r in ret: + qml.apply(r) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res + + res = jax.jit(cost, static_argnums=1)(params, cache=None) + + assert isinstance(res, expected_type) + if expected_type is tuple: + for r in res: + assert r.shape == out_dim + else: + assert res.shape == out_dim + + def test_qnode_sample(self, execute_kwargs): + """Tests computing multiple expectation values in a tape.""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + grad_meth = ( + execute_kwargs["gradient_kwargs"]["method"] + if "gradient_kwargs" in execute_kwargs + else "" + ) + if "adjoint" in grad_meth or "backprop" in grad_meth: + pytest.skip("Adjoint does not support probs") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.sample(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q, shots=dev.shots) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res + + res = jax.jit(cost, static_argnums=1)(params, cache=None) + assert res.shape == (dev.shots,) + + def test_multiple_expvals_grad(self, execute_kwargs): + """Tests computing multiple expectation values in a tape.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res[0] + res[1] + + res = jax.jit(jax.grad(cost), static_argnums=1)(params, cache=None) + assert res.shape == (3,) + + def test_multi_tape_jacobian_probs_expvals(self, execute_kwargs): + """Test the jacobian computation with multiple tapes with probability + and expectation value computations.""" + adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian" + if adjoint: + pytest.skip("The adjoint diff method doesn't support probabilities.") + + def cost(x, y, device, interface, ek): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return qml.execute([tape1, tape2], device, **ek, interface=interface)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + x_ = np.array(0.543) + y_ = np.array(-0.654) + + res = cost(x, y, dev, interface="jax-jit", ek=execute_kwargs) + + exp = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) + + for r, e in zip(res, exp): + assert jax.numpy.allclose(r, e, atol=1e-7) + + +def test_diff_method_None_jit(): + """Test that jitted execution works when `gradient_fn=None`.""" + + dev = qml.device("default.qubit.jax", wires=1, shots=10) + + @jax.jit + def wrapper(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return qml.execute([tape], dev, gradient_fn=None) + + assert jax.numpy.allclose(wrapper(jax.numpy.array(0.0))[0], 1.0) diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_jax_jit_qnode_legacy.py similarity index 68% rename from tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_jax_jit_qnode_legacy.py index 1c0b8ad1f59..f109ddc0a4c 100644 --- a/tests/interfaces/default_qubit_2_integration/test_jax_jit_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_jax_jit_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2021 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. @@ -12,35 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the JAX-JIT interface with a QNode""" -import copy - # pylint: disable=too-many-arguments,too-few-public-methods -from functools import partial - import pytest -from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit -# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - [DefaultQubit(), "backprop", True, False], - [DefaultQubit(), "finite-diff", False, False], - [DefaultQubit(), "parameter-shift", False, False], - [DefaultQubit(), "adjoint", True, False], - [DefaultQubit(), "adjoint", True, True], - [ParamShiftDerivativesDevice(), "device", False, True], - [DefaultQubit(), "adjoint", False, False], - [DefaultQubit(), "spsa", False, False], - [DefaultQubit(), "hadamard", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", False, True], - [qml.device("lightning.qubit", wires=5), "adjoint", True, False], - [qml.device("lightning.qubit", wires=5), "adjoint", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", True, True], - [qml.device("lightning.qubit", wires=5), "parameter-shift", False, False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] interface_and_qubit_device_and_diff_method = [ ["auto"] + inner_list for inner_list in qubit_device_and_diff_method @@ -57,33 +43,33 @@ @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with JAX integrates with the PennyLane stack""" - def test_execution_with_interface( - self, dev, diff_method, grad_on_execution, interface, device_vjp - ): + def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): """Test execution works with the interface""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention") + if diff_method == "backprop": + pytest.skip("Test does not support backprop") + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)) - a = jax.numpy.array(0.1) + a = np.array(0.1, requires_grad=True) jax.jit(circuit)(a) assert circuit.interface == interface @@ -96,9 +82,7 @@ def circuit(a): assert isinstance(grad, jax.Array) assert grad.shape == () - def test_changing_trainability( - self, dev, diff_method, grad_on_execution, interface, device_vjp, tol - ): + def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -107,12 +91,13 @@ def test_changing_trainability( a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) + dev = qml.device(dev_name, wires=2) + @qnode( dev, interface=interface, diff_method="parameter-shift", grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -145,21 +130,21 @@ def circuit(a, b): circuit(a, b) assert circuit.qtape.trainable_params == [1] - def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): """Test classical processing within the quantum tape""" a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) c = jax.numpy.array(0.3) - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -174,24 +159,21 @@ def circuit(a, b, c): assert len(res) == 2 - def test_matrix_parameter( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the jax interface works correctly with a matrix parameter""" - - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - U = jax.numpy.array([[0, 1], [1, 0]]) a = jax.numpy.array(0.1) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -204,19 +186,16 @@ def circuit(U, a): if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [1] - def test_differentiable_expand( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that operation and nested tape expansion is differentiable""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + gradient_kwargs = { + "sampler_rng": SEED_FOR_SPSA, + "num_directions": 10, + } tol = TOL_FOR_SPSA class U3(qml.U3): @@ -228,6 +207,13 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) @@ -236,7 +222,6 @@ def decomposition(self): diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, p): @@ -263,12 +248,15 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): """Test setting jacobian options""" if diff_method != "finite-diff": pytest.skip("Test only applies to finite diff.") + a = np.array([0.1, 0.2], requires_grad=True) + dev = qml.device(dev_name, wires=1) + @qnode( dev, interface=interface, @@ -276,7 +264,6 @@ def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, h=1e-8, approx_order=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -291,37 +278,36 @@ def circuit(a): @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestVectorValuedQNode: """Test that using vector-valued QNodes with JAX integrate with the PennyLane stack""" - def test_diff_expval_expval( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_expval_expval(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test jacobian calculation""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") + gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -363,30 +349,29 @@ def circuit(a, b): assert res[1][1].shape == () assert np.allclose(res[1][1], expected[1][1], atol=tol, rtol=0) - def test_jacobian_no_evaluate( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test jacobian calculation when no prior circuit evaluation has been performed""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") - gradient_kwargs = {} + gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -427,21 +412,22 @@ def circuit(a, b): assert r.shape == () assert np.allclose(r, e, atol=tol, rtol=0) - def test_diff_single_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with a single prob output""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -450,7 +436,6 @@ def test_diff_single_probs( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -480,21 +465,22 @@ def circuit(x, y): assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0) assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0) - def test_diff_multi_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA + num_wires = 3 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -503,7 +489,6 @@ def test_diff_multi_probs( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -566,21 +551,22 @@ def circuit(x, y): assert jac[1][1].shape == (4,) assert np.allclose(jac[1][1], expected_1[1], atol=tol, rtol=0) - def test_diff_expval_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -589,7 +575,6 @@ def test_diff_expval_probs( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -643,19 +628,23 @@ def circuit(x, y): assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) def test_diff_expval_probs_sub_argnums( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs with less trainable parameters (argnums) than parameters.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") kwargs = {} - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + + num_wires = 2 + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -664,7 +653,6 @@ def test_diff_expval_probs_sub_argnums( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -697,21 +685,19 @@ def circuit(x, y): assert jac[1][0].shape == (2,) assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "hadamard": + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "hadamard": pytest.skip("Hadamard does not support var") elif diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=3) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -720,7 +706,6 @@ def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, i diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -781,8 +766,8 @@ class TestShotsIntegration: remains differentiable.""" def test_diff_method_None(self, interface): - """Test device works with diff_method=None.""" - dev = DefaultQubit() + """Test jax device works with diff_method=None.""" + dev = qml.device("default.qubit.jax", wires=1, shots=10) @jax.jit @qml.qnode(dev, diff_method=None, interface=interface) @@ -792,40 +777,52 @@ def circuit(x): assert jax.numpy.allclose(circuit(jax.numpy.array(0.0)), 1) - @pytest.mark.skip("jax.jit does not work with sample") - def test_changing_shots(self, interface): + def test_changing_shots(self, interface, mocker, tol): """Test that changing shots works on execution""" + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) + @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - circuit(a, b) + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - assert res.shape == (100, 2) # pylint:disable=comparison-with-callable + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # no additional calls def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" + dev = qml.device("default.qubit.legacy", wires=2, shots=1) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) + @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - jit_cost_fn = jax.jit(cost_fn, static_argnames=["shots"]) - res = jax.grad(jit_cost_fn, argnums=[0, 1])(a, b, shots=30000) + # TODO: jit when https://github.com/PennyLaneAI/pennylane/issues/3474 is resolved + res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000) + assert dev.shots == 1 expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] assert np.allclose(res, expected, atol=0.1, rtol=0) @@ -833,79 +830,44 @@ def cost_fn(a, b): def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" # pylint: disable=unused-argument + dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = jax.numpy.array([0.543, -0.654]) spy = mocker.spy(qml, "execute") - @qnode(DefaultQubit(), interface=interface) + # We're choosing interface="jax" such that backprop can be used in the + # test later + @qnode(dev, interface="jax") def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - cost_fn(a, b, shots=100) # pylint:disable=unexpected-keyword-arg # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn == qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift cost_fn(a, b) - assert cost_fn.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + # if we set the shots to None, backprop can now be used + cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert cost_fn.gradient_fn == "backprop" - @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)]) - def test_shot_vectors_single_measurements(self, interface, shots): - """Test jax-jit can work with shot vectors.""" - - dev = qml.device("default.qubit", shots=shots, seed=4747) - - @jax.jit - @qml.qnode(dev, interface=interface, diff_method="parameter-shift") - def circuit(x): - qml.RX(x, wires=0) - return qml.var(qml.PauliZ(0)) - - res = circuit(0.5) - expected = 1 - np.cos(0.5) ** 2 - assert qml.math.allclose(res[0], expected, atol=1e-2) - assert qml.math.allclose(res[1], expected, atol=3e-2) - - g = jax.jacobian(circuit)(0.5) - - expected_g = 2 * np.cos(0.5) * np.sin(0.5) - assert qml.math.allclose(g[0], expected_g, atol=2e-2) - assert qml.math.allclose(g[1], expected_g, atol=2e-2) - - @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)]) - def test_shot_vectors_multiple_measurements(self, interface, shots): - """Test jax-jit can work with shot vectors.""" - - dev = qml.device("default.qubit", shots=shots, seed=987548) - - @jax.jit - @qml.qnode(dev, interface=interface, diff_method="parameter-shift") - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)), qml.probs(wires=0) - - res = circuit(0.5) - assert qml.math.allclose(res[0][0], np.cos(0.5), atol=5e-3) - assert qml.math.allclose(res[1][0], np.cos(0.5), atol=5e-3) - expected_probs = np.array([np.cos(0.25) ** 2, np.sin(0.25) ** 2]) - assert qml.math.allclose(res[0][1], expected_probs, atol=5e-3) - assert qml.math.allclose(res[1][1], expected_probs, atol=5e-3) + cost_fn(a, b) + assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_sampling(self, dev_name, diff_method, grad_on_execution, interface): """Test sampling works as expected""" if grad_on_execution: pytest.skip("Sampling not possible with forward grad_on_execution differentiation.") @@ -913,19 +875,17 @@ def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp, interfa if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev = qml.device(dev_name, wires=2, shots=10) + + @qnode( + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) - return qml.sample(qml.Z(0)), qml.sample(qml.s_prod(2, qml.X(0) @ qml.Y(1))) + return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = jax.jit(circuit, static_argnames="shots")(shots=10) + res = jax.jit(circuit)() assert isinstance(res, tuple) @@ -934,7 +894,7 @@ def circuit(): assert isinstance(res[1], jax.Array) assert res[1].shape == (10,) - def test_counts(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_counts(self, dev_name, diff_method, grad_on_execution, interface): """Test counts works as expected""" if grad_on_execution: pytest.skip("Sampling not possible with forward grad_on_execution differentiation.") @@ -942,12 +902,10 @@ def test_counts(self, dev, diff_method, grad_on_execution, device_vjp, interface if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") + dev = qml.device(dev_name, wires=2, shots=10) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(): qml.Hadamard(wires=[0]) @@ -958,9 +916,9 @@ def circuit(): with pytest.raises( NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts." ): - jax.jit(circuit, static_argnames="shots")(shots=10) + jax.jit(circuit)() else: - res = jax.jit(circuit, static_argnames="shots")(shots=10) + res = jax.jit(circuit)() assert isinstance(res, tuple) @@ -969,32 +927,28 @@ def circuit(): assert isinstance(res[1], dict) assert len(res[1]) == 2 - def test_chained_qnodes(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_chained_qnodes(self, dev_name, diff_method, grad_on_execution, interface): """Test that the gradient of chained QNodes works without error""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit2(data, weights): qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) @@ -1020,76 +974,35 @@ def cost(weights): assert len(res) == 2 - def test_postselection_differentiation( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention") - elif dev.name == "lightning.qubit": - pytest.xfail("lightning qubit does not support postselection.") - - @qml.qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - phi = jax.numpy.array(1.23) - theta = jax.numpy.array(4.56) - - assert np.allclose(jax.jit(circuit)(phi, theta), jax.jit(expected_circuit)(theta)) - - gradient = jax.jit(jax.grad(circuit, argnums=[0, 1]))(phi, theta) - exp_theta_grad = jax.jit(jax.grad(expected_circuit))(theta) - assert np.allclose(gradient, [0.0, exp_theta_grad]) - @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegrationHigherOrder: """Tests that ensure various qubit circuits integrate correctly when computing higher-order derivatives""" - def test_second_derivative( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_second_derivative(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test second derivative calculation of a scalar-valued QNode""" + gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: - pytest.skip("Adjoint does not support second derivatives.") + if diff_method == "adjoint": + pytest.skip("Adjoint does not second derivative.") elif diff_method == "spsa": - gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - gradient_kwargs["num_directions"] = 20 - gradient_kwargs["h"] = H_FOR_SPSA + gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA, "num_directions": 10} tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1120,25 +1033,31 @@ def circuit(x): else: assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test hessian calculation of a scalar-valued QNode""" gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { "h": H_FOR_SPSA, - "num_directions": 40, + "num_directions": 20, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1172,12 +1091,10 @@ def circuit(x): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test hessian calculation of a vector-valued QNode""" gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1187,12 +1104,18 @@ def test_hessian_vector_valued( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1236,11 +1159,11 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, dev, diff_method, interface, device_vjp, grad_on_execution, tol + self, dev_name, diff_method, interface, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1250,12 +1173,18 @@ def test_hessian_vector_valued_postprocessing( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1302,11 +1231,11 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" gradient_kwargs = {} - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1316,13 +1245,19 @@ def test_hessian_vector_valued_separate_args( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, max_diff=2, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -1347,7 +1282,6 @@ def circuit(a, b): ] ) assert np.allclose(g, expected_g.T, atol=tol, rtol=0) - hess = jax.jit(jax.jacobian(jac_fn, argnums=[0, 1]))(a, b) expected_hess = np.array( @@ -1367,24 +1301,23 @@ def circuit(a, b): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_state(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 - if dev.name == "lightning.qubit" and diff_method == "adjoint": - pytest.xfail("lightning.qubit does not support adjoint with the state.") + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) - if not dev.wires: - dev = copy.copy(dev) - dev._wires = qml.wires.Wires([0, 1]) # pylint:disable=protected-access @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1394,7 +1327,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is np.dtype("complex128") + assert res.dtype is np.dtype("complex128") # pylint:disable=no-member probs = jax.numpy.abs(res) ** 2 return probs[0] + probs[2] @@ -1408,13 +1341,9 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector( - self, state, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_projector(self, state, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the variance of a projector is correctly returned""" gradient_kwargs = {} - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") if diff_method == "adjoint": pytest.skip("Adjoint does not support projectors") elif diff_method == "hadamard": @@ -1423,6 +1352,7 @@ def test_projector( gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=2) P = jax.numpy.array(state) x, y = 0.765, -0.654 @@ -1431,7 +1361,6 @@ def test_projector( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1454,9 +1383,97 @@ def circuit(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) +# TODO: Add CV test when return types and custom diff are compatible +@pytest.mark.parametrize( + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": H_FOR_SPSA}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +@pytest.mark.parametrize("interface", ["auto", "jax-jit", "jax"]) +class TestCV: + """Tests for CV integration""" + + def test_first_order_observable(self, diff_method, kwargs, interface, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + r = 0.543 + phi = -0.654 + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + res = circuit(r, phi) + expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = jax.grad(circuit, argnums=[0, 1])(r, phi) + expected = np.array( + [ + 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, + 2 * np.sinh(2 * r) * np.sin(2 * phi), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_second_order_observable(self, diff_method, kwargs, interface, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + n = 0.12 + a = 0.765 + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) + + res = circuit(n, a) + expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = jax.grad(circuit, argnums=[0, 1])(n, a) + expected = np.array([2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +# TODO: add support for fwd grad_on_execution to JAX-JIT +@pytest.mark.parametrize("interface", ["auto", "jax-jit"]) +def test_adjoint_reuse_device_state(mocker, interface): + """Tests that the jax interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=1) + + @qnode(dev, interface=interface, diff_method="adjoint") + def circ(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + jax.grad(circ)(1.0) + assert circ.device.num_executions == 1 + + spy.assert_called_with(mocker.ANY, use_device_state=True) + + @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly @@ -1464,13 +1481,20 @@ class TestTapeExpansion: @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface + self, dev_name, diff_method, grad_on_execution, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa"): pytest.skip("Only supports gradient transforms") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + class PhaseShift(qml.PhaseShift): grad_method = None @@ -1483,7 +1507,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, - device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1494,19 +1517,16 @@ def circuit(x, y): x = jax.numpy.array(0.5) y = jax.numpy.array(0.7) circuit(x, y) - jax.grad(circuit, argnums=[0])(x, y) @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, max_diff, device_vjp, interface, mocker, tol + self, dev_name, diff_method, grad_on_execution, max_diff, interface, mocker, tol ): """Test that the Hamiltonian is not expanded if there are non-commuting groups and the number of shots is None and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - if dev.name == "param_shift.qubit": - pytest.xfail("gradients transforms have a different vjp shape convention.") if diff_method == "adjoint": pytest.skip("The adjoint method does not yet support Hamiltonians") elif diff_method == "hadamard": @@ -1519,17 +1539,16 @@ def test_hamiltonian_expansion_analytic( } tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=3, shots=None) spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] - @jax.jit @qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1573,16 +1592,14 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, interface, max_diff, mocker ): - """Test that the Hamiltonian is correctly measured if there + """Test that the Hamiltonian is expanded if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} tol = 0.3 - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") if diff_method in ("adjoint", "backprop", "finite-diff"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "hadamard": @@ -1591,6 +1608,8 @@ def test_hamiltonian_finite_shots( gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA, "h": H_FOR_SPSA, "num_directions": 20} tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=3, shots=50000) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1599,7 +1618,6 @@ def test_hamiltonian_finite_shots( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1615,12 +1633,13 @@ def circuit(data, weights, coeffs): c = jax.numpy.array([-0.6543, 0.24, 0.54]) # test output - res = circuit(d, w, c, shots=50000) # pylint:disable=unexpected-keyword-arg + res = circuit(d, w, c) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) + spy.assert_called() # test gradients - grad = jax.grad(circuit, argnums=[1, 2])(d, w, c, shots=50000) + grad = jax.grad(circuit, argnums=[1, 2])(d, w, c) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1645,21 +1664,27 @@ def circuit(data, weights, coeffs): # assert np.allclose(grad2_w_c, expected, atol=tol) def test_vmap_compared_param_broadcasting( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when vectorized=True is specified for the callback when caching is disabled.""" - if ( - dev.name == "default.qubit" - and diff_method == "adjoint" - and grad_on_execution - and not device_vjp - ): - pytest.xfail("adjoint is incompatible with parameter broadcasting.") interface = "jax-jit" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not yet support Hamiltonians") + elif diff_method == "hadamard": + pytest.skip("The Hadamard method does not yet support Hamiltonians") + + if diff_method == "backprop": + pytest.skip( + "The backprop method does not yet support parameter-broadcasting with Hamiltonians" + ) + phys_qubits = 2 + if diff_method == "hadamard": + phys_qubits = 3 n_configs = 5 pars_q = np.random.rand(n_configs, 2) + dev = qml.device(dev_name, wires=tuple(range(phys_qubits)), shots=None) def minimal_circ(params): @qml.qnode( @@ -1667,80 +1692,42 @@ def minimal_circ(params): interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, cache=None, ) def _measure_operator(): qml.RY(params[..., 0], wires=0) qml.RY(params[..., 1], wires=1) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) + op = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)]) + return qml.expval(op) res = _measure_operator() return res - res1 = jax.jit(minimal_circ)(pars_q) - res2 = jax.jit(jax.vmap(minimal_circ))(pars_q) - assert np.allclose(res1, res2, tol) + assert np.allclose( + jax.jit(minimal_circ)(pars_q), jax.jit(jax.vmap(minimal_circ))(pars_q), tol + ) def test_vmap_compared_param_broadcasting_multi_output( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when vectorized=True is specified for the callback when caching is disabled and when multiple output values are returned.""" - if ( - dev.name == "default.qubit" - and diff_method == "adjoint" - and grad_on_execution - and not device_vjp - ): - pytest.xfail("adjoint is incompatible with parameter broadcasting.") interface = "jax-jit" + if diff_method == "adjoint": + pytest.skip("The adjoint method does not yet support Hamiltonians") + elif diff_method == "hadamard": + pytest.skip("The Hadamard method does not yet support Hamiltonians") - n_configs = 5 - pars_q = np.random.rand(n_configs, 2) - - def minimal_circ(params): - @qml.qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - cache=None, + if diff_method == "backprop": + pytest.skip( + "The backprop method does not yet support parameter-broadcasting with Hamiltonians" ) - def _measure_operator(): - qml.RY(params[..., 0], wires=0) - qml.RY(params[..., 1], wires=1) - return qml.expval(qml.Z(0) @ qml.Z(1)), qml.expval(qml.X(0) @ qml.X(1)) - - res = _measure_operator() - return res - - res1, res2 = jax.jit(minimal_circ)(pars_q) - vres1, vres2 = jax.jit(jax.vmap(minimal_circ))(pars_q) - assert np.allclose(res1, vres1, tol) - assert np.allclose(res2, vres2, tol) - - def test_vmap_compared_param_broadcasting_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): - """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when - vectorized=True is specified for the callback when caching is disabled and when multiple output values - are returned.""" - if ( - dev.name == "default.qubit" - and diff_method == "adjoint" - and grad_on_execution - and not device_vjp - ): - pytest.xfail("adjoint is incompatible with parameter broadcasting.") - elif dev.name == "lightning.qubit" and diff_method == "adjoint": - pytest.xfail("lightning adjoign cannot differentiate probabilities.") - interface = "jax-jit" + phys_qubits = 2 n_configs = 5 pars_q = np.random.rand(n_configs, 2) + dev = qml.device(dev_name, wires=tuple(range(phys_qubits)), shots=None) def minimal_circ(params): @qml.qnode( @@ -1748,13 +1735,14 @@ def minimal_circ(params): interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, cache=None, ) def _measure_operator(): qml.RY(params[..., 0], wires=0) qml.RY(params[..., 1], wires=1) - return qml.probs(wires=0), qml.probs(wires=1) + op1 = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)]) + op2 = qml.Hamiltonian([1.0], [qml.PauliX(0) @ qml.PauliX(1)]) + return qml.expval(op1), qml.expval(op2) res = _measure_operator() return res @@ -1770,25 +1758,23 @@ def _measure_operator(): @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestJIT: """Test JAX JIT integration with the QNode and automatic resolution of the correct JAX interface variant.""" - def test_gradient( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface - ): + def test_gradient(self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface): """Test derivative calculation of a scalar valued QNode""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + gradient_kwargs = {} - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if device_vjp and jacobian == jax.jacfwd: - pytest.skip("device vjps not compatible with forward diff.") - elif diff_method == "spsa": + if diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA @@ -1797,7 +1783,6 @@ def test_gradient( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x): @@ -1821,9 +1806,7 @@ def circuit(x): "ignore:Requested adjoint differentiation to be computed with finite shots." ) @pytest.mark.parametrize("shots", [10, 1000]) - def test_hermitian( - self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface - ): + def test_hermitian(self, dev_name, diff_method, grad_on_execution, shots, jacobian, interface): """Test that the jax device works with qml.Hermitian and jitting even when shots>0. @@ -1831,6 +1814,13 @@ def test_hermitian( to different reasons, hence the parametrization in the test. """ # pylint: disable=unused-argument + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + if diff_method == "backprop": pytest.skip("Backpropagation is unsupported if shots > 0.") @@ -1840,11 +1830,7 @@ def test_hermitian( projector = np.array(qml.matrix(qml.PauliZ(0) @ qml.PauliZ(1))) @qml.qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circ(projector): return qml.expval(qml.Hermitian(projector, wires=range(2))) @@ -1857,20 +1843,23 @@ def circ(projector): ) @pytest.mark.parametrize("shots", [10, 1000]) def test_probs_obs_none( - self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface + self, dev_name, diff_method, grad_on_execution, shots, jacobian, interface ): """Test that the jax device works with qml.probs, a MeasurementProcess that has obs=None even when shots>0.""" # pylint: disable=unused-argument + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + if diff_method in ["backprop", "adjoint"]: pytest.skip("Backpropagation is unsupported if shots > 0.") @qml.qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(): return qml.probs(wires=0) @@ -1881,30 +1870,23 @@ def circuit(): # reason="Non-trainable parameters are not being correctly unwrapped by the interface" # ) def test_gradient_subset( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface + self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface ): """Test derivative calculation of a scalar valued QNode with respect to a subset of arguments""" - if diff_method == "spsa" and not grad_on_execution and not device_vjp: + if diff_method == "spsa" and not grad_on_execution: pytest.xfail(reason="incorrect jacobian results") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if diff_method == "device" and not grad_on_execution and device_vjp: - pytest.xfail(reason="various runtime-related errors") - - if diff_method == "adjoint" and device_vjp and jacobian is jax.jacfwd: - pytest.xfail(reason="TypeError applying forward-mode autodiff.") + if diff_method == "hadamard" and not grad_on_execution: + pytest.xfail(reason="XLA raised wire error") a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) + dev = qml.device(dev_name, wires=1) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b, c): qml.RY(a, wires=0) @@ -1921,17 +1903,20 @@ def circuit(a, b, c): assert np.allclose(g, expected_g, atol=tol, rtol=0) def test_gradient_scalar_cost_vector_valued_qnode( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface + self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface ): """Test derivative calculation of a scalar valued cost function that uses the output of a vector-valued QNode""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + gradient_kwargs = {} - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - elif jacobian == jax.jacfwd and device_vjp: - pytest.skip("device vjps are not compatible with forward differentiation.") + if diff_method == "adjoint": + pytest.xfail(reason="Adjoint does not support probs.") elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA @@ -1941,7 +1926,6 @@ def test_gradient_scalar_cost_vector_valued_qnode( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1974,26 +1958,21 @@ def cost(x, y, idx): assert np.allclose(g0, expected_g[0][idx], atol=tol, rtol=0) assert np.allclose(g1, expected_g[1][idx], atol=tol, rtol=0) - # pylint: disable=unused-argument def test_matrix_parameter( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface + self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface ): """Test that the JAX-JIT interface works correctly with a matrix parameter""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("device vjps are not compatible with forward differentiation.") - # pylint: disable=unused-argument + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circ(p, U): qml.QubitUnitary(U, wires=0) @@ -2005,7 +1984,7 @@ def circ(p, U): res = jax.jit(circ)(p, U) assert np.allclose(res, -np.cos(p), atol=tol, rtol=0) - jac_fn = jax.jit(jacobian(circ, argnums=0)) + jac_fn = jax.jit(jax.grad(circ, argnums=(0))) res = jac_fn(p, U) assert np.allclose(res, np.sin(p), atol=tol, rtol=0) @@ -2013,31 +1992,27 @@ def circ(p, U): @pytest.mark.parametrize("shots", [None, 10000]) @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestReturn: """Class to test the shape of the Grad/Jacobian with different return types.""" def test_grad_single_measurement_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For one measurement and one param, the gradient is a float.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -2046,30 +2021,27 @@ def circuit(a): a = jax.numpy.array(0.1) - grad = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + grad = jax.jit(jacobian(circuit))(a) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -2079,9 +2051,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - grad = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( - a, b, shots=shots - ) + grad = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) assert isinstance(grad, tuple) assert len(grad) == 2 @@ -2089,24 +2059,21 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2115,31 +2082,31 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - grad = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + grad = jax.jit(jacobian(circuit))(a) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -2148,31 +2115,30 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -2182,7 +2148,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")(a, b, shots=shots) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) assert isinstance(jac, tuple) @@ -2193,25 +2159,24 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2219,33 +2184,29 @@ def circuit(a): return qml.probs(wires=[0, 1]) a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2253,9 +2214,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( - par_0, par_1, shots=shots - ) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) assert isinstance(jac, tuple) @@ -2274,24 +2233,20 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2300,7 +2255,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2312,29 +2267,23 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_var_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of var.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or only diagonal measurements") + + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2342,9 +2291,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( - par_0, par_1, shots=shots - ) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2364,26 +2311,20 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_var_var_multiple_params_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of var.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or only diagonal measurements") + + dev = qml.device(dev_name, wires=2, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2392,7 +2333,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2404,24 +2345,23 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a single params return an array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if device_vjp and jacobian == jax.jacfwd: - pytest.skip("device vjp not compatible with forward differentiation.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -2430,7 +2370,7 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2442,24 +2382,23 @@ def circuit(a): assert jac[1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -2469,7 +2408,7 @@ def circuit(a, b): a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")(a, b, shots=shots) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2489,24 +2428,23 @@ def circuit(a, b): assert jac[1][1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transforms have a different vjp shape convention.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -2515,7 +2453,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) + jac = jax.jit(jacobian(circuit))(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2536,17 +2474,23 @@ def circuit(a): @pytest.mark.parametrize("hessian", hessian_fn) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestReturnHessian: """Class to test the shape of the Hessian with different return types.""" def test_hessian_expval_multiple_params( - self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - if diff_method in {"adjoint", "device"}: + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") par_0 = jax.numpy.array(0.1) @@ -2558,7 +2502,6 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2586,13 +2529,20 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2601,7 +2551,6 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2615,10 +2564,12 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_var_multiple_params( - self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - if diff_method in {"adjoint", "device"}: + dev = qml.device(dev_name, wires=2) + + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") @@ -2632,7 +2583,6 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2660,14 +2610,16 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") + dev = qml.device(dev_name, wires=2) + params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2676,7 +2628,6 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2690,10 +2641,17 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - if diff_method in {"adjoint", "device"}: + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of non commuting obs.") @@ -2707,7 +2665,6 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2734,15 +2691,22 @@ def circuit(x, y): assert h_comp.shape == (2,) def test_hessian_probs_expval_multiple_param_array( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of non commuting obs.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2751,7 +2715,6 @@ def test_hessian_probs_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2770,10 +2733,12 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - if diff_method in {"adjoint", "device"}: + dev = qml.device(dev_name, wires=2) + + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") @@ -2787,7 +2752,6 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2814,14 +2778,16 @@ def circuit(x, y): assert h_comp.shape == (2,) def test_hessian_probs_var_multiple_param_array( - self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, hessian, grad_on_execution, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - if diff_method in {"adjoint", "device"}: + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") + dev = qml.device(dev_name, wires=2) + params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2830,7 +2796,6 @@ def test_hessian_probs_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2853,9 +2818,15 @@ def circuit(x): @pytest.mark.parametrize("diff_method", ["parameter-shift", "hadamard"]) def test_jax_device_hessian_shots(hessian, diff_method): """The hessian of multiple measurements with a multiple param array return a single array.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 - @partial(jax.jit, static_argnames="shots") - @qml.qnode(DefaultQubit(), diff_method=diff_method, max_diff=2) + dev = qml.device("default.qubit.jax", wires=num_wires, shots=10000) + + @jax.jit + @qml.qnode(dev, diff_method=diff_method, max_diff=2) def circuit(x): qml.RY(x[0], wires=0) qml.RX(x[1], wires=0) @@ -2864,7 +2835,7 @@ def circuit(x): x = jax.numpy.array([1.0, 2.0]) a, b = x - hess = jax.jit(hessian(circuit), static_argnames="shots")(x, shots=10000) + hess = jax.jit(hessian(circuit))(x) expected_hess = [ [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)], @@ -2878,33 +2849,28 @@ def circuit(x): @pytest.mark.parametrize("argnums", [0, 1, [0, 1]]) @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestSubsetArgnums: def test_single_measurement( self, interface, - dev, + dev_name, diff_method, grad_on_execution, - device_vjp, jacobian, argnums, jit_inside, tol, ): """Test single measurement with different diff methods with argnums.""" + + dev = qml.device(dev_name, wires=3) + kwargs = {} - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transform have a different vjp shape convention.") if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) @qml.qnode( dev, @@ -2912,7 +2878,6 @@ def test_single_measurement( diff_method=diff_method, grad_on_execution=grad_on_execution, cache=False, - device_vjp=device_vjp, **kwargs, ) def circuit(a, b): @@ -2942,34 +2907,27 @@ def circuit(a, b): def test_multi_measurements( self, interface, - dev, + dev_name, diff_method, grad_on_execution, - device_vjp, jacobian, argnums, jit_inside, tol, ): """Test multiple measurements with different diff methods with argnums.""" - if jacobian == jax.jacfwd and device_vjp: - pytest.skip("jacfwd is not compatible with device_vjp=True.") - if "lightning" in dev.name: - pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if dev.name == "param_shift.qubit": - pytest.xfail("gradient transform have a different vjp shape convention.") + dev = qml.device(dev_name, wires=3) kwargs = {} if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) @qml.qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **kwargs, ) def circuit(a, b): @@ -2995,72 +2953,3 @@ def circuit(a, b): else: assert np.allclose(jac[0], expected[0], atol=tol) assert np.allclose(jac[1], expected[1], atol=tol) - - -class TestSinglePrecision: - """Tests for compatibility with single precision mode.""" - - # pylint: disable=import-outside-toplevel - def test_type_conversion_fallback(self): - """Test that if the type isn't int, float, or complex, we still have a fallback.""" - from pennylane.workflow.interfaces.jax_jit import _jax_dtype - - assert _jax_dtype(bool) == jax.numpy.dtype(bool) - - @pytest.mark.parametrize("diff_method", ("adjoint", "parameter-shift")) - def test_float32_return(self, diff_method): - """Test that jax jit works when float64 mode is disabled.""" - jax.config.update("jax_enable_x64", False) - - try: - - @jax.jit - @qml.qnode(qml.device("default.qubit"), diff_method=diff_method) - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - grad = jax.grad(circuit)(jax.numpy.array(0.1)) - assert qml.math.allclose(grad, -np.sin(0.1)) - finally: - jax.config.update("jax_enable_x64", True) - jax.config.update("jax_enable_x64", True) - - @pytest.mark.parametrize("diff_method", ("adjoint", "finite-diff")) - def test_complex64_return(self, diff_method): - """Test that jax jit works with differentiating the state.""" - jax.config.update("jax_enable_x64", False) - - try: - tol = 2e-2 if diff_method == "finite-diff" else 1e-6 - - @jax.jit - @qml.qnode(qml.device("default.qubit", wires=1), diff_method=diff_method) - def circuit(x): - qml.RX(x, wires=0) - return qml.state() - - j = jax.jacobian(circuit, holomorphic=True)(jax.numpy.array(0.1 + 0j)) - assert qml.math.allclose(j, [-np.sin(0.05) / 2, -np.cos(0.05) / 2 * 1j], atol=tol) - - finally: - jax.config.update("jax_enable_x64", True) - jax.config.update("jax_enable_x64", True) - - def test_int32_return(self): - """Test that jax jit forward execution works with samples and int32""" - - jax.config.update("jax_enable_x64", False) - - try: - - @jax.jit - @qml.qnode(qml.device("default.qubit", shots=10), diff_method=qml.gradients.param_shift) - def circuit(x): - qml.RX(x, wires=0) - return qml.sample(wires=0) - - _ = circuit(jax.numpy.array(0.1)) - finally: - jax.config.update("jax_enable_x64", True) - jax.config.update("jax_enable_x64", True) diff --git a/tests/interfaces/legacy_devices_integration/test_jax_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_legacy.py new file mode 100644 index 00000000000..e9b7fdd8904 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_jax_legacy.py @@ -0,0 +1,876 @@ +# Copyright 2018-2021 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. +"""Unit tests for the JAX-Python interface""" +import numpy as np + +# pylint: disable=protected-access,too-few-public-methods +import pytest + +import pennylane as qml +from pennylane import execute +from pennylane.gradients import param_shift +from pennylane.typing import TensorLike + +pytestmark = pytest.mark.jax + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) + + +class TestJaxExecuteUnitTests: + """Unit tests for jax execution""" + + def test_jacobian_options(self, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + device, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + )[0] + + jax.grad(cost)(a, device=dev) + + for args in spy.call_args_list: + assert args[1]["shifts"] == [(np.pi / 4,)] * 2 + + def test_incorrect_grad_on_execution(self): + """Test that an error is raised if an gradient transform + is used with grad_on_execution=True""" + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + device, + gradient_fn=param_shift, + grad_on_execution=True, + )[0] + + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + jax.grad(cost)(a, device=dev) + + def test_unknown_interface(self): + """Test that an error is raised if the interface is unknown""" + a = jax.numpy.array([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + device, + gradient_fn=param_shift, + interface="None", + )[0] + + with pytest.raises(ValueError, match="Unknown interface"): + cost(a, device=dev) + + def test_grad_on_execution(self, mocker): + """Test that grad_on_execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=2) + spy = mocker.spy(dev, "execute_and_gradients") + + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + ) + + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), 0)], + [qml.expval(qml.PauliZ(0))], + ) + + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + ) + return execute( + [tape1, tape2, tape3], + dev, + gradient_fn="device", + gradient_kwargs={ + "method": "adjoint_jacobian", + "use_device_state": True, + }, + ) + + a = jax.numpy.array([0.1, 0.2]) + res = cost(a) + + x, y = a + assert np.allclose(res[0][0], np.cos(x) * np.cos(y)) + assert np.allclose(res[0][1], 1) + assert np.allclose(res[1], np.cos(0.5)) + assert np.allclose(res[2], np.cos(x) * np.cos(y)) + + # adjoint method only performs a single device execution per tape, but gets both result and gradient + assert dev.num_executions == 3 + spy.assert_not_called() + + g = jax.jacobian(cost)(a) + spy.assert_called() + expected_g = (-np.sin(x) * np.cos(y), -np.cos(x) * np.sin(y)) + assert qml.math.allclose(g[0][0], expected_g) + assert qml.math.allclose(g[0][1], np.zeros(2)) + assert qml.math.allclose(g[1], np.zeros(2)) + assert qml.math.allclose(g[2], expected_g) + + def test_no_grad_on_execution(self, mocker): + """Test that `grad_on_execution=False` uses the `device.execute_and_gradients`.""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "execute_and_gradients") + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + a = jax.numpy.array([0.1, 0.2]) + cost(a) + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + jax.grad(cost)(a) + spy_gradients.assert_called() + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cachesize): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn=param_shift, + cachesize=cachesize, + )[0] + + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cachesize=2) + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn=param_shift, + cache=cache, + )[0] + + custom_cache = {} + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_custom_cache_multiple(self, mocker): + """Test the use of a custom cache object with multiple tapes""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + def cost(a, b, cache): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + res = execute( + [tape1, tape2], + dev, + gradient_fn=param_shift, + cache=cache, + ) + return res[0] + + custom_cache = {} + jax.grad(cost)(a, b, cache=custom_cache) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_caching_param_shift(self, tol): + """Test that, when using parameter-shift transform, + caching produces the optimum number of evaluations.""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn=param_shift, + cache=cache, + )[0] + + # Without caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) + params = jax.numpy.array([0.1, 0.2]) + jax.grad(cost)(params, cache=None) + assert dev.num_executions == 5 + + # With caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + jac_fn = jax.grad(cost) + grad1 = jac_fn(params, cache=True) + assert dev.num_executions == 5 + + # Check that calling the cost function again + # continues to evaluate the device (that is, the cache + # is emptied between calls) + grad2 = jac_fn(params, cache=True) + assert dev.num_executions == 10 + assert np.allclose(grad1, grad2, atol=tol, rtol=0) + + # Check that calling the cost function again + # with different parameters produces a different Jacobian + grad2 = jac_fn(2 * params, cache=True) + assert dev.num_executions == 15 + assert not np.allclose(grad1, grad2, atol=tol, rtol=0) + + def test_caching_adjoint_backward(self): + """Test that caching produces the optimum number of adjoint evaluations + when no grad on execution.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], + dev, + gradient_fn="device", + cache=cache, + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + )[0] + + # Without caching, 2 evaluations are required. + # 1 for the forward pass, and one per output dimension + # on the backward pass. + jax.grad(cost)(params, cache=None) + assert dev.num_executions == 2 + + # With caching, also 2 evaluations are required. One + # for the forward pass, and one for the backward pass. + dev._num_executions = 0 + jac_fn = jax.grad(cost) + jac_fn(params, cache=True) + assert dev.num_executions == 2 + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + }, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestJaxExecuteIntegration: + """Test the jax interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, execute_kwargs): + """Test execution""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, b): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return execute([tape1, tape2], dev, **execute_kwargs) + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + res = cost(a, b) + + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_scalar_jacobian(self, execute_kwargs, tol): + """Test scalar jacobian calculation""" + a = jax.numpy.array(0.1) + dev = qml.device("default.qubit.legacy", wires=2) + + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, **execute_kwargs)[0] + + res = jax.grad(cost)(a) + assert res.shape == () + + # compare to standard tape jacobian + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(dev.batch_execute(tapes)) + + assert expected.shape == () + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_reusing_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + assert tape.trainable_params == [0, 1] + + def cost(a, b): + # An explicit call to _update() is required here to update the + # trainable parameters in between tape executions. + # This is different from how the autograd interface works. + # Unless the update is issued, the validation check related to the + # number of provided parameters fails in the tape: (len(params) != + # required_length) and the tape produces incorrect results. + tape._update() + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return execute([new_tape], dev, **execute_kwargs)[0] + + jac_fn = jax.grad(cost) + jac = jac_fn(a, b) + + a = jax.numpy.array(0.54) + b = jax.numpy.array(0.8) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = [np.cos(2 * a)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac_fn = jax.grad(lambda a, b: cost(2 * a, b)) + jac = jac_fn(a, b) + expected = -2 * np.sin(2 * a) + assert np.allclose(jac, expected, atol=tol, rtol=0) + + def test_grad_with_different_grad_on_execution(self, execute_kwargs): + """Test jax grad for adjoint diff method with different execution kwargs.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] + return res + + results = jax.grad(cost)(params, cache=None) + for r, e in zip(results, expected_results): + assert jax.numpy.allclose(r, e, atol=1e-7) + + def test_classical_processing_single_tape(self, execute_kwargs): + """Test classical processing within the quantum tape for a single tape""" + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + c = jax.numpy.array(0.3) + + def cost(a, b, c, device): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a * c, wires=0) + qml.RZ(b, wires=0) + qml.RX(c + c**2 + jax.numpy.sin(a), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = jax.grad(cost, argnums=(0, 1, 2))(a, b, c, device=dev) + assert len(res) == 3 + + def test_classical_processing_multiple_tapes(self, execute_kwargs): + """Test classical processing within the quantum tape for multiple + tapes""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.RY(x[0], wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.Hadamard(0) + qml.CRX(2 * x[0] * x[1], wires=[0, 1]) + qml.RX(2 * x[1], wires=[1]) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + return result[0] + result[1] - 7 * result[1] + + res = jax.grad(cost_fn)(params) + assert res.shape == (2,) + + def test_multiple_tapes_output(self, execute_kwargs): + """Test the output types for the execution of multiple quantum tapes""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.RY(x[0], wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.Hadamard(0) + qml.CRX(2 * x[0] * x[1], wires=[0, 1]) + qml.RX(2 * x[1], wires=[1]) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + return execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + + res = cost_fn(params) + assert isinstance(res, TensorLike) + assert all(isinstance(r, jax.numpy.ndarray) for r in res) + assert all(r.shape == () for r in res) + + def test_matrix_parameter(self, execute_kwargs, tol): + """Test that the jax interface works correctly + with a matrix parameter""" + a = jax.numpy.array(0.1) + U = jax.numpy.array([[0, 1], [1, 0]]) + + def cost(a, U, device): + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + tape.trainable_params = [0] + return execute([tape], device, **execute_kwargs)[0] + + dev = qml.device("default.qubit.legacy", wires=2) + res = cost(a, U, device=dev) + assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + + jac_fn = jax.grad(cost, argnums=0) + res = jac_fn(a, U, device=dev) + assert np.allclose(res, np.sin(a), atol=tol, rtol=0) + + def test_differentiable_expand(self, execute_kwargs, tol): + """Test that operation and nested tapes expansion + is differentiable""" + + class U3(qml.U3): + def decomposition(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + def cost_fn(a, p, device): + with qml.queuing.AnnotatedQueue() as q_tape: + qml.RX(a, wires=0) + U3(*p, wires=0) + qml.expval(qml.PauliX(0)) + + tape = qml.tape.QuantumScript.from_queue(q_tape) + tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name)) + return execute([tape], device, **execute_kwargs)[0] + + a = jax.numpy.array(0.1) + p = jax.numpy.array([0.1, 0.2, 0.3]) + + dev = qml.device("default.qubit.legacy", wires=1) + res = cost_fn(a, p, device=dev) + expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( + np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + jac_fn = jax.grad(cost_fn, argnums=1) + res = jac_fn(a, p, device=dev) + expected = jax.numpy.array( + [ + np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), + np.cos(p[1]) * np.cos(p[2]) * np.sin(a) + - np.sin(p[1]) + * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), + np.sin(a) + * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_independent_expval(self, execute_kwargs): + """Tests computing an expectation value that is independent of trainable + parameters.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.grad(cost)(params, cache=None) + assert res.shape == (3,) + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestVectorValued: + """Test vector-valued jacobian returns for the JAX Python interface.""" + + def test_multiple_expvals(self, execute_kwargs): + """Tests computing multiple expectation values in a tape.""" + + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = qml.execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jacobian(cost)(params, cache=None) + + assert isinstance(res, tuple) + assert len(res) == 2 + + assert res[0].shape == (3,) + assert isinstance(res[0], jax.numpy.ndarray) + + assert res[1].shape == (3,) + assert isinstance(res[1], jax.numpy.ndarray) + + def test_multiple_expvals_single_par(self, execute_kwargs): + """Tests computing multiple expectation values in a tape with a single + trainable parameter.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.1]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = qml.execute([tape], dev, cache=cache, **execute_kwargs) + return res[0] + + res = jax.jacobian(cost)(params, cache=None) + + assert isinstance(res, tuple) + + assert isinstance(res[0], jax.numpy.ndarray) + assert res[0].shape == (1,) + + assert isinstance(res[1], jax.numpy.ndarray) + assert res[1].shape == (1,) + + def test_multi_tape_fwd(self, execute_kwargs): + """Test the forward evaluation of a cost function that uses the output + of multiple tapes that be vector-valued.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = jax.numpy.array([0.3, 0.2]) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.expval(qml.PauliY(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[1], wires=[0]) + qml.RX(x[1], wires=[0]) + qml.RX(-x[1], wires=[0]) + qml.expval(qml.PauliY(0)) + qml.expval(qml.PauliY(1)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = qml.execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + return result[0] + result[1][0] + + expected = -jax.numpy.sin(params[0]) + -jax.numpy.sin(params[1]) + res = cost_fn(params) + assert jax.numpy.allclose(expected, res) + + def test_multi_tape_jacobian(self, execute_kwargs): + """Test the jacobian computation with multiple tapes.""" + + def cost(x, y, device, interface, ek): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + return qml.execute([tape1, tape2], device, **ek, interface=interface) + + dev = qml.device("default.qubit.legacy", wires=2) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + x_ = np.array(0.543) + y_ = np.array(-0.654) + + exec_jax = cost(x, y, dev, interface="jax-python", ek=execute_kwargs) + exec_autograd = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) + + assert np.allclose(exec_jax, exec_autograd) + + res = jax.jacobian(cost, argnums=(0, 1))( + x, y, dev, interface="jax-python", ek=execute_kwargs + ) + + import autograd.numpy as anp + + def cost_stack(x, y, device, interface, ek): + return anp.hstack(cost(x, y, device, interface, ek)) + + exp = qml.jacobian(cost_stack, argnum=(0, 1))( + x_, y_, dev, interface="autograd", ek=execute_kwargs + ) + res_0 = jax.numpy.array([res[0][0][0], res[0][1][0], res[1][0][0], res[1][1][0]]) + res_1 = jax.numpy.array([res[0][0][1], res[0][1][1], res[1][0][1], res[1][1][1]]) + + assert np.allclose(res_0, exp[0]) + assert np.allclose(res_1, exp[1]) + + def test_multi_tape_jacobian_probs_expvals(self, execute_kwargs): + """Test the jacobian computation with multiple tapes with probability + and expectation value computations.""" + + adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian" + if adjoint: + pytest.skip("The adjoint diff method doesn't support probabilities.") + + def cost(x, y, device, interface, ek): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + return qml.execute([tape1, tape2], device, **ek, interface=interface) + + dev = qml.device("default.qubit.legacy", wires=2) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + x_ = np.array(0.543) + y_ = np.array(-0.654) + + exec_jax = cost(x, y, dev, interface="jax-python", ek=execute_kwargs) + exec_autograd = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) + + assert all( + np.allclose(exec_jax[i][j], exec_autograd[i][j]) for i in range(2) for j in range(2) + ) + + res = jax.jacobian(cost, argnums=(0, 1))( + x, y, dev, interface="jax-python", ek=execute_kwargs + ) + + assert isinstance(res, TensorLike) + assert len(res) == 2 + + for r, exp_shape in zip(res, [(), (2,)]): + assert isinstance(r, tuple) + assert len(r) == 2 + assert len(r[0]) == 2 + assert isinstance(r[0][0], jax.numpy.ndarray) + assert r[0][0].shape == exp_shape + assert isinstance(r[0][1], jax.numpy.ndarray) + assert r[0][1].shape == exp_shape + assert len(r[1]) == 2 + assert isinstance(r[1][0], jax.numpy.ndarray) + assert r[1][0].shape == exp_shape + assert isinstance(r[1][1], jax.numpy.ndarray) + assert r[1][1].shape == exp_shape diff --git a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_jax_qnode_legacy.py similarity index 74% rename from tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_jax_qnode_legacy.py index 3ea650f96e8..fa88e315155 100644 --- a/tests/interfaces/default_qubit_2_integration/test_jax_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_jax_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2021 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. @@ -12,39 +12,27 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the JAX-Python interface with a QNode""" -# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg, use-implicit-booleaness-not-comparison - -from itertools import product - -import numpy as np +# pylint: disable=too-many-arguments,too-few-public-methods,too-many-public-methods import pytest import pennylane as qml +from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit - -device_seed = 42 - -# device, diff_method, grad_on_execution, device_vjp -device_and_diff_method = [ - [DefaultQubit(seed=device_seed), "backprop", True, False], - [DefaultQubit(seed=device_seed), "finite-diff", False, False], - [DefaultQubit(seed=device_seed), "parameter-shift", False, False], - [DefaultQubit(seed=device_seed), "adjoint", True, False], - [DefaultQubit(seed=device_seed), "adjoint", False, False], - [DefaultQubit(seed=device_seed), "adjoint", True, True], - [DefaultQubit(seed=device_seed), "adjoint", False, True], - [DefaultQubit(seed=device_seed), "spsa", False, False], - [DefaultQubit(seed=device_seed), "hadamard", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", False, True], - [qml.device("lightning.qubit", wires=5), "adjoint", True, True], - [qml.device("lightning.qubit", wires=5), "adjoint", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", True, False], +from pennylane.tape import QuantumScript + +qubit_device_and_diff_method = [ + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] -interface_and_device_and_diff_method = [ - ["auto"] + inner_list for inner_list in device_and_diff_method -] + [["jax"] + inner_list for inner_list in device_and_diff_method] +interface_and_qubit_device_and_diff_method = [ + ["auto"] + inner_list for inner_list in qubit_device_and_diff_method +] + [["jax"] + inner_list for inner_list in qubit_device_and_diff_method] pytestmark = pytest.mark.jax @@ -58,49 +46,46 @@ @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", interface_and_device_and_diff_method + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with JAX integrates with the PennyLane stack""" - def test_execution_with_interface( - self, dev, diff_method, grad_on_execution, interface, device_vjp - ): + def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)) - a = jax.numpy.array(0.1) + a = np.array(0.1, requires_grad=True) circuit(a) assert circuit.interface == interface - # jax doesn't set trainable parameters on regular execution - assert circuit.qtape.trainable_params == [] + # the tape is able to deduce trainable parameters + assert circuit.qtape.trainable_params == [0] # gradients should work grad = jax.grad(circuit)(a) assert isinstance(grad, jax.Array) - # the tape is able to deduce trainable parameters - assert circuit.qtape.trainable_params == [0] assert grad.shape == () - def test_changing_trainability( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): # pylint:disable=unused-argument + def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -109,7 +94,14 @@ def test_changing_trainability( a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - @qnode(dev, interface=interface, diff_method="parameter-shift") + dev = qml.device(dev_name, wires=2) + + @qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -135,18 +127,27 @@ def circuit(a, b): expected = [-np.sin(a) + np.sin(a) * np.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): + # trainability also updates on evaluation + a = np.array(0.54, requires_grad=False) + b = np.array(0.8, requires_grad=True) + circuit(a, b) + assert circuit.qtape.trainable_params == [1] + + def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): """Test classical processing within the quantum tape""" a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) c = jax.numpy.array(0.3) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -161,20 +162,21 @@ def circuit(a, b, c): assert len(res) == 2 - def test_matrix_parameter( - self, dev, diff_method, grad_on_execution, interface, device_vjp, tol - ): + def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the jax interface works correctly with a matrix parameter""" U = jax.numpy.array([[0, 1], [1, 0]]) a = jax.numpy.array(0.1) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -187,32 +189,39 @@ def circuit(U, a): if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [1] - def test_differentiable_expand( - self, dev, diff_method, grad_on_execution, interface, device_vjp, tol - ): + def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that operation and nested tape expansion is differentiable""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 10 + spsa_kwargs = dict( + sampler_rng=np.random.default_rng(SEED_FOR_SPSA), + num_directions=10, + ) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA - class U3(qml.U3): # pylint:disable=too-few-public-methods + class U3(qml.U3): def decomposition(self): theta, phi, lam = self.data wires = self.wires - return [ - qml.Rot(lam, theta, -lam, wires=wires), - qml.PhaseShift(phi + lam, wires=wires), - ] + with qml.queuing.AnnotatedQueue() as q_tape: + qml.Rot(lam, theta, -lam, wires=wires) + qml.PhaseShift(phi + lam, wires=wires) + + tape = QuantumScript.from_queue(q_tape) + return tape.operations + + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) @@ -241,16 +250,23 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): # pylint:disable=unused-argument + def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): """Test setting jacobian options""" if diff_method != "finite-diff": pytest.skip("Test only applies to finite diff.") - a = jax.numpy.array([0.1, 0.2]) + a = np.array([0.1, 0.2], requires_grad=True) + + dev = qml.device(dev_name, wires=1) - @qnode(dev, interface=interface, diff_method="finite-diff", h=1e-8, approx_order=2) + @qnode( + dev, + interface=interface, + diff_method="finite-diff", + h=1e-8, + approx_order=2, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -260,31 +276,30 @@ def circuit(a): @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestVectorValuedQNode: """Test that using vector-valued QNodes with JAX integrate with the PennyLane stack""" - def test_diff_expval_expval( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_expval_expval(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test jacobian calculation""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode(dev, **kwargs) def circuit(a, b): @@ -295,7 +310,7 @@ def circuit(a, b): res = circuit(a, b) - assert circuit.qtape.trainable_params == [] + assert circuit.qtape.trainable_params == [0, 1] assert isinstance(res, tuple) assert len(res) == 2 @@ -305,7 +320,6 @@ def circuit(a, b): res = jax.jacobian(circuit, argnums=[0, 1])(a, b) expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - assert circuit.qtape.trainable_params == [0, 1] assert isinstance(res, tuple) assert len(res) == 2 @@ -325,20 +339,11 @@ def circuit(a, b): assert res[1][1].shape == () assert np.allclose(res[1][1], expected[1][1], atol=tol, rtol=0) - def test_jacobian_no_evaluate( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test jacobian calculation when no prior circuit evaluation has been performed""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA @@ -346,6 +351,13 @@ def test_jacobian_no_evaluate( a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -361,10 +373,11 @@ def circuit(a, b): expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - assert isinstance(res[0][0], jax.numpy.ndarray) - for i, j in product((0, 1), (0, 1)): - assert res[i][j].shape == () - assert np.allclose(res[i][j], expected[i][j], atol=tol, rtol=0) + for _res, _exp in zip(res, expected): + for r, e in zip(_res, _exp): + assert isinstance(r, jax.numpy.ndarray) + assert r.shape == () + assert np.allclose(r, e, atol=tol, rtol=0) # call the Jacobian with new parameters a = jax.numpy.array(0.6) @@ -377,27 +390,30 @@ def circuit(a, b): expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - for i, j in product((0, 1), (0, 1)): - assert res[i][j].shape == () - assert np.allclose(res[i][j], expected[i][j], atol=tol, rtol=0) + for _res, _exp in zip(res, expected): + for r, e in zip(_res, _exp): + assert isinstance(r, jax.numpy.ndarray) + assert r.shape == () + assert np.allclose(r, e, atol=tol, rtol=0) - def test_diff_single_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with a single prob output""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -429,24 +445,24 @@ def circuit(x, y): assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0) assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0) - def test_diff_multi_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + num_wires = 3 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -468,11 +484,11 @@ def circuit(x, y): ] assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == (2,) # pylint:disable=comparison-with-callable + assert res[0].shape == (2,) assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (4,) # pylint:disable=comparison-with-callable + assert res[1].shape == (4,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -511,23 +527,24 @@ def circuit(x, y): assert jac[1][1].shape == (4,) assert np.allclose(jac[1][1], expected_1[1], atol=tol, rtol=0) - def test_diff_expval_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -544,11 +561,11 @@ def circuit(x, y): assert len(res) == 2 assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == () # pylint:disable=comparison-with-callable + assert res[0].shape == () assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) # pylint:disable=comparison-with-callable + assert res[1].shape == (2,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -582,28 +599,32 @@ def circuit(x, y): assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) def test_diff_expval_probs_sub_argnums( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs with less trainable parameters (argnums) than parameters.""" kwargs = {} - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "spsa": tol = TOL_FOR_SPSA + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) - if diff_method == "adjoint": - x = x + 0j - y = y + 0j - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -612,8 +633,6 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1]) - if "lightning" in dev.name: - pytest.xfail("lightning does not support measuring probabilities with adjoint.") jac = jax.jacobian(circuit, argnums=[0])(x, y) expected = [ @@ -638,24 +657,21 @@ def circuit(x, y): assert jac[1][0].shape == (2,) assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - } - - if diff_method == "hadamard": + kwargs = dict( + diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + if diff_method == "adjoint": + pytest.skip("Adjoint does not support probs") + elif diff_method == "hadamard": pytest.skip("Hadamard does not support var") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=3) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -674,11 +690,11 @@ def circuit(x, y): ] assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == () # pylint:disable=comparison-with-callable + assert res[0].shape == () assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) # pylint:disable=comparison-with-callable + assert res[1].shape == (2,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -718,40 +734,53 @@ class TestShotsIntegration: remains differentiable.""" def test_diff_method_None(self, interface): - """Test device works with diff_method=None.""" + """Test jax device works with diff_method=None.""" + dev = qml.device("default.qubit.jax", wires=1, shots=10) - @qml.qnode(DefaultQubit(), diff_method=None, interface=interface) + @qml.qnode(dev, diff_method=None, interface=interface) def circuit(x): qml.RX(x, wires=0) return qml.expval(qml.PauliZ(0)) - assert jax.numpy.allclose(circuit(jax.numpy.array(0.0), shots=10), 1) + assert jax.numpy.allclose(circuit(jax.numpy.array(0.0)), 1) - def test_changing_shots(self, interface): + def test_changing_shots(self, interface, mocker, tol): """Test that changing shots works on execution""" + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) + @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - circuit(a, b) + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 - res = circuit(a, b, shots=100) - assert res.shape == (100, 2) # pylint: disable=comparison-with-callable + res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(a, b) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # no additional calls def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" + dev = qml.device("default.qubit.legacy", wires=2, shots=1) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) + @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -759,40 +788,51 @@ def cost_fn(a, b): return qml.expval(qml.PauliY(1)) res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000) + assert dev.shots == 1 expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] assert np.allclose(res, expected, atol=0.1, rtol=0) - def test_update_diff_method(self, interface, mocker): + def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" + # pylint: disable=unused-argument + dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = jax.numpy.array([0.543, -0.654]) spy = mocker.spy(qml, "execute") - @qnode(DefaultQubit(), interface=interface) + # We're choosing interface="jax" such that backprop can be used in the + # test later + @qnode(dev, interface="jax") def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn == qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift - # if we use the default shots value of None, backprop can now be used cost_fn(a, b) - assert cost_fn.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift + + # if we set the shots to None, backprop can now be used + cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert cost_fn.gradient_fn == "backprop" + + cost_fn(a, b) + assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift -@pytest.mark.parametrize("dev,diff_method,grad_on_execution, device_vjp", device_and_diff_method) +@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp): + def test_sampling(self, dev_name, diff_method, grad_on_execution): """Test sampling works as expected""" if grad_on_execution is True: pytest.skip("Sampling not possible with grad_on_execution differentiation.") @@ -800,48 +840,43 @@ def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp): if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - @qnode( - dev, - diff_method=diff_method, - interface="jax", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + dev = qml.device(dev_name, wires=2, shots=10) + + @qnode(dev, diff_method=diff_method, interface="jax", grad_on_execution=grad_on_execution) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert isinstance(res, tuple) assert isinstance(res[0], jax.Array) - assert res[0].shape == (10,) # pylint:disable=comparison-with-callable + assert res[0].shape == (10,) assert isinstance(res[1], jax.Array) - assert res[1].shape == (10,) # pylint:disable=comparison-with-callable + assert res[1].shape == (10,) - def test_counts(self, dev, diff_method, grad_on_execution, device_vjp): + def test_counts(self, dev_name, diff_method, grad_on_execution): """Test counts works as expected""" if grad_on_execution is True: pytest.skip("Sampling not possible with grad_on_execution differentiation.") if diff_method == "adjoint": - pytest.skip("Adjoint errors with finite shots") + pytest.skip("Adjoint warns with finite shots") - @qnode( - dev, - diff_method=diff_method, - interface="jax", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + dev = qml.device(dev_name, wires=2, shots=10) + + @qnode(dev, diff_method=diff_method, interface="jax", grad_on_execution=grad_on_execution) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) - return qml.counts(qml.PauliZ(0)), qml.counts(qml.PauliX(1)) + return ( + qml.counts(qml.PauliZ(0), all_outcomes=True), + qml.counts(qml.PauliX(1), all_outcomes=True), + ) - res = circuit(shots=10) + res = circuit() assert isinstance(res, tuple) @@ -850,32 +885,25 @@ def circuit(): assert isinstance(res[1], dict) assert len(res[1]) == 2 - def test_chained_qnodes(self, dev, diff_method, grad_on_execution, device_vjp): + def test_chained_qnodes(self, dev_name, diff_method, grad_on_execution): """Test that the gradient of chained QNodes works without error""" - # pylint:disable=too-few-public-methods + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] - @qnode( - dev, - interface="jax", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="jax", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) - @qnode( - dev, - interface="jax", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + @qnode(dev, interface="jax", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit2(data, weights): qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) Template(weights, wires=[0, 1]) @@ -900,72 +928,37 @@ def cost(weights): assert len(res) == 2 - def test_postselection_differentiation(self, dev, diff_method, grad_on_execution, device_vjp): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") - - @qml.qnode( - dev, - diff_method=diff_method, - interface="jax", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface="jax", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - phi = jax.numpy.array(1.23) - theta = jax.numpy.array(4.56) - - assert np.allclose(circuit(phi, theta), expected_circuit(theta)) - - gradient = jax.grad(circuit, argnums=[0, 1])(phi, theta) - exp_theta_grad = jax.grad(expected_circuit)(theta) - assert np.allclose(gradient, [0.0, exp_theta_grad]) - @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegrationHigherOrder: """Tests that ensure various qubit circuits integrate correctly when computing higher-order derivatives""" - def test_second_derivative( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_second_derivative(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test second derivative calculation of a scalar-valued QNode""" - kwargs = { - "diff_method": diff_method, - "interface": interface, - "grad_on_execution": grad_on_execution, - "device_vjp": device_vjp, - "max_diff": 2, - } - + kwargs = dict( + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + max_diff=2, + ) if diff_method == "adjoint": pytest.skip("Adjoint does not second derivative.") elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + spsa_kwargs = dict( + sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=100, h=0.001 + ) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(x): qml.RY(x[0], wires=0) @@ -994,7 +987,7 @@ def circuit(x): else: assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test hessian calculation of a scalar-valued QNode""" gradient_kwargs = {} if diff_method == "adjoint": @@ -1007,12 +1000,18 @@ def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interfac } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1046,15 +1045,12 @@ def circuit(x): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test hessian calculation of a vector-valued QNode""" gradient_kwargs = {} if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": - qml.math.random.seed(42) gradient_kwargs = { "h": H_FOR_SPSA, "num_directions": 20, @@ -1062,12 +1058,18 @@ def test_hessian_vector_valued( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1111,7 +1113,7 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, dev, diff_method, interface, grad_on_execution, device_vjp, tol + self, dev_name, diff_method, interface, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" gradient_kwargs = {} @@ -1125,12 +1127,18 @@ def test_hessian_vector_valued_postprocessing( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1177,7 +1185,7 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + self, dev_name, diff_method, grad_on_execution, interface, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" gradient_kwargs = {} @@ -1191,12 +1199,18 @@ def test_hessian_vector_valued_separate_args( } tol = TOL_FOR_SPSA + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1222,7 +1236,6 @@ def circuit(a, b): ] ) assert np.allclose(g, expected_g.T, atol=tol, rtol=0) - hess = jax.jacobian(jac_fn, argnums=[0, 1])(a, b) expected_hess = np.array( @@ -1242,21 +1255,23 @@ def circuit(a, b): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): + def test_state(self, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("Lightning does not support state adjoint differentiation.") + dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1266,7 +1281,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is np.dtype("complex128") # pylint:disable=no-member + assert res.dtype is np.dtype("complex128") probs = jax.numpy.abs(res) ** 2 return probs[0] + probs[2] @@ -1280,19 +1295,18 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector( - self, state, dev, diff_method, grad_on_execution, device_vjp, interface, tol - ): + def test_projector(self, state, dev_name, diff_method, grad_on_execution, interface, tol): """Test that the variance of a projector is correctly returned""" gradient_kwargs = {} if diff_method == "adjoint": - pytest.skip("adjoint supports all expvals or only diagonal measurements.") - if diff_method == "hadamard": + pytest.skip("Adjoint does not support projectors") + elif diff_method == "hadamard": pytest.skip("Hadamard does not support var.") elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=2) P = jax.numpy.array(state) x, y = 0.765, -0.654 @@ -1301,7 +1315,6 @@ def test_projector( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1325,7 +1338,94 @@ def circuit(x, y): @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": H_FOR_SPSA}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +@pytest.mark.parametrize("interface", ["jax", "jax-python"]) +class TestCV: + """Tests for CV integration""" + + def test_first_order_observable(self, diff_method, kwargs, interface, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + r = 0.543 + phi = -0.654 + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + res = circuit(r, phi) + expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = jax.grad(circuit, argnums=[0, 1])(r, phi) + expected = np.array( + [ + 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, + 2 * np.sinh(2 * r) * np.sin(2 * phi), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_second_order_observable(self, diff_method, kwargs, interface, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + n = 0.12 + a = 0.765 + + @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) + + res = circuit(n, a) + expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res = jax.grad(circuit, argnums=[0, 1])(n, a) + expected = np.array([2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]) + assert np.allclose(res, expected, atol=tol, rtol=0) + + +@pytest.mark.parametrize("interface", ["auto", "jax", "jax-python"]) +def test_adjoint_reuse_device_state(mocker, interface): + """Tests that the jax interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=1) + + @qnode(dev, interface=interface, diff_method="adjoint") + def circ(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + jax.grad(circ)(1.0) + assert circ.device.num_executions == 1 + + spy.assert_called_with(mocker.ANY, use_device_state=True) + + +@pytest.mark.parametrize( + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly @@ -1333,14 +1433,21 @@ class TestTapeExpansion: @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface + self, dev_name, diff_method, grad_on_execution, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + class PhaseShift(qml.PhaseShift): grad_method = None def decomposition(self): @@ -1352,7 +1459,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, - device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1368,7 +1474,7 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, max_diff, interface, device_vjp, mocker, tol + self, dev_name, diff_method, grad_on_execution, max_diff, interface, mocker, tol ): """Test that the Hamiltonian is not expanded if there are non-commuting groups and the number of shots is None @@ -1386,6 +1492,7 @@ def test_hamiltonian_expansion_analytic( } tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=3, shots=None) spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @@ -1394,7 +1501,6 @@ def test_hamiltonian_expansion_analytic( interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=max_diff, **gradient_kwargs, ) @@ -1439,11 +1545,11 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff, mocker + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, interface, max_diff, mocker ): - """Test that the Hamiltonian is correctly measured (and not expanded) - if there are non-commuting groups and the number of shots is finite + """Test that the Hamiltonian is expanded if there + are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} tol = 0.3 @@ -1459,6 +1565,7 @@ def test_hamiltonian_finite_shots( } tol = TOL_FOR_SPSA + dev = qml.device(dev_name, wires=3, shots=50000) spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @@ -1467,7 +1574,6 @@ def test_hamiltonian_finite_shots( interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, max_diff=max_diff, **gradient_kwargs, ) @@ -1484,13 +1590,13 @@ def circuit(data, weights, coeffs): c = jax.numpy.array([-0.6543, 0.24, 0.54]) # test output - res = circuit(d, w, c, shots=50000) + res = circuit(d, w, c) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_not_called() + spy.assert_called() # test gradients - grad = jax.grad(circuit, argnums=[1, 2])(d, w, c, shots=50000) + grad = jax.grad(circuit, argnums=[1, 2])(d, w, c) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1520,24 +1626,27 @@ def circuit(data, weights, coeffs): @pytest.mark.parametrize("shots", [None, 10000]) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) -class TestReturn: # pylint:disable=too-many-public-methods +class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" def test_grad_single_measurement_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """For one measurement and one param, the gradient is a float.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -1546,24 +1655,27 @@ def circuit(a): a = jax.numpy.array(0.1) - grad = jax.grad(circuit)(a, shots=shots) + grad = jax.grad(circuit)(a) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, shots, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -1573,7 +1685,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - grad = jax.grad(circuit, argnums=[0, 1])(a, b, shots=shots) + grad = jax.grad(circuit, argnums=[0, 1])(a, b) assert isinstance(grad, tuple) assert len(grad) == 2 @@ -1581,18 +1693,21 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, shots, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1601,14 +1716,14 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - grad = jax.grad(circuit)(a, shots=shots) + grad = jax.grad(circuit)(a) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == (2,) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" @@ -1618,12 +1733,15 @@ def test_jacobian_single_measurement_param_probs( if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -1632,14 +1750,14 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4,) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1648,12 +1766,15 @@ def test_jacobian_single_measurement_probs_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -1663,7 +1784,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - jac = jacobian(circuit, argnums=[0, 1])(a, b, shots=shots) + jac = jacobian(circuit, argnums=[0, 1])(a, b) assert isinstance(jac, tuple) @@ -1675,7 +1796,7 @@ def circuit(a, b): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1684,12 +1805,15 @@ def test_jacobian_single_measurement_probs_multiple_param_single_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1697,29 +1821,26 @@ def circuit(a): return qml.probs(wires=[0, 1]) a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4, 2) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_expval_expval_multiple_params( - self, - dev, - diff_method, - grad_on_execution, - jacobian, - shots, - interface, - device_vjp, + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if device_vjp and jacobian is jax.jacfwd: - pytest.skip("forward pass can't be done with registered vjp.") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -1729,7 +1850,6 @@ def test_jacobian_expval_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1737,7 +1857,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(jac, tuple) @@ -1757,22 +1877,21 @@ def circuit(x, y): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if device_vjp and jacobian is jax.jacfwd: - pytest.skip("forward pass can't be done with registered vjp.") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1781,7 +1900,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1794,16 +1913,18 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_var_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": - pytest.skip("adjoint supports either all measurements or only diagonal measurements.") - if diff_method == "hadamard": + pytest.skip("Test does not supports adjoint because of var.") + elif diff_method == "hadamard": pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -1813,7 +1934,6 @@ def test_jacobian_var_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1821,7 +1941,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1842,22 +1962,20 @@ def circuit(x, y): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_var_var_multiple_params_array( - self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or all diagonal measurements.") - if diff_method == "hadamard": + pytest.skip("Test does not supports adjoint because of var.") + elif diff_method == "hadamard": pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1866,7 +1984,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1879,22 +1997,24 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a single params return an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - if diff_method == "adjoint" and jacobian == jax.jacfwd: - pytest.skip("jacfwd doesn't like complex numbers") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -1903,7 +2023,7 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1916,32 +2036,33 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - if diff_method == "adjoint" and jacobian == jax.jacfwd: - pytest.skip("jacfwd doesn't like complex numbers") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) - jac = jacobian(circuit, argnums=[0, 1])(a, b, shots=shots) + jac = jacobian(circuit, argnums=[0, 1])(a, b) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1962,22 +2083,23 @@ def circuit(a, b): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - if "lightning" in dev.name: - pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - if diff_method == "adjoint" and jacobian == jax.jacfwd: - pytest.skip("jacfwd doesn't like complex numbers") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1986,7 +2108,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a, shots=shots) + jac = jacobian(circuit)(a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1998,12 +2120,19 @@ def circuit(a): assert jac[1].shape == (4, 2) def test_hessian_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2016,7 +2145,6 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2024,7 +2152,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2042,7 +2170,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2050,6 +2178,13 @@ def test_hessian_expval_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2058,7 +2193,6 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2066,13 +2200,13 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit)(params, shots=shots) + hess = jax.hessian(circuit)(params) assert isinstance(hess, jax.numpy.ndarray) assert hess.shape == (2, 2) def test_hessian_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2081,6 +2215,7 @@ def test_hessian_var_multiple_params( pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -2091,7 +2226,6 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2099,7 +2233,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2117,7 +2251,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2127,6 +2261,8 @@ def test_hessian_var_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2135,7 +2271,6 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2143,15 +2278,22 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit)(params, shots=shots) + hess = jax.hessian(circuit)(params) assert isinstance(hess, jax.numpy.ndarray) assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2169,7 +2311,6 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2177,7 +2318,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2213,7 +2354,7 @@ def circuit(x, y): assert hess[1][1][1].shape == (2,) def test_hessian_expval_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2223,13 +2364,19 @@ def test_hessian_expval_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + params = jax.numpy.array([0.1, 0.2]) @qnode( dev, interface=interface, diff_method=diff_method, - device_vjp=device_vjp, max_diff=2, grad_on_execution=grad_on_execution, ) @@ -2239,7 +2386,7 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit)(params, shots=shots) + hess = jax.hessian(circuit)(params) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2251,7 +2398,7 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2261,6 +2408,8 @@ def test_hessian_probs_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = qml.numpy.array(0.1) par_1 = qml.numpy.array(0.2) @@ -2270,7 +2419,6 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2278,7 +2426,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2314,7 +2462,7 @@ def circuit(x, y): assert hess[1][1][1].shape == (2,) def test_hessian_var_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots, interface + self, dev_name, diff_method, grad_on_execution, shots, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2324,6 +2472,8 @@ def test_hessian_var_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2332,7 +2482,6 @@ def test_hessian_var_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2340,7 +2489,7 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit)(params, shots=shots) + hess = jax.hessian(circuit)(params) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2352,11 +2501,14 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) -def test_no_ops(): +@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) +def test_no_ops(dev_name): """Test that the return value of the QNode matches in the interface even if there are no ops""" - @qml.qnode(DefaultQubit(), interface="jax") + dev = qml.device(dev_name, wires=1) + + @qml.qnode(dev, interface="jax") def circuit(): qml.Hadamard(wires=0) return qml.state() diff --git a/tests/interfaces/legacy_devices_integration/test_jax_qnode_shot_vector_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_qnode_shot_vector_legacy.py new file mode 100644 index 00000000000..0ef5c8291de --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_jax_qnode_shot_vector_legacy.py @@ -0,0 +1,933 @@ +# Copyright 2022 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. +"""Integration tests for using the jax interface with shot vectors and with a QNode""" +# pylint: disable=too-many-arguments,too-many-public-methods +import pytest +from flaky import flaky + +import pennylane as qml +from pennylane import numpy as np +from pennylane import qnode + +pytestmark = pytest.mark.jax + +jax = pytest.importorskip("jax") +jax.config.update("jax_enable_x64", True) + +all_shots = [(1, 20, 100), (1, (20, 1), 100), (1, (5, 4), 100)] + +qubit_device_and_diff_method = [ + ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], + ["default.qubit.legacy", "parameter-shift", {}], + ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}], +] + +interface_and_qubit_device_and_diff_method = [ + ["jax"] + inner_list for inner_list in qubit_device_and_diff_method +] + +TOLS = { + "finite-diff": 0.3, + "parameter-shift": 1e-2, + "spsa": 0.32, +} + +jacobian_fn = [jax.jacobian, jax.jacrev, jax.jacfwd] + + +@pytest.mark.parametrize("shots", all_shots) +@pytest.mark.parametrize( + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method +) +class TestReturnWithShotVectors: + """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jac_single_measurement_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For one measurement and one param, the gradient is a float.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array(0.1) + + jac = jacobian(circuit)(a) + + assert isinstance(jac, tuple) + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, jax.numpy.ndarray) + assert j.shape == () + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jac_single_measurement_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For one measurement and multiple param, the gradient is a tuple of arrays.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + jac = jacobian(circuit, argnums=[0, 1])(a, b) + + assert isinstance(jac, tuple) + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 + assert j[0].shape == () + assert j[1].shape == () + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_single_measurement_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For one measurement and multiple param as a single array params, the gradient is an array.""" + dev = qml.device(dev_name, wires=1, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, jax.numpy.ndarray) + assert j.shape == (2,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_single_measurement_param_probs( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For a multi dimensional measurement (probs), check that a single array is returned with the correct + dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.probs(wires=[0, 1]) + + a = jax.numpy.array(0.1) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, jax.numpy.ndarray) + assert j.shape == (4,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_single_measurement_probs_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with + the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.probs(wires=[0, 1]) + + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) + + jac = jacobian(circuit, argnums=[0, 1])(a, b) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == (4,) + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (4,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_single_measurement_probs_multiple_param_single_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with + the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.probs(wires=[0, 1]) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, jax.numpy.ndarray) + assert j.shape == (4, 2) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_expval_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=1, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) + + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + + assert isinstance(j[0], tuple) + assert len(j[0]) == 2 + assert isinstance(j[0][0], jax.numpy.ndarray) + assert j[0][0].shape == () + assert isinstance(j[0][1], jax.numpy.ndarray) + assert j[0][1].shape == () + + assert isinstance(j[1], tuple) + assert len(j[1]) == 2 + assert isinstance(j[1][0], jax.numpy.ndarray) + assert j[1][0].shape == () + assert isinstance(j[1][1], jax.numpy.ndarray) + assert j[1][1].shape == () + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_expval_expval_multiple_params_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 # measurements + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == (2,) + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (2,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_var_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) + + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 + + assert isinstance(j[0], tuple) + assert len(j[0]) == 2 + assert isinstance(j[0][0], jax.numpy.ndarray) + assert j[0][0].shape == () + assert isinstance(j[0][1], jax.numpy.ndarray) + assert j[0][1].shape == () + + assert isinstance(j[1], tuple) + assert len(j[1]) == 2 + assert isinstance(j[1][0], jax.numpy.ndarray) + assert j[1][0].shape == () + assert isinstance(j[1][1], jax.numpy.ndarray) + assert j[1][1].shape == () + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_var_var_multiple_params_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 # measurements + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == (2,) + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (2,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_multiple_measurement_single_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a single params return an array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = jax.numpy.array(0.1) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(jac, tuple) + assert len(j) == 2 + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == () + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (4,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_multiple_measurement_multiple_param( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = np.array(0.1, requires_grad=True) + b = np.array(0.2, requires_grad=True) + + jac = jacobian(circuit, argnums=[0, 1])(a, b) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 + + assert isinstance(j[0], tuple) + assert len(j[0]) == 2 + assert isinstance(j[0][0], jax.numpy.ndarray) + assert j[0][0].shape == () + assert isinstance(j[0][1], jax.numpy.ndarray) + assert j[0][1].shape == () + + assert isinstance(j[1], tuple) + assert len(j[1]) == 2 + assert isinstance(j[1][0], jax.numpy.ndarray) + assert j[1][0].shape == (4,) + assert isinstance(j[1][1], jax.numpy.ndarray) + assert j[1][1].shape == (4,) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_jacobian_multiple_measurement_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(a): + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + a = jax.numpy.array([0.1, 0.2]) + + jac = jacobian(circuit)(a) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(jac) == num_copies + for j in jac: + assert isinstance(j, tuple) + assert len(j) == 2 # measurements + + assert isinstance(j[0], jax.numpy.ndarray) + assert j[0].shape == (2,) + + assert isinstance(j[1], jax.numpy.ndarray) + assert j[1].shape == (4, 2) + + def test_hessian_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of single a measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(hess, tuple) + assert len(h) == 2 + + assert isinstance(h[0], tuple) + assert len(h[0]) == 2 + assert isinstance(h[0][0], jax.numpy.ndarray) + assert h[0][0].shape == () + assert h[0][1].shape == () + + assert isinstance(h[1], tuple) + assert len(h[1]) == 2 + assert isinstance(h[1][0], jax.numpy.ndarray) + assert h[1][0].shape == () + assert h[1][1].shape == () + + def test_hessian_expval_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = jax.numpy.array([0.1, 0.2]) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + hess = jax.hessian(circuit)(params) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, jax.numpy.ndarray) + assert h.shape == (2, 2) + + def test_hessian_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of single a measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], tuple) + assert len(h[0]) == 2 + assert isinstance(h[0][0], jax.numpy.ndarray) + assert h[0][0].shape == () + assert h[0][1].shape == () + + assert isinstance(h[1], tuple) + assert len(h[1]) == 2 + assert isinstance(h[1][0], jax.numpy.ndarray) + assert h[1][0].shape == () + assert h[1][1].shape == () + + def test_hessian_var_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = jax.numpy.array([0.1, 0.2]) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + hess = jax.hessian(circuit)(params) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, jax.numpy.ndarray) + assert h.shape == (2, 2) + + def test_hessian_probs_expval_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = jax.numpy.array(0.1) + par_1 = jax.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], tuple) + assert len(h[0]) == 2 + assert isinstance(h[0][0], tuple) + assert len(h[0][0]) == 2 + assert isinstance(h[0][0][0], jax.numpy.ndarray) + assert h[0][0][0].shape == () + assert isinstance(h[0][0][1], jax.numpy.ndarray) + assert h[0][0][1].shape == () + assert isinstance(h[0][1], tuple) + assert len(h[0][1]) == 2 + assert isinstance(h[0][1][0], jax.numpy.ndarray) + assert h[0][1][0].shape == () + assert isinstance(h[0][1][1], jax.numpy.ndarray) + assert h[0][1][1].shape == () + + assert isinstance(h[1], tuple) + assert len(h[1]) == 2 + assert isinstance(h[1][0], tuple) + assert len(h[1][0]) == 2 + assert isinstance(h[1][0][0], jax.numpy.ndarray) + assert h[1][0][0].shape == (2,) + assert isinstance(h[1][0][1], jax.numpy.ndarray) + assert h[1][0][1].shape == (2,) + assert isinstance(h[1][1], tuple) + assert len(h[1][1]) == 2 + assert isinstance(h[1][1][0], jax.numpy.ndarray) + assert h[1][1][0].shape == (2,) + assert isinstance(h[1][1][1], jax.numpy.ndarray) + assert h[1][1][1].shape == (2,) + + def test_hessian_expval_probs_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of multiple measurements with a multiple param array return a single array.""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because second order diff.") + + dev = qml.device(dev_name, wires=2, shots=shots) + + params = jax.numpy.array([0.1, 0.2]) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + hess = jax.hessian(circuit)(params) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], jax.numpy.ndarray) + assert h[0].shape == (2, 2) + + assert isinstance(h[1], jax.numpy.ndarray) + assert h[1].shape == (2, 2, 2) + + def test_hessian_probs_var_multiple_params( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + par_0 = qml.numpy.array(0.1) + par_1 = qml.numpy.array(0.2) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], tuple) + assert len(h[0]) == 2 + assert isinstance(h[0][0], tuple) + assert len(h[0][0]) == 2 + assert isinstance(h[0][0][0], jax.numpy.ndarray) + assert h[0][0][0].shape == () + assert isinstance(h[0][0][1], jax.numpy.ndarray) + assert h[0][0][1].shape == () + assert isinstance(h[0][1], tuple) + assert len(h[0][1]) == 2 + assert isinstance(h[0][1][0], jax.numpy.ndarray) + assert h[0][1][0].shape == () + assert isinstance(h[0][1][1], jax.numpy.ndarray) + assert h[0][1][1].shape == () + + assert isinstance(h[1], tuple) + assert len(h[1]) == 2 + assert isinstance(h[1][0], tuple) + assert len(h[1][0]) == 2 + assert isinstance(h[1][0][0], jax.numpy.ndarray) + assert h[1][0][0].shape == (2,) + assert isinstance(h[1][0][1], jax.numpy.ndarray) + assert h[1][0][1].shape == (2,) + assert isinstance(h[1][1], tuple) + assert len(h[1][1]) == 2 + assert isinstance(h[1][1][0], jax.numpy.ndarray) + assert h[1][1][0].shape == (2,) + assert isinstance(h[1][1][1], jax.numpy.ndarray) + assert h[1][1][1].shape == (2,) + + def test_hessian_var_probs_multiple_param_array( + self, dev_name, diff_method, gradient_kwargs, shots, interface + ): + """The hessian of multiple measurements with a multiple param array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + + params = jax.numpy.array([0.1, 0.2]) + + @qnode(dev, interface=interface, diff_method=diff_method, max_diff=2, **gradient_kwargs) + def circuit(x): + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) + + hess = jax.hessian(circuit)(params) + + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(hess) == num_copies + for h in hess: + assert isinstance(h, tuple) + assert len(h) == 2 + + assert isinstance(h[0], jax.numpy.ndarray) + assert h[0].shape == (2, 2) + + assert isinstance(h[1], jax.numpy.ndarray) + assert h[1].shape == (2, 2, 2) + + +@pytest.mark.parametrize("shots", all_shots) +class TestReturnShotVectorsDevice: + """Test for shot vectors with device method adjoint_jacobian.""" + + def test_jac_adjoint_fwd_error(self, shots): + """Test that an error is raised for adjoint forward.""" + dev = qml.device("default.qubit.legacy", wires=1, shots=shots) + + with pytest.warns( + UserWarning, match="Requested adjoint differentiation to be computed with finite shots." + ): + + @qnode(dev, interface="jax", diff_method="adjoint", grad_on_execution=True) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array(0.1) + + if isinstance(shots, tuple): + with pytest.raises( + qml.QuantumFunctionError, + match="Adjoint does not support shot vectors.", + ): + jax.jacobian(circuit)(a) + + def test_jac_adjoint_bwd_error(self, shots): + """Test that an error is raised for adjoint backward.""" + dev = qml.device("default.qubit.legacy", wires=1, shots=shots) + + with pytest.warns( + UserWarning, match="Requested adjoint differentiation to be computed with finite shots." + ): + + @qnode(dev, interface="jax", diff_method="adjoint", grad_on_execution=False) + def circuit(a): + qml.RY(a, wires=0) + qml.RX(0.2, wires=0) + return qml.expval(qml.PauliZ(0)) + + a = jax.numpy.array(0.1) + + with pytest.raises( + qml.QuantumFunctionError, + match="Adjoint does not support shot vectors.", + ): + jax.jacobian(circuit)(a) + + +qubit_device_and_diff_method = [ + ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], + ["default.qubit.legacy", "parameter-shift", {}], +] + +shots_large = [(1000000, 900000, 800000), (1000000, (900000, 2))] + + +@flaky(max_runs=5) +@pytest.mark.parametrize("shots", shots_large) +@pytest.mark.parametrize( + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method +) +class TestReturnShotVectorIntegration: + """Tests for the integration of shots with the Jax interface.""" + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_single_expectation_value( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """Tests correct output shape and evaluation for a tape + with a single expval output""" + dev = qml.device(dev_name, wires=2, shots=shots) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) + + expected = np.array([[-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]]) + all_res = jacobian(circuit, argnums=[0, 1])(x, y) + + assert isinstance(all_res, tuple) + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(all_res) == num_copies + + for res in all_res: + assert isinstance(res[0], jax.numpy.ndarray) + assert res[0].shape == () + + assert isinstance(res[1], jax.numpy.ndarray) + assert res[1].shape == () + tol = TOLS[diff_method] + assert np.allclose(res, expected, atol=tol, rtol=0) + + @pytest.mark.parametrize("jacobian", jacobian_fn) + def test_prob_expectation_values( + self, dev_name, diff_method, gradient_kwargs, shots, jacobian, interface + ): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + dev = qml.device(dev_name, wires=2, shots=shots) + x = jax.numpy.array(0.543) + y = jax.numpy.array(-0.654) + + @qnode(dev, interface=interface, diff_method=diff_method, **gradient_kwargs) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) + + all_res = jacobian(circuit, argnums=[0, 1])(x, y) + + tol = TOLS[diff_method] + + assert isinstance(all_res, tuple) + num_copies = sum( + [1 for x in shots if isinstance(x, int)] + [x[1] for x in shots if isinstance(x, tuple)] + ) + assert len(all_res) == num_copies + + for res in all_res: + assert isinstance(res, tuple) + assert len(res) == 2 + + assert isinstance(res[0], tuple) + assert len(res[0]) == 2 + assert np.allclose(res[0][0], -np.sin(x), atol=tol, rtol=0) + assert isinstance(res[0][0], jax.numpy.ndarray) + assert np.allclose(res[0][1], 0, atol=tol, rtol=0) + assert isinstance(res[0][1], jax.numpy.ndarray) + + assert isinstance(res[1], tuple) + assert len(res[1]) == 2 + assert np.allclose( + res[1][0], + [ + -(np.cos(y / 2) ** 2 * np.sin(x)) / 2, + -(np.sin(x) * np.sin(y / 2) ** 2) / 2, + (np.sin(x) * np.sin(y / 2) ** 2) / 2, + (np.cos(y / 2) ** 2 * np.sin(x)) / 2, + ], + atol=tol, + rtol=0, + ) + assert isinstance(res[1][0], jax.numpy.ndarray) + assert np.allclose( + res[1][1], + [ + -(np.cos(x / 2) ** 2 * np.sin(y)) / 2, + (np.cos(x / 2) ** 2 * np.sin(y)) / 2, + (np.sin(x / 2) ** 2 * np.sin(y)) / 2, + -(np.sin(x / 2) ** 2 * np.sin(y)) / 2, + ], + atol=tol, + rtol=0, + ) + assert isinstance(res[1][1], jax.numpy.ndarray) diff --git a/tests/interfaces/legacy_devices_integration/test_set_shots_legacy.py b/tests/interfaces/legacy_devices_integration/test_set_shots_legacy.py new file mode 100644 index 00000000000..6e9739a631f --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_set_shots_legacy.py @@ -0,0 +1,47 @@ +# Copyright 2018-2023 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. +""" +Tests for workflow.set_shots +""" + + +import pennylane as qml +from pennylane.measurements import Shots +from pennylane.workflow import set_shots + + +def test_set_with_shots_class(): + """Test that shots can be set on the old device interface with a Shots class.""" + + dev = qml.devices.DefaultQubitLegacy(wires=1) + with set_shots(dev, Shots(10)): + assert dev.shots == 10 + + assert dev.shots is None + + shot_tuples = Shots((10, 10)) + with set_shots(dev, shot_tuples): + assert dev.shots == 20 + assert dev.shot_vector == list(shot_tuples.shot_vector) + + assert dev.shots is None + + +def test_shots_not_altered_if_False(): + """Test a value of False can be passed to shots, indicating to not override + shots on the device.""" + + dev = qml.devices.DefaultQubitLegacy(wires=1) + with set_shots(dev, False): + assert dev.shots is None diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_tensorflow_autograph_qnode_shot_vector_legacy.py similarity index 78% rename from tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_tensorflow_autograph_qnode_shot_vector_legacy.py index c831fce26cc..f5b6e37a85b 100644 --- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_autograph_qnode_shot_vector_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_tensorflow_autograph_qnode_shot_vector_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2023 Xanadu Quantum Technologies Inc. +# Copyright 2022 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. @@ -12,13 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the TF interface with shot vectors and with a QNode""" -# pylint: disable=too-many-arguments,too-few-public-methods,unexpected-keyword-arg +# pylint: disable=too-many-arguments,too-few-public-methods,redefined-outer-name import pytest import pennylane as qml from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf @@ -27,25 +26,36 @@ shots_and_num_copies = [((1, (5, 2), 10), 4)] shots_and_num_copies_hess = [((10, (5, 1)), 2)] + +kwargs = { + "finite-diff": {"h": 10e-2}, + "parameter-shift": {}, + "spsa": {"h": 10e-2, "num_directions": 30}, +} + qubit_device_and_diff_method = [ - [DefaultQubit(), "finite-diff", {"h": 10e-2}], - [DefaultQubit(), "parameter-shift", {}], - [ - DefaultQubit(), - "spsa", - {"h": 10e-2, "num_directions": 20, "sampler_rng": np.random.default_rng(42)}, - ], + ["default.qubit.legacy", "finite-diff"], + ["default.qubit.legacy", "parameter-shift"], + ["default.qubit.legacy", "spsa"], ] TOLS = { "finite-diff": 0.3, "parameter-shift": 1e-2, - "spsa": 0.5, + "spsa": 0.3, } +@pytest.fixture +def gradient_kwargs(request): + diff_method = request.node.funcargs["diff_method"] + return kwargs[diff_method] | ( + {"sampler_rng": np.random.default_rng(42)} if diff_method == "spsa" else {} + ) + + @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) -@pytest.mark.parametrize("dev,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -54,13 +64,14 @@ class TestReturnWithShotVectors: """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" def test_jac_single_measurement_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and one param, the gradient is a float.""" + dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.expval(qml.PauliZ(0)) @@ -68,7 +79,7 @@ def circuit(a, **_): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -77,13 +88,14 @@ def circuit(a, **_): assert jac.shape == (num_copies,) def test_jac_single_measurement_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b, **_): + def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)) @@ -92,7 +104,7 @@ def circuit(a, b, **_): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -104,13 +116,14 @@ def circuit(a, b, **_): assert j.shape == (num_copies,) def test_jacobian_single_measurement_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) @@ -118,7 +131,7 @@ def circuit(a, **_): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -127,14 +140,15 @@ def circuit(a, **_): assert jac.shape == (num_copies, 2) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.probs(wires=[0, 1]) @@ -142,7 +156,7 @@ def circuit(a, **_): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -151,14 +165,15 @@ def circuit(a, **_): assert jac.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b, **_): + def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.probs(wires=[0, 1]) @@ -167,7 +182,7 @@ def circuit(a, b, **_): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -179,14 +194,15 @@ def circuit(a, b, **_): assert j.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) @@ -194,7 +210,7 @@ def circuit(a, **_): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -203,22 +219,24 @@ def circuit(a, **_): assert jac.shape == (num_copies, 4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The gradient of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = tf.Variable(1.5, dtype=tf.float64) par_1 = tf.Variable(0.7, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, max_diff=1, **gradient_kwargs) - def circuit(x, y, **_): + def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) with tf.GradientTape() as tape: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, (par_0, par_1)) @@ -230,13 +248,14 @@ def circuit(x, y, **_): assert j.shape == (num_copies, 2) def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) qml.RY(a[2], wires=0) @@ -245,7 +264,7 @@ def circuit(a, **_): a = tf.Variable([0.7, 0.9, 1.1], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, a) @@ -254,13 +273,14 @@ def circuit(a, **_): assert jac.shape == (num_copies, 2, 3) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a single params return an array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -268,7 +288,7 @@ def circuit(a, **_): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -277,13 +297,14 @@ def circuit(a, **_): assert jac.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b, **_): + def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -292,7 +313,7 @@ def circuit(a, b, **_): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, (a, b)) @@ -304,13 +325,14 @@ def circuit(a, b, **_): assert j.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, **_): + def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -318,7 +340,7 @@ def circuit(a, **_): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -329,7 +351,7 @@ def circuit(a, **_): @pytest.mark.slow @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) -@pytest.mark.parametrize("dev,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -338,7 +360,7 @@ class TestReturnShotVectorHessian: """Class to test the shape of the Hessian with different return types and shot vectors.""" def test_hessian_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The hessian of a single measurement with multiple params return a tuple of arrays.""" @@ -346,12 +368,14 @@ def test_hessian_expval_multiple_params( # TODO: Find out why. pytest.skip("SPSA gradient does not support this particular test case") + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = tf.Variable(1.5, dtype=tf.float64) par_1 = tf.Variable(0.7, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, max_diff=2, **gradient_kwargs) - def circuit(x, y, **_): + def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) @@ -359,7 +383,7 @@ def circuit(x, y, **_): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack(res) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -374,11 +398,11 @@ def circuit(x, y, **_): assert h.shape == (2, num_copies) -shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)] +shots_and_num_copies = [((20000, 18000, 16000), 3), ((20000, (18000, 2)), 3)] @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) -@pytest.mark.parametrize("dev,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev_name,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -387,23 +411,24 @@ class TestReturnShotVectorIntegration: """Tests for the integration of shots with the TF interface.""" def test_single_expectation_value( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """Tests correct output shape and evaluation for a tape with a single expval output""" + dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(x, y, **_): + def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) with tf.GradientTape() as tape: - res = circuit(x, y, shots=shots) + res = circuit(x, y) res = qml.math.stack(res) all_res = tape.jacobian(res, (x, y)) @@ -420,23 +445,24 @@ def circuit(x, y, **_): assert np.allclose(res, exp, atol=tol, rtol=0) def test_prob_expectation_values( - self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(x, y, **_): + def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) with tf.GradientTape() as tape: - res = circuit(x, y, shots=shots) + res = circuit(x, y) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) all_res = tape.jacobian(res, (x, y)) diff --git a/tests/interfaces/legacy_devices_integration/test_tensorflow_legacy.py b/tests/interfaces/legacy_devices_integration/test_tensorflow_legacy.py new file mode 100644 index 00000000000..057f93d0038 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_tensorflow_legacy.py @@ -0,0 +1,1028 @@ +# Copyright 2018-2021 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. +"""Unit tests for the TensorFlow interface""" +# pylint: disable=protected-access,too-few-public-methods +import numpy as np +import pytest + +import pennylane as qml +from pennylane import execute +from pennylane.gradients import finite_diff, param_shift + +pytestmark = pytest.mark.tf + +tf = pytest.importorskip("tensorflow", minversion="2.1") + + +class TestTensorFlowExecuteUnitTests: + """Unit tests for TensorFlow execution""" + + def test_jacobian_options(self, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = tf.Variable([0.1, 0.2], dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=1) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute( + [tape], + dev, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + interface="tf", + )[0] + + res = t.jacobian(res, a) + + for args in spy.call_args_list: + assert args[1]["shifts"] == [(np.pi / 4,)] * 2 + + def test_incorrect_grad_on_execution(self): + """Test that an error is raised if a gradient transform + is used with grad_on_execution""" + a = tf.Variable([0.1, 0.2]) + + dev = qml.device("default.qubit.legacy", wires=1) + + with tf.GradientTape(): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + execute([tape], dev, gradient_fn=param_shift, grad_on_execution=True, interface="tf") + + def test_grad_on_execution(self, mocker): + """Test that grad on execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + a = tf.Variable([0.1, 0.2]) + spy = mocker.spy(dev, "execute_and_gradients") + + with tf.GradientTape(): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + interface="tf", + ) + + # adjoint method only performs a single device execution, but gets both result and gradient + assert dev.num_executions == 1 + spy.assert_called() + + def test_no_grad_execution(self, mocker): + """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") + a = tf.Variable([0.1, 0.2]) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + interface="tf", + )[0] + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + t.jacobian(res, a) + spy_gradients.assert_called() + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + a = tf.Variable([0.1, 0.2]) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, gradient_fn=param_shift, cachesize=2, interface="tf")[0] + + t.jacobian(res, a) + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + a = tf.Variable([0.1, 0.2]) + custom_cache = {} + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, gradient_fn=param_shift, cache=custom_cache, interface="tf")[ + 0 + ] + + t.jacobian(res, a) + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + unwrapped_tape = qml.transforms.convert_to_numpy_parameters(tape)[0][0] + h = unwrapped_tape.hash + + assert h in cache + assert np.allclose(cache[h], res) + + def test_caching_param_shift(self): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum.""" + dev = qml.device("default.qubit.legacy", wires=1) + a = tf.Variable([0.1, 0.2], dtype=tf.float64) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="tf")[0] + + # Without caching, and non-vectorized, 9 evaluations are required to compute + # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) + with tf.GradientTape(persistent=True) as t: + res = cost(a, cache=None) + t.jacobian(res, a, experimental_use_pfor=False) + assert dev.num_executions == 9 + + # With caching, and non-vectorized, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + with tf.GradientTape(persistent=True) as t: + res = cost(a, cache=True) + t.jacobian(res, a) + assert dev.num_executions == 5 + + # In vectorized mode, 5 evaluations are required to compute + # the Jacobian regardless of caching: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + with tf.GradientTape() as t: + res = cost(a, cache=None) + t.jacobian(res, a) + assert dev.num_executions == 5 + + @pytest.mark.parametrize("num_params", [2, 3]) + def test_caching_param_shift_hessian(self, num_params, tol): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum + when computing Hessians.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = tf.Variable(np.arange(1, num_params + 1) / 10, dtype=tf.float64) + + N = params.shape[0] + + def cost(x, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + + for i in range(2, num_params): + qml.RZ(x[i], wires=[i % 2]) + + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + return execute( + [tape], dev, gradient_fn=param_shift, cache=cache, interface="tf", max_diff=2 + )[0] + + # No caching: number of executions is not ideal + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + res = cost(params, cache=False) + grad = t1.gradient(res, params) + hess1 = t2.jacobian(grad, params) + + if num_params == 2: + # compare to theoretical result + x, y, *_ = params * 1.0 + expected = np.array( + [ + [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(expected, hess1, atol=tol, rtol=0) + + nonideal_runs = dev.num_executions + + # Use caching: number of executions is ideal + dev._num_executions = 0 + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + res = cost(params, cache=True) + grad = t1.gradient(res, params) + hess2 = t2.jacobian(grad, params) + + assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + expected_runs_ideal = 1 # forward pass + expected_runs_ideal += 2 * N # Jacobian + expected_runs_ideal += N + 1 # Hessian diagonal + expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal + assert dev.num_executions == expected_runs_ideal + assert expected_runs_ideal < nonideal_runs + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift, "interface": "tf"}, + {"gradient_fn": param_shift, "interface": "auto"}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + "interface": "tf", + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + "interface": "tf", + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + "interface": "auto", + }, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + "interface": "auto", + }, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestTensorFlowExecuteIntegration: + """Test the TensorFlow interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, execute_kwargs): + """Test execution""" + dev = qml.device("default.qubit.legacy", wires=1) + a = tf.Variable(0.1) + b = tf.Variable(0.2) + + with tf.GradientTape(): + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + res = execute([tape1, tape2], dev, **execute_kwargs) + + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + assert isinstance(res[0], tf.Tensor) + assert isinstance(res[1], tf.Tensor) + + def test_scalar_jacobian(self, execute_kwargs, tol): + """Test scalar jacobian calculation""" + a = tf.Variable(0.1, dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + + res = t.jacobian(res, a) + assert res.shape == () + + # compare to standard tape jacobian + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(dev.batch_execute(tapes)) + + assert expected.shape == () + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_jacobian(self, execute_kwargs, tol): + """Test jacobian calculation""" + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, max_diff=2, **execute_kwargs)[0] + res = tf.stack(res) + + expected = [np.cos(a), -np.cos(a) * np.sin(b)] + assert np.allclose(res, expected, atol=tol, rtol=0) + + (agrad, bgrad) = t.jacobian(res, [a, b]) + assert agrad.shape == (2,) + assert bgrad.shape == (2,) + + expected = [[-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]] + assert np.allclose(expected, [agrad, bgrad], atol=tol, rtol=0) + + def test_tape_no_parameters(self, execute_kwargs, tol): + """Test that a tape with no parameters is correctly + ignored during the gradient computation""" + dev = qml.device("default.qubit.legacy", wires=1) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) + x, y = 1.0 * params + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.expval(qml.PauliX(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(0.5, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + with qml.queuing.AnnotatedQueue() as q3: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape3 = qml.tape.QuantumScript.from_queue(q3) + res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs)) + res = tf.stack(res) + + expected = 1 + np.cos(0.5) + np.cos(x) * np.cos(y) + assert np.allclose(res, expected, atol=tol, rtol=0) + + grad = t.gradient(res, params) + expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)] + assert np.allclose(grad, expected, atol=tol, rtol=0) + + def test_reusing_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + assert tape.trainable_params == [0, 1] + res = execute([tape], dev, **execute_kwargs)[0] + res = tf.stack(res) + + t.jacobian(res, [a, b]) + + a = tf.Variable(0.54, dtype=tf.float64) + b = tf.Variable(0.8, dtype=tf.float64) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + with tf.GradientTape() as t: + tape = tape.bind_new_parameters([2 * a, b], [0, 1]) + res2 = execute([tape], dev, **execute_kwargs)[0] + res2 = tf.stack(res2) + + expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac2 = t.jacobian(res2, [a, b]) + expected = [ + [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], + [0, -tf.cos(2 * a) * tf.cos(b)], + ] + assert np.allclose(jac2, expected, atol=tol, rtol=0) + + def test_reusing_pre_constructed_quantum_tape(self, execute_kwargs, tol): + """Test re-using a quantum tape that was previously constructed + *outside of* a gradient tape, by passing new parameters""" + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + with tf.GradientTape() as t: + tape = tape.bind_new_parameters([a, b], [0, 1]) + assert tape.trainable_params == [0, 1] + res = execute([tape], dev, **execute_kwargs)[0] + res = qml.math.stack(res) + + t.jacobian(res, [a, b]) + + a = tf.Variable(0.54, dtype=tf.float64) + b = tf.Variable(0.8, dtype=tf.float64) + + with tf.GradientTape() as t: + tape = tape.bind_new_parameters([2 * a, b], [0, 1]) + res2 = execute([tape], dev, **execute_kwargs)[0] + res2 = qml.math.stack(res2) + + expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] + assert np.allclose(res2, expected, atol=tol, rtol=0) + + jac2 = t.jacobian(res2, [a, b]) + expected = [ + [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], + [0, -tf.cos(2 * a) * tf.cos(b)], + ] + assert np.allclose(jac2, expected, atol=tol, rtol=0) + + def test_classical_processing(self, execute_kwargs): + """Test classical processing within the quantum tape""" + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.constant(0.2, dtype=tf.float64) + c = tf.Variable(0.3, dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=1) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a * c, wires=0) + qml.RZ(b, wires=0) + qml.RX(c + c**2 + tf.sin(a), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [0, 2] + assert tape.get_parameters() == [a * c, c + c**2 + tf.sin(a)] + + res = t.jacobian(res, [a, b, c]) + assert isinstance(res[0], tf.Tensor) + assert res[1] is None + assert isinstance(res[2], tf.Tensor) + + def test_no_trainable_parameters(self, execute_kwargs): + """Test evaluation and Jacobian if there are no trainable parameters""" + b = tf.constant(0.2, dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(0.2, wires=0) + qml.RX(b, wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + res = qml.math.stack(res) + + assert res.shape == (2,) + assert isinstance(res, tf.Tensor) + + res = t.jacobian(res, b) + assert res is None + + @pytest.mark.parametrize("U", [tf.constant([[0, 1], [1, 0]]), np.array([[0, 1], [1, 0]])]) + def test_matrix_parameter(self, execute_kwargs, U, tol): + """Test that the TF interface works correctly + with a matrix parameter""" + a = tf.Variable(0.1, dtype=tf.float64) + + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [1] + + assert np.allclose(res, -tf.cos(a), atol=tol, rtol=0) + + res = t.jacobian(res, a) + assert np.allclose(res, tf.sin(a), atol=tol, rtol=0) + + def test_differentiable_expand(self, execute_kwargs, tol): + """Test that operation and nested tape expansion + is differentiable""" + + class U3(qml.U3): + def decomposition(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + dev = qml.device("default.qubit.legacy", wires=1) + a = np.array(0.1) + p = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) + + with tf.GradientTape() as tape: + with qml.queuing.AnnotatedQueue() as q_qtape: + qml.RX(a, wires=0) + U3(p[0], p[1], p[2], wires=0) + qml.expval(qml.PauliX(0)) + + qtape = qml.tape.QuantumScript.from_queue(q_qtape) + res = execute([qtape], dev, **execute_kwargs)[0] + + expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( + tf.cos(p[2]) * tf.sin(p[1]) + tf.cos(p[0]) * tf.cos(p[1]) * tf.sin(p[2]) + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = tape.jacobian(res, p) + expected = np.array( + [ + tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), + tf.cos(p[1]) * tf.cos(p[2]) * tf.sin(a) + - tf.sin(p[1]) + * (tf.cos(a) * tf.sin(p[0]) + tf.cos(p[0]) * tf.sin(a) * tf.sin(p[2])), + tf.sin(a) + * (tf.cos(p[0]) * tf.cos(p[1]) * tf.cos(p[2]) - tf.sin(p[1]) * tf.sin(p[2])), + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_probability_differentiation(self, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob outputs""" + + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + x = tf.Variable(0.543, dtype=tf.float64) + y = tf.Variable(-0.654, dtype=tf.float64) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + res = qml.math.stack(res) + + expected = np.array( + [ + [tf.cos(x / 2) ** 2, tf.sin(x / 2) ** 2], + [(1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = t.jacobian(res, [x, y]) + expected = np.array( + [ + [ + [-tf.sin(x) / 2, tf.sin(x) / 2], + [-tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], + ], + [ + [0, 0], + [-tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], + ], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_ragged_differentiation(self, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + x = tf.Variable(0.543, dtype=tf.float64) + y = tf.Variable(-0.654, dtype=tf.float64) + + with tf.GradientTape() as t: + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + res = tf.experimental.numpy.hstack(res) + + expected = np.array( + [tf.cos(x), (1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = t.jacobian(res, [x, y]) + expected = np.array( + [ + [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], + [0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) + + def test_sampling(self, execute_kwargs): + """Test sampling works as expected""" + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + + dev = qml.device("default.qubit.legacy", wires=2, shots=10) + + with tf.GradientTape(): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(tf.Variable(0.1), wires=0) + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.sample(qml.PauliZ(0)) + qml.sample(qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute([tape], dev, **execute_kwargs)[0] + res = qml.math.stack(res) + + assert res.shape == (2, 10) + assert isinstance(res, tf.Tensor) + + +@pytest.mark.parametrize("interface", ["auto", "tf"]) +class TestHigherOrderDerivatives: + """Test that the TensorFlow execute function can be differentiated""" + + @pytest.mark.slow + @pytest.mark.parametrize( + "params", + [ + tf.Variable([0.543, -0.654], dtype=tf.float64), + tf.Variable([0, -0.654], dtype=tf.float64), + tf.Variable([-2.0, 0], dtype=tf.float64), + ], + ) + def test_parameter_shift_hessian(self, params, tol, interface): + """Tests that the output of the parameter-shift transform + can be differentiated using tensorflow, yielding second derivatives.""" + dev = qml.device("default.qubit.tf", wires=2) + x, y = params * 1.0 + + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(params[0], wires=[0]) + qml.RY(params[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(params[0], wires=0) + qml.RY(params[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = execute( + [tape1, tape2], dev, gradient_fn=param_shift, interface=interface, max_diff=2 + ) + res = result[0] + result[1][0] + + expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) + + grad = t1.gradient(res, params) + expected = np.array( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert np.allclose(grad, expected, atol=tol, rtol=0) + + hess = t2.jacobian(grad, params) + expected = np.array( + [ + [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(hess, expected, atol=tol, rtol=0) + + def test_hessian_vector_valued(self, tol, interface): + """Test hessian calculation of a vector valued QNode""" + dev = qml.device("default.qubit.tf", wires=1) + params = tf.Variable([0.543, -0.654], dtype=tf.float64) + + with tf.GradientTape() as t2: + with tf.GradientTape(persistent=True) as t1: + with qml.queuing.AnnotatedQueue() as q: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute( + [tape], dev, gradient_fn=param_shift, interface=interface, max_diff=2 + )[0] + res = tf.stack(res) + + g = t1.jacobian(res, params, experimental_use_pfor=False) + + hess = t2.jacobian(g, params) + + a, b = params * 1.0 + + expected_res = [ + 0.5 + 0.5 * tf.cos(a) * tf.cos(b), + 0.5 - 0.5 * tf.cos(a) * tf.cos(b), + ] + assert np.allclose(res, expected_res, atol=tol, rtol=0) + + expected_g = [ + [-0.5 * tf.sin(a) * tf.cos(b), -0.5 * tf.cos(a) * tf.sin(b)], + [0.5 * tf.sin(a) * tf.cos(b), 0.5 * tf.cos(a) * tf.sin(b)], + ] + assert np.allclose(g, expected_g, atol=tol, rtol=0) + + expected_hess = [ + [ + [-0.5 * tf.cos(a) * tf.cos(b), 0.5 * tf.sin(a) * tf.sin(b)], + [0.5 * tf.sin(a) * tf.sin(b), -0.5 * tf.cos(a) * tf.cos(b)], + ], + [ + [0.5 * tf.cos(a) * tf.cos(b), -0.5 * tf.sin(a) * tf.sin(b)], + [-0.5 * tf.sin(a) * tf.sin(b), 0.5 * tf.cos(a) * tf.cos(b)], + ], + ] + + np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) + + def test_adjoint_hessian(self, interface): + """Since the adjoint hessian is not a differentiable transform, + higher-order derivatives are not supported.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = tf.Variable([0.543, -0.654], dtype=tf.float64) + + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + with qml.queuing.AnnotatedQueue() as q: + qml.RX(params[0], wires=[0]) + qml.RY(params[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + res = execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + interface=interface, + )[0] + + grad = t1.gradient(res, params) + assert grad is not None + assert grad.dtype == tf.float64 + assert grad.shape == params.shape + + hess = t2.jacobian(grad, params) + assert hess is None + + def test_max_diff(self, tol, interface): + """Test that setting the max_diff parameter blocks higher-order + derivatives""" + dev = qml.device("default.qubit.tf", wires=2) + params = tf.Variable([0.543, -0.654], dtype=tf.float64) + x, y = params * 1.0 + + with tf.GradientTape() as t2: + with tf.GradientTape() as t1: + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(params[0], wires=[0]) + qml.RY(params[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(params[0], wires=0) + qml.RY(params[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + result = execute( + [tape1, tape2], dev, gradient_fn=param_shift, max_diff=1, interface=interface + ) + res = result[0] + result[1][0] + + expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) + + grad = t1.gradient(res, params) + + expected = np.array( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert np.allclose(grad, expected, atol=tol, rtol=0) + + hess = t2.jacobian(grad, params) + assert hess is None + + +execute_kwargs_hamiltonian = [ + {"gradient_fn": param_shift, "interface": "tensorflow"}, + {"gradient_fn": finite_diff, "interface": "tensorflow"}, + {"gradient_fn": param_shift, "interface": "auto"}, + {"gradient_fn": finite_diff, "interface": "auto"}, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +class TestHamiltonianWorkflows: + """Test that tapes ending with expectations + of Hamiltonians provide correct results and gradients""" + + @pytest.fixture + def cost_fn(self, execute_kwargs): + """Cost function for gradient tests""" + + def _cost_fn(weights, coeffs1, coeffs2, dev=None): + obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] + H1 = qml.Hamiltonian(coeffs1, obs1) + + obs2 = [qml.PauliZ(0)] + H2 = qml.Hamiltonian(coeffs2, obs2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(weights[0], wires=0) + qml.RY(weights[1], wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(H1) + qml.expval(H2) + + tape = qml.tape.QuantumScript.from_queue(q) + return tf.stack(execute([tape], dev, **execute_kwargs)[0]) + + return _cost_fn + + @staticmethod + def cost_fn_expected(weights, coeffs1, coeffs2): + """Analytic value of cost_fn above""" + a, b, c = coeffs1.numpy() + d = coeffs2.numpy()[0] + x, y = weights.numpy() + return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + + @staticmethod + def cost_fn_jacobian_expected(weights, coeffs1, coeffs2): + """Analytic jacobian of cost_fn above""" + a, b, c = coeffs1.numpy() + d = coeffs2.numpy()[0] + x, y = weights.numpy() + return np.array( + [ + [ + -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), + b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), + np.cos(x), + np.cos(x) * np.sin(y), + -(np.sin(x) * np.sin(y)), + 0, + ], + [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + ] + ) + + def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): + # pylint: disable=unused-argument + weights = tf.Variable([0.4, 0.5], dtype=tf.float64) + coeffs1 = tf.constant([0.1, 0.2, 0.3], dtype=tf.float64) + coeffs2 = tf.constant([0.7], dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as tape: + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = tape.jacobian(res, [weights, coeffs1, coeffs2]) + expected = self.cost_fn_jacobian_expected(weights, coeffs1, coeffs2) + assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0) + assert res[1] is None + assert res[2] is None + + def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): + # pylint: disable=unused-argument + weights = tf.Variable([0.4, 0.5], dtype=tf.float64) + coeffs1 = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) + coeffs2 = tf.Variable([0.7], dtype=tf.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + with tf.GradientTape() as tape: + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res, expected, atol=tol, rtol=0) + + res = tape.jacobian(res, [weights, coeffs1, coeffs2]) + expected = self.cost_fn_jacobian_expected(weights, coeffs1, coeffs2) + assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0) + assert np.allclose(res[1], expected[:, 2:5], atol=tol, rtol=0) + assert np.allclose(res[2], expected[:, 5:], atol=tol, rtol=0) diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_legacy.py similarity index 70% rename from tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_legacy.py index e3c4597ae06..5b03f9da10f 100644 --- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2018-2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2020 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. @@ -13,31 +13,26 @@ # limitations under the License. """Integration tests for using the TensorFlow interface with a QNode""" import numpy as np - -# pylint: disable=too-many-arguments,too-few-public-methods,comparison-with-callable, use-implicit-booleaness-not-comparison import pytest import pennylane as qml from pennylane import qnode -from pennylane.devices import DefaultQubit + +# pylint: disable=too-many-arguments,too-few-public-methods, use-dict-literal, use-implicit-booleaness-not-comparison + pytestmark = pytest.mark.tf tf = pytest.importorskip("tensorflow") -# device, diff_method, grad_on_execution, device_vjp + qubit_device_and_diff_method = [ - [DefaultQubit(), "finite-diff", False, False], - [DefaultQubit(), "parameter-shift", False, False], - [DefaultQubit(), "backprop", True, False], - [DefaultQubit(), "adjoint", True, False], - [DefaultQubit(), "adjoint", False, False], - [DefaultQubit(), "adjoint", False, True], - [DefaultQubit(), "spsa", False, False], - [DefaultQubit(), "hadamard", False, False], - [qml.device("lightning.qubit", wires=4), "adjoint", False, True], - [qml.device("lightning.qubit", wires=4), "adjoint", False, False], - [qml.device("lightning.qubit", wires=4), "adjoint", True, True], - [qml.device("lightning.qubit", wires=4), "adjoint", True, False], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] TOL_FOR_SPSA = 1.0 @@ -50,25 +45,25 @@ @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with TensorFlow integrates with the PennyLane stack""" - def test_execution_with_interface( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): + def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -89,7 +84,7 @@ def circuit(a): # with the interface, the tape returns tensorflow tensors assert isinstance(res, tf.Tensor) - assert res.shape == () + assert res.shape == tuple() # the tape is able to deduce trainable parameters assert circuit.qtape.trainable_params == [0] @@ -97,20 +92,23 @@ def circuit(a): # gradients should work grad = tape.gradient(res, a) assert isinstance(grad, tf.Tensor) - assert grad.shape == () + assert grad.shape == tuple() - def test_interface_swap(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_interface_swap(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test that the TF interface can be applied to a QNode with a pre-existing interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface="autograd", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -137,44 +135,51 @@ def circuit(a): assert np.allclose(res1, res2, atol=tol, rtol=0) assert np.allclose(grad1, grad2, atol=tol, rtol=0) - def test_drawing(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_drawing(self, dev_name, diff_method, grad_on_execution, interface): """Test circuit drawing when using the TF interface""" x = tf.Variable(0.1, dtype=tf.float64) y = tf.Variable([0.2, 0.3], dtype=tf.float64) z = tf.Variable(0.4, dtype=tf.float64) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(p1, p2=y, **kwargs): qml.RX(p1, wires=0) qml.RY(p2[0] * p2[1], wires=1) qml.RX(kwargs["p3"], wires=0) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) + return qml.state() result = qml.draw(circuit)(p1=x, p3=z) - expected = "0: ──RX(0.10)──RX(0.40)─╭●── \n1: ──RY(0.06)───────────╰X── " + expected = "0: ──RX(0.10)──RX(0.40)─╭●── State\n1: ──RY(0.06)───────────╰X── State" assert result == expected - def test_jacobian(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_jacobian(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test jacobian calculation""" - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + a = tf.Variable(0.1, dtype=tf.float64) b = tf.Variable(0.2, dtype=tf.float64) @@ -185,7 +190,7 @@ def circuit(a, b): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.stack(res) @@ -197,17 +202,61 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, [a, b]) expected = [[-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)]] assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_jacobian_dtype(self, dev_name, diff_method, grad_on_execution, interface): + """Test calculating the jacobian with a different datatype""" + if diff_method == "backprop": + pytest.skip("Test does not support backprop") + + a = tf.Variable(0.1, dtype=tf.float32) + b = tf.Variable(0.2, dtype=tf.float32) + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, r_dtype=np.float32) + + @qnode( + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) + def circuit(a, b): + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + + with tf.GradientTape() as tape: + res = circuit(a, b) + res = tf.stack(res) + + assert circuit.qtape.trainable_params == [0, 1] + + assert isinstance(res, tf.Tensor) + assert res.shape == (2,) + assert res.dtype is tf.float32 + + res = tape.jacobian(res, [a, b]) + assert [r.dtype is tf.float32 for r in res] + + def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): """Test setting finite-difference jacobian options""" if diff_method not in {"finite-diff", "spsa"}: pytest.skip("Test only works with finite diff and spsa.") a = tf.Variable([0.1, 0.2]) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, interface=interface, @@ -215,21 +264,18 @@ def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, approx_order=2, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) - tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + tape.jacobian(res, a) - def test_changing_trainability( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method in ["backprop", "adjoint", "spsa"]: @@ -238,16 +284,21 @@ def test_changing_trainability( a = tf.Variable(0.1, dtype=tf.float64) b = tf.Variable(0.2, dtype=tf.float64) + num_wires = 2 + diff_kwargs = {} - if diff_method == "finite-diff": + if diff_method == "hadamard": + num_wires = 3 + elif diff_method == "finite-diff": diff_kwargs = {"approx_order": 2, "strategy": "center"} + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, **diff_kwargs, ) def circuit(a, b): @@ -256,7 +307,7 @@ def circuit(a, b): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.stack(res) @@ -266,7 +317,7 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - jac = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, [a, b]) expected = [ [-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)], @@ -277,7 +328,7 @@ def circuit(a, b): a = tf.Variable(0.54, dtype=tf.float64) b = tf.constant(0.8, dtype=tf.float64) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.stack(res) @@ -287,22 +338,25 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) expected = [-tf.sin(a), tf.sin(a) * tf.sin(b)] assert np.allclose(jac, expected, atol=tol, rtol=0) - def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): """Test classical processing within the quantum tape""" a = tf.Variable(0.1, dtype=tf.float64) b = tf.constant(0.2, dtype=tf.float64) c = tf.Variable(0.3, dtype=tf.float64) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(x, y, z): qml.RY(x * z, wires=0) @@ -310,30 +364,30 @@ def circuit(x, y, z): qml.RX(z + z**2 + tf.sin(a), wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b, c) if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [0, 2] assert circuit.qtape.get_parameters() == [a * c, c + c**2 + tf.sin(a)] - res = tape.jacobian(res, [a, b, c], experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, [a, b, c]) assert isinstance(res[0], tf.Tensor) assert res[1] is None assert isinstance(res[2], tf.Tensor) - def test_no_trainable_parameters( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): + def test_no_trainable_parameters(self, dev_name, diff_method, grad_on_execution, interface): """Test evaluation if there are no trainable parameters""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(a, b): qml.RY(a, wires=0) @@ -344,7 +398,7 @@ def circuit(a, b): a = 0.1 b = tf.constant(0.2, dtype=tf.float64) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.stack(res) @@ -355,30 +409,31 @@ def circuit(a, b): assert isinstance(res, tf.Tensor) # can't take the gradient with respect to "a" since it's a Python scalar - grad = tape.jacobian(res, b, experimental_use_pfor=not device_vjp) + grad = tape.jacobian(res, b) assert grad is None @pytest.mark.parametrize("U", [tf.constant([[0, 1], [1, 0]]), np.array([[0, 1], [1, 0]])]) - def test_matrix_parameter( - self, dev, diff_method, grad_on_execution, device_vjp, U, tol, interface - ): + def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, U, tol, interface): """Test that the TF interface works correctly with a matrix parameter""" a = tf.Variable(0.1, dtype=tf.float64) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(U, a): qml.QubitUnitary(U, wires=0) qml.RY(a, wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(U, a) if diff_method == "finite-diff": @@ -386,23 +441,18 @@ def circuit(U, a): assert np.allclose(res, -tf.cos(a), atol=tol, rtol=0) - res = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, a) assert np.allclose(res, tf.sin(a), atol=tol, rtol=0) - def test_differentiable_expand( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test that operation and nested tapes expansion is differentiable""" - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA class U3(qml.U3): @@ -414,6 +464,13 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + a = np.array(0.1) p = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) @@ -423,7 +480,7 @@ def circuit(a, p): U3(p[0], p[1], p[2], wires=0) return qml.expval(qml.PauliX(0)) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, p) expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( @@ -431,7 +488,7 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, p, experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, p) expected = np.array( [ tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), @@ -450,9 +507,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and differentiates it.""" - def test_changing_shots(self, interface): + def test_changing_shots(self, mocker, tol, interface): """Test that changing shots works on execution""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -461,21 +518,31 @@ def circuit(weights): qml.RY(weights[0], wires=0) qml.RX(weights[1], wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - circuit(weights) + res = circuit(weights) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 - res = circuit(weights, shots=100) # pylint: disable=unexpected-keyword-arg - assert res.shape == (100, 2) + circuit(weights, shots=100) # pylint: disable=unexpected-keyword-arg + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(weights) + assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) + spy.assert_called_once() def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" # pylint: disable=unexpected-keyword-arg - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -490,6 +557,7 @@ def circuit(weights): res = circuit(weights, shots=[10000, 10000, 10000]) res = tf.transpose(tf.stack(res)) + assert dev.shots is None assert len(res) == 3 jacobian = tape.jacobian(res, weights) @@ -500,7 +568,7 @@ def test_multiple_gradient_integration(self, tol, interface): """Test that temporarily setting the shots works for gradient computations, even if the QNode has been re-evaluated with a different number of shots in the meantime.""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -514,7 +582,7 @@ def circuit(weights): with tf.GradientTape() as tape: res1 = circuit(weights) - assert qml.math.shape(res1) == () + assert qml.math.shape(res1) == tuple() res2 = circuit(weights, shots=[(1, 1000)]) # pylint: disable=unexpected-keyword-arg assert qml.math.shape(res2) == (1000,) @@ -525,7 +593,7 @@ def circuit(weights): def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=100) weights = tf.Variable([0.543, -0.654], dtype=tf.float64) spy = mocker.spy(qml, "execute") @@ -537,45 +605,116 @@ def circuit(weights): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - circuit(weights, shots=100) # pylint:disable=unexpected-keyword-arg # since we are using finite shots, parameter-shift will # be chosen - assert circuit.gradient_fn == qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + circuit(weights) + assert circuit.gradient_fn is qml.gradients.param_shift - # if we use the default shots value of None, backprop can now be used circuit(weights) - assert circuit.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert circuit.gradient_fn is qml.gradients.param_shift + + # if we set the shots to None, backprop can now be used + circuit(weights, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert circuit.gradient_fn == "backprop" + + circuit(weights) + assert circuit.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + + +@pytest.mark.parametrize("interface", ["auto", "tf"]) +class TestAdjoint: + """Specific integration tests for the adjoint method""" + + def test_reuse_state(self, mocker, interface): + """Tests that the TF interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=2) + + @qnode(dev, diff_method="adjoint", interface=interface) + def circ(x): + qml.RX(x[0], wires=0) + qml.RY(x[1], wires=1) + qml.CNOT(wires=(0, 1)) + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + weights = tf.Variable([0.1, 0.2], dtype=tf.float64) + x, y = 1.0 * weights + + with tf.GradientTape() as tape: + res = tf.reduce_sum(circ(weights)) + + grad = tape.gradient(res, weights) + expected_grad = [-tf.sin(x), tf.cos(y)] + + assert np.allclose(grad, expected_grad) + assert circ.device.num_executions == 1 + spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) + + def test_reuse_state_multiple_evals(self, mocker, tol, interface): + """Tests that the TF interface reuses the device state for adjoint differentiation, + even where there are intermediate evaluations.""" + dev = qml.device("default.qubit.legacy", wires=2) + + x_val = 0.543 + y_val = -0.654 + x = tf.Variable(x_val, dtype=tf.float64) + y = tf.Variable(y_val, dtype=tf.float64) + + @qnode(dev, diff_method="adjoint", interface=interface) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + with tf.GradientTape() as tape: + res1 = circuit(x, y) + + assert np.allclose(res1, np.cos(x_val), atol=tol, rtol=0) + + # intermediate evaluation with different values + circuit(tf.math.tan(x), tf.math.cosh(y)) + + # the adjoint method will continue to compute the correct derivative + grad = tape.gradient(res1, x) + assert np.allclose(grad, -np.sin(x_val), atol=tol, rtol=0) + assert dev.num_executions == 2 + spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" def test_probability_differentiation( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + self, dev_name, diff_method, grad_on_execution, tol, interface ): """Tests correct output shape and evaluation for a tape with multiple probs outputs""" - - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") - - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @@ -586,7 +725,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.probs(wires=[0]), qml.probs(wires=[1]) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(x, y) res = tf.stack(res) @@ -598,7 +737,7 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [x, y], experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, [x, y]) expected = np.array( [ [ @@ -613,26 +752,25 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_ragged_differentiation( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_ragged_differentiation(self, dev_name, diff_method, grad_on_execution, tol, interface): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") - - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } - if diff_method == "spsa": + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @@ -643,7 +781,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1]) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(x, y) res = tf.experimental.numpy.hstack(res) @@ -656,7 +794,7 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [x, y], experimental_use_pfor=not device_vjp) + res = tape.jacobian(res, [x, y]) expected = np.array( [ [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], @@ -665,20 +803,24 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_second_derivative( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_second_derivative(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test second derivative calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -708,18 +850,24 @@ def circuit(x): ] assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_hessian(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -748,20 +896,24 @@ def circuit(x): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface - ): + def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test hessian calculation of a vector valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -804,19 +956,25 @@ def circuit(x): np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) def test_hessian_vector_valued_postprocessing( - self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + self, dev_name, diff_method, grad_on_execution, tol, interface ): """Test hessian calculation of a vector valued QNode with post-processing""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=0) @@ -855,18 +1013,24 @@ def circuit(x): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_ragged(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_hessian_ragged(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test hessian calculation of a ragged QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -917,21 +1081,23 @@ def circuit(x): ] np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) - def test_state(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): + def test_state(self, dev_name, diff_method, grad_on_execution, tol, interface): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -941,6 +1107,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) + assert res.dtype is tf.complex128 probs = tf.math.abs(res) ** 2 return probs[0] + probs[2] @@ -955,27 +1122,21 @@ def cost_fn(x, y): assert np.allclose(grad, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - @pytest.mark.parametrize("dtype", ("int32", "int64")) - def test_projector( - self, state, dev, diff_method, grad_on_execution, device_vjp, tol, interface, dtype - ): + def test_projector(self, state, dev_name, diff_method, grad_on_execution, tol, interface): """Test that the variance of a projector is correctly returned""" - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "interface": interface, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + ) if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or all diagonal measurements.") - if diff_method == "hadamard": + pytest.skip("Adjoint does not support projectors") + elif diff_method == "hadamard": pytest.skip("Variance not implemented yet.") elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - P = tf.constant(state, dtype=dtype) + dev = qml.device(dev_name, wires=2) + P = tf.constant(state) x, y = 0.765, -0.654 weights = tf.Variable([x, y], dtype=tf.float64) @@ -1000,69 +1161,95 @@ def circuit(weights): ] assert np.allclose(grad, expected, atol=tol, rtol=0) - def test_postselection_differentiation( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") +@pytest.mark.parametrize( + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": 0.05}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +class TestCV: + """Tests for CV integration""" - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) + def test_first_order_observable(self, diff_method, kwargs, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + r = tf.Variable(0.543, dtype=tf.float64) + phi = tf.Variable(-0.654, dtype=tf.float64) + + @qnode(dev, diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + with tf.GradientTape() as tape: + res = circuit(r, phi) + + expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 + assert np.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + grad = tape.gradient(res, [r, phi]) + expected = [ + 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, + 2 * np.sinh(2 * r) * np.sin(2 * phi), + ] + assert np.allclose(grad, expected, atol=tol, rtol=0) - phi = tf.Variable(1.23) - theta = tf.Variable(4.56) + def test_second_order_observable(self, diff_method, kwargs, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + n = tf.Variable(0.12, dtype=tf.float64) + a = tf.Variable(0.765, dtype=tf.float64) - assert np.allclose(circuit(phi, theta), expected_circuit(theta)) + @qnode(dev, diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) - with tf.GradientTape() as res_tape: - res = circuit(phi, theta) - gradient = res_tape.gradient(res, [phi, theta]) + with tf.GradientTape() as tape: + res = circuit(n, a) - with tf.GradientTape() as expected_tape: - expected = expected_circuit(theta) - exp_theta_grad = expected_tape.gradient(expected, theta) + expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) + assert np.allclose(res, expected, atol=tol, rtol=0) - assert np.allclose(gradient, [0.0, exp_theta_grad]) + # circuit jacobians + grad = tape.gradient(res, [n, a]) + expected = [2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)] + assert np.allclose(grad, expected, atol=tol, rtol=0) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the TF interface""" - def test_gradient_expansion(self, dev, diff_method, grad_on_execution, device_vjp, interface): + def test_gradient_expansion(self, dev_name, diff_method, grad_on_execution, interface): """Test that a *supported* operation with no gradient recipe is expanded for both parameter-shift and finite-differences, but not for execution.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) class PhaseShift(qml.PhaseShift): grad_method = None @@ -1076,7 +1263,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, ) def circuit(x): qml.Hadamard(wires=0) @@ -1088,6 +1274,7 @@ def circuit(x): with tf.GradientTape() as t2: with tf.GradientTape() as t1: loss = circuit(x) + res = t1.gradient(loss, x) assert np.allclose(res, -3 * np.sin(3 * x)) @@ -1099,13 +1286,20 @@ def circuit(x): @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface + self, dev_name, diff_method, grad_on_execution, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + class PhaseShift(qml.PhaseShift): grad_method = None @@ -1118,7 +1312,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, - device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1136,24 +1329,24 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, tol, interface + self, dev_name, diff_method, grad_on_execution, max_diff, tol, interface ): """Test that if there are non-commuting groups and the number of shots is None the first and second order gradients are correctly evaluated""" - kwargs = { - "diff_method": diff_method, - "grad_on_execution": grad_on_execution, - "max_diff": max_diff, - "interface": interface, - "device_vjp": device_vjp, - } + kwargs = dict( + diff_method=diff_method, + grad_on_execution=grad_on_execution, + max_diff=max_diff, + interface=interface, + ) if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint/hadamard method does not yet support Hamiltonians") elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} + dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1200,14 +1393,14 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, max_diff, mocker, interface ): - """Test that the Hamiltonian is correctly measured if there + """Test that the Hamiltonian is expanded if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.1 + tol = 0.3 if diff_method in ("adjoint", "backprop", "hadamard"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": @@ -1220,6 +1413,8 @@ def test_hamiltonian_finite_shots( elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} + dev = qml.device(dev_name, wires=3, shots=50000) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1228,7 +1423,6 @@ def test_hamiltonian_finite_shots( grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1244,19 +1438,17 @@ def circuit(data, weights, coeffs): c = tf.Variable([-0.6543, 0.24, 0.54], dtype=tf.float64) # test output - with tf.GradientTape(persistent=True) as _t2: + with tf.GradientTape(persistent=True) as t2: with tf.GradientTape() as t1: - res = circuit(d, w, c, shots=50000) # pylint:disable=unexpected-keyword-arg + res = circuit(d, w, c) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) + spy.assert_called() # test gradients grad = t1.gradient(res, [d, w, c]) - if diff_method in ["finite-diff", "spsa"]: - pytest.skip(f"{diff_method} not compatible") - expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1266,36 +1458,33 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[2], expected_c, atol=tol) # test second-order derivatives - # TODO: figure out why grad2_c is np.zeros((3,3)) instead of None - # if diff_method == "parameter-shift" and max_diff == 2: - # grad2_c = _t2.jacobian(grad[2], c) - # print(grad2_c, grad[2], c) - # assert grad2_c is None + if diff_method == "parameter-shift" and max_diff == 2: + grad2_c = t2.jacobian(grad[2], c) + assert grad2_c is None - # grad2_w_c = _t2.jacobian(grad[1], c) - # expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ - # 0, - # -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), - # -np.sin(d[1] + w[1]), - # ] - # assert np.allclose(grad2_w_c, expected, atol=tol) + grad2_w_c = t2.jacobian(grad[1], c) + expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ + 0, + -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), + -np.sin(d[1] + w[1]), + ] + assert np.allclose(grad2_w_c, expected, atol=tol) class TestSample: """Tests for the sample integration""" - # pylint:disable=unexpected-keyword-arg - def test_sample_dimension(self): """Test sampling works as expected""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert isinstance(res, tuple) assert len(res) == 2 @@ -1307,14 +1496,15 @@ def circuit(): def test_sampling_expval(self): """Test sampling works as expected if combined with expectation values""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert len(res) == 2 assert isinstance(res, tuple) @@ -1325,67 +1515,76 @@ def circuit(): def test_sample_combination(self): """Test the output of combining expval, var and sample""" + n_sample = 10 - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) + + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit(shots=10) + result = circuit() assert len(result) == 3 - assert result[0].shape == (10,) + assert result[0].shape == (n_sample,) assert result[1].shape == () assert result[2].shape == () assert isinstance(result[0], tf.Tensor) assert isinstance(result[1], tf.Tensor) assert isinstance(result[2], tf.Tensor) - assert result[0].dtype is tf.float64 # pylint:disable=no-member - assert result[1].dtype is tf.float64 # pylint:disable=no-member - assert result[2].dtype is tf.float64 # pylint:disable=no-member + assert result[0].dtype is tf.float64 + assert result[1].dtype is tf.float64 + assert result[2].dtype is tf.float64 def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" + n_sample = 10 + + dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit(shots=10) + result = circuit() assert isinstance(result, tf.Tensor) - assert np.array_equal(result.shape, (10,)) + assert np.array_equal(result.shape, (n_sample,)) def test_multi_wire_sample_regular_shape(self): """Test the return type and shape of sampling multiple wires where a rectangular array is expected""" + n_sample = 10 + + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") + @qnode(dev, diff_method="parameter-shift", interface="tf") def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit(shots=10) + result = circuit() result = tf.stack(result) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tf.Tensor) - assert np.array_equal(result.shape, (3, 10)) + assert np.array_equal(result.shape, (3, n_sample)) assert result.dtype == tf.float64 def test_counts(self): """Test counts works as expected for TF""" + dev = qml.device("default.qubit.legacy", wires=2, shots=100) - # pylint:disable=unsubscriptable-object,no-member - @qnode(DefaultQubit(), interface="tf") + @qnode(dev, interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.counts(qml.PauliZ(0)) - res = circuit(shots=100) + res = circuit() assert isinstance(res, dict) assert list(res.keys()) == [-1, 1] @@ -1415,11 +1614,12 @@ class TestAutograph: def test_autograph_gradients(self, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated""" + dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) + @qnode(dev, diff_method="parameter-shift", interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1428,7 +1628,7 @@ def circuit(x, y): with tf.GradientTape() as tape: p0, p1 = circuit(x, y) - loss = p0[0] + p1[1] # pylint:disable=unsubscriptable-object + loss = p0[0] + p1[1] expected = tf.cos(x / 2) ** 2 + (1 - tf.cos(x) * tf.cos(y)) / 2 assert np.allclose(loss, expected, atol=tol, rtol=0) @@ -1440,11 +1640,12 @@ def circuit(x, y): def test_autograph_jacobian(self, decorator, interface, tol): """Test that a parameter-shift vector-valued QNode can be compiled using @tf.function, and differentiated""" + dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface) + @qnode(dev, diff_method="parameter-shift", max_diff=1, interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1482,14 +1683,10 @@ def circuit(x, y): def test_autograph_adjoint_single_param(self, grad_on_execution, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" + dev = qml.device("default.qubit.legacy", wires=1) @decorator - @qnode( - DefaultQubit(), - diff_method="adjoint", - interface=interface, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution) def circuit(x): qml.RY(x, wires=0) return qml.expval(qml.PauliZ(0)) @@ -1510,14 +1707,10 @@ def circuit(x): def test_autograph_adjoint_multi_params(self, grad_on_execution, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" + dev = qml.device("default.qubit.legacy", wires=1) @decorator - @qnode( - DefaultQubit(), - diff_method="adjoint", - interface=interface, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution) def circuit(x): qml.RY(x[0], wires=0) qml.RX(x[1], wires=0) @@ -1536,44 +1729,16 @@ def circuit(x): expected_g = [-tf.sin(a) * tf.cos(b), -tf.cos(a) * tf.sin(b)] assert np.allclose(g, expected_g, atol=tol, rtol=0) - @pytest.mark.xfail - @pytest.mark.parametrize("grad_on_execution", [True, False]) - def test_autograph_adjoint_multi_out(self, grad_on_execution, decorator, interface, tol): - """Test that a parameter-shift QNode can be compiled - using @tf.function, and differentiated to second order""" - - @decorator - @qnode( - DefaultQubit(), - diff_method="adjoint", - interface=interface, - grad_on_execution=grad_on_execution, - ) - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(0)) - - x = tf.Variable(0.5, dtype=tf.float64) - - with tf.GradientTape() as tape: - res = qml.math.hstack(circuit(x)) - g = tape.jacobian(res, x) - a = x * 1.0 - - expected_res = [tf.cos(a), tf.sin(a)] - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - expected_g = [-tf.sin(a), tf.cos(a)] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - def test_autograph_ragged_differentiation(self, decorator, interface, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + dev = qml.device("default.qubit.legacy", wires=2) + x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface) + @qnode(dev, diff_method="parameter-shift", max_diff=1, interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1605,11 +1770,12 @@ def circuit(x, y): def test_autograph_hessian(self, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" + dev = qml.device("default.qubit.legacy", wires=1) a = tf.Variable(0.543, dtype=tf.float64) b = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=2, interface=interface) + @qnode(dev, diff_method="parameter-shift", max_diff=2, interface=interface) def circuit(x, y): qml.RY(x, wires=0) qml.RX(y, wires=0) @@ -1635,15 +1801,35 @@ def circuit(x, y): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) + def test_autograph_sample(self, decorator, interface): + """Test that a QNode returning raw samples can be compiled + using @tf.function""" + dev = qml.device("default.qubit", wires=2, shots=100) + x = tf.Variable(0.543, dtype=tf.float64) + y = tf.Variable(-0.654, dtype=tf.float64) + + @decorator + @qnode(dev, diff_method="parameter-shift", interface=interface) + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.sample() + + result = circuit(x, y) + result = np.array(result).flatten() + assert np.all([r in [1, 0] for r in result]) + def test_autograph_state(self, decorator, interface, tol): """Test that a parameter-shift QNode returning a state can be compiled using @tf.function""" + dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) # TODO: fix this for diff_method=None @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) + @qnode(dev, diff_method="parameter-shift", interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1660,15 +1846,16 @@ def circuit(x, y): def test_autograph_dimension(self, decorator, interface): """Test sampling works as expected""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) @decorator - @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) - def circuit(**_): + @qnode(dev, diff_method="parameter-shift", interface=interface) + def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=10) # pylint:disable=unexpected-keyword-arg + res = circuit() res = tf.stack(res) assert res.shape == (2, 10) @@ -1676,23 +1863,24 @@ def circuit(**_): @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution,device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" def test_grad_single_measurement_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For one measurement and one param, the gradient is a float.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -1710,16 +1898,19 @@ def circuit(a): assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -1740,16 +1931,19 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) @@ -1768,64 +1962,68 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") + + num_wires = 2 - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable(0.1, dtype=tf.float64) + a = tf.Variable(0.1) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) assert isinstance(jac, tf.Tensor) assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) + a = tf.Variable(0.1) + b = tf.Variable(0.2) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) - jac = tape.jacobian(res, (a, b), experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, (a, b)) assert isinstance(jac, tuple) @@ -1836,94 +2034,101 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """For a multi dimensional measurement (probs), check that a single array is returned.""" + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable([0.1, 0.2], dtype=tf.float64) + a = tf.Variable([0.1, 0.2]) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) assert isinstance(jac, tf.Tensor) assert jac.shape == (4, 2) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The jacobian of multiple measurements with a single params return an array.""" + num_wires = 2 - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable(0.1, dtype=tf.float64) + a = tf.Variable(0.1) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) assert isinstance(jac, tf.Tensor) assert jac.shape == (5,) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) + a = tf.Variable(0.1) + b = tf.Variable(0.2) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a, b) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, (a, b), experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, (a, b)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1935,40 +2140,50 @@ def circuit(a, b): assert jac[1].shape == (5,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("state adjoint diff not supported with lightning") + if diff_method == "adjoint": + pytest.skip("Test does not supports adjoint because of probabilities.") + + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable([0.1, 0.2], dtype=tf.float64) + a = tf.Variable([0.1, 0.2]) - with tf.GradientTape(persistent=device_vjp) as tape: + with tf.GradientTape() as tape: res = circuit(a) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + jac = tape.jacobian(res, a) assert isinstance(jac, tf.Tensor) assert jac.shape == (5, 2) def test_hessian_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -1981,7 +2196,6 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2008,13 +2222,20 @@ def circuit(x, y): assert hess[1].shape == (2,) def test_hessian_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2023,7 +2244,6 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2042,10 +2262,10 @@ def circuit(x): assert isinstance(hess, tf.Tensor) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, interface - ): + def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution, interface): """The hessian of single a measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2061,7 +2281,6 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2088,7 +2307,7 @@ def circuit(x, y): assert hess[1].shape == (2,) def test_hessian_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2097,6 +2316,8 @@ def test_hessian_var_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because of variance.") + dev = qml.device(dev_name, wires=2) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2105,7 +2326,6 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2125,9 +2345,16 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + num_wires = 2 + + if diff_method == "hadamard": + pytest.skip("Test does not support hadamard because multiple measurements.") + + dev = qml.device(dev_name, wires=num_wires) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2140,7 +2367,6 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2168,7 +2394,7 @@ def circuit(x, y): assert hess[1].shape == (6,) def test_hessian_probs_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" @@ -2178,6 +2404,10 @@ def test_hessian_probs_expval_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because multiple measurements.") + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2186,7 +2416,6 @@ def test_hessian_probs_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2207,9 +2436,11 @@ def circuit(x): assert hess.shape == (3, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") if diff_method == "hadamard": @@ -2224,7 +2455,6 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2252,7 +2482,7 @@ def circuit(x, y): assert hess[1].shape == (6,) def test_hessian_probs_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, interface + self, dev_name, diff_method, grad_on_execution, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2260,6 +2490,8 @@ def test_hessian_probs_var_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because of variance.") + dev = qml.device(dev_name, wires=2) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2268,7 +2500,6 @@ def test_hessian_probs_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2289,49 +2520,17 @@ def circuit(x): assert hess.shape == (3, 2, 2) -def test_no_ops(): +@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) +def test_no_ops(dev_name): """Test that the return value of the QNode matches in the interface even if there are no ops""" - @qml.qnode(DefaultQubit(), interface="tf") + dev = qml.device(dev_name, wires=1) + + @qml.qnode(dev, interface="tf") def circuit(): qml.Hadamard(wires=0) return qml.state() res = circuit() assert isinstance(res, tf.Tensor) - - -def test_error_device_vjp_jacobian(): - """Test a ValueError is raised if a jacobian is attempted to be computed with device_vjp=True.""" - - dev = qml.device("default.qubit") - - @qml.qnode(dev, diff_method="adjoint", device_vjp=True) - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(0)) - - x = tf.Variable(0.1) - - with tf.GradientTape() as tape: - y = qml.math.hstack(circuit(x)) - - with pytest.raises(ValueError): - tape.jacobian(y, x) - - -def test_error_device_vjp_state_float32(): - """Test a ValueError is raised is state differentiation is attemped with float32 parameters.""" - - dev = qml.device("default.qubit") - - @qml.qnode(dev, diff_method="adjoint", device_vjp=True) - def circuit(x): - qml.RX(x, wires=0) - return qml.probs(wires=0) - - x = tf.Variable(0.1, dtype=tf.float32) - with pytest.raises(ValueError, match="tensorflow with adjoint differentiation of the state"): - with tf.GradientTape(): - circuit(x) diff --git a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_shot_vector_legacy.py similarity index 82% rename from tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_shot_vector_legacy.py index a3a5606fc78..bea3056b7c5 100644 --- a/tests/interfaces/default_qubit_2_integration/test_tensorflow_qnode_shot_vector_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_tensorflow_qnode_shot_vector_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2023 Xanadu Quantum Technologies Inc. +# Copyright 2022 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. @@ -12,13 +12,12 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the TF interface with shot vectors and with a QNode""" -# pylint: disable=too-many-arguments,unexpected-keyword-arg +# pylint: disable=too-many-arguments import pytest import pennylane as qml from pennylane import numpy as np from pennylane import qnode -from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf @@ -28,9 +27,9 @@ shots_and_num_copies_hess = [((10, (5, 1)), 2)] qubit_device_and_diff_method = [ - [DefaultQubit(), "finite-diff", {"h": 10e-2}], - [DefaultQubit(), "parameter-shift", {}], - [DefaultQubit(), "spsa", {"h": 10e-2, "num_directions": 20}], + ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], + ["default.qubit.legacy", "parameter-shift", {}], + ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}], ] TOLS = { @@ -46,15 +45,16 @@ @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) @pytest.mark.parametrize( - "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnWithShotVectors: """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" def test_jac_single_measurement_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and one param, the gradient is a float.""" + dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -65,7 +65,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -74,9 +74,10 @@ def circuit(a): assert jac.shape == (num_copies,) def test_jac_single_measurement_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -88,7 +89,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -100,9 +101,10 @@ def circuit(a, b): assert j.shape == (num_copies,) def test_jacobian_single_measurement_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -113,7 +115,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -122,10 +124,11 @@ def circuit(a): assert jac.shape == (num_copies, 2) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -136,7 +139,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -145,10 +148,11 @@ def circuit(a): assert jac.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -160,7 +164,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -172,10 +176,11 @@ def circuit(a, b): assert j.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -186,7 +191,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -195,9 +200,10 @@ def circuit(a): assert jac.shape == (num_copies, 4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The gradient of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1) par_1 = tf.Variable(0.2) @@ -210,7 +216,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) with tf.GradientTape() as tape: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, (par_0, par_1)) @@ -222,9 +228,10 @@ def circuit(x, y): assert j.shape == (num_copies, 2) def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -236,7 +243,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2, 0.3]) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, a) @@ -245,9 +252,10 @@ def circuit(a): assert jac.shape == (num_copies, 2, 3) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a single params return an array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -258,7 +266,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -267,9 +275,10 @@ def circuit(a): assert jac.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -281,7 +290,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b, shots=shots) + res = circuit(a, b) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, (a, b)) @@ -293,9 +302,10 @@ def circuit(a, b): assert j.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -306,7 +316,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a, shots=shots) + res = circuit(a) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -318,15 +328,16 @@ def circuit(a): @pytest.mark.slow @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) @pytest.mark.parametrize( - "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnShotVectorHessian: """Class to test the shape of the Hessian with different return types and shot vectors.""" def test_hessian_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of a single measurement with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1, dtype=tf.float64) par_1 = tf.Variable(0.2, dtype=tf.float64) @@ -340,7 +351,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack(res) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -355,9 +366,10 @@ def circuit(x, y): assert h.shape == (2, num_copies) def test_hessian_expval_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of single measurement with a multiple params array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) params = tf.Variable([0.1, 0.2], dtype=tf.float64) @@ -370,7 +382,7 @@ def circuit(x): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(params, shots=shots) + res = circuit(params) res = qml.math.stack(res) jac = tape2.jacobian(res, params, experimental_use_pfor=False) @@ -381,9 +393,10 @@ def circuit(x): assert hess.shape == (num_copies, 2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1, dtype=tf.float64) par_1 = tf.Variable(0.2, dtype=tf.float64) @@ -397,7 +410,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1, shots=shots) + res = circuit(par_0, par_1) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -412,10 +425,12 @@ def circuit(x, y): assert h.shape == (2, num_copies, 3) def test_hessian_expval_probs_multiple_param_array( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode(dev, diff_method=diff_method, interface=interface, max_diff=2, **gradient_kwargs) @@ -427,7 +442,7 @@ def circuit(x): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(params, shots=shots) + res = circuit(params) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape2.jacobian(res, params, experimental_use_pfor=False) @@ -438,22 +453,22 @@ def circuit(x): assert hess.shape == (num_copies, 3, 2, 2) -shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)] +shots_and_num_copies = [((20000, 18000, 16000), 3), ((20000, (18000, 2)), 3)] @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) @pytest.mark.parametrize( - "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnShotVectorIntegration: """Tests for the integration of shots with the TF interface.""" def test_single_expectation_value( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """Tests correct output shape and evaluation for a tape with a single expval output""" - + dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543) y = tf.Variable(-0.654) @@ -465,7 +480,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) with tf.GradientTape() as tape: - res = circuit(x, y, shots=shots) + res = circuit(x, y) res = qml.math.stack(res) all_res = tape.jacobian(res, (x, y)) @@ -482,11 +497,11 @@ def circuit(x, y): assert np.allclose(res, exp, atol=tol, rtol=0) def test_prob_expectation_values( - self, dev, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - + dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543) y = tf.Variable(-0.654) @@ -498,7 +513,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) with tf.GradientTape() as tape: - res = circuit(x, y, shots=shots) + res = circuit(x, y) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) all_res = tape.jacobian(res, (x, y)) diff --git a/tests/interfaces/legacy_devices_integration/test_torch_legacy.py b/tests/interfaces/legacy_devices_integration/test_torch_legacy.py new file mode 100644 index 00000000000..5aac97fb0e3 --- /dev/null +++ b/tests/interfaces/legacy_devices_integration/test_torch_legacy.py @@ -0,0 +1,1308 @@ +# Copyright 2018-2022 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. +"""Unit tests for the Torch interface""" +# pylint: disable=protected-access,too-few-public-methods +import numpy as np +import pytest + +import pennylane as qml +from pennylane import execute +from pennylane.gradients import finite_diff, param_shift +from pennylane.typing import TensorLike + +pytestmark = pytest.mark.torch +torch = pytest.importorskip("torch") +torch_functional = pytest.importorskip("torch.autograd.functional") +torch_cuda = pytest.importorskip("torch.cuda") + + +@pytest.mark.parametrize("interface", ["torch", "auto"]) +class TestTorchExecuteUnitTests: + """Unit tests for torch execution""" + + def test_jacobian_options(self, interface, mocker): + """Test setting jacobian options""" + spy = mocker.spy(qml.gradients, "param_shift") + + a = torch.tensor([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute( + [tape], + dev, + gradient_fn=param_shift, + gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, + interface=interface, + )[0] + + res.backward() + + for args in spy.call_args_list: + assert args[1]["shift"] == [(np.pi / 4,)] * 2 + + def test_incorrect_grad_on_execution(self, interface): + """Test that an error is raised if a gradient transform + is used with grad_on_execution=True""" + a = torch.tensor([0.1, 0.2], requires_grad=True) + + dev = qml.device("default.qubit.legacy", wires=1) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + with pytest.raises( + ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" + ): + execute( + [tape], dev, gradient_fn=param_shift, grad_on_execution=True, interface=interface + ) + + def test_grad_on_execution_reuse_state(self, interface, mocker): + """Test that grad_on_execution uses the `device.execute_and_gradients` pathway + while reusing the quantum state.""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(dev, "execute_and_gradients") + + a = torch.tensor([0.1, 0.2], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + interface=interface, + ) + + # adjoint method only performs a single device execution, but gets both result and gradient + assert dev.num_executions == 1 + spy.assert_called() + + def test_grad_on_execution(self, interface, mocker): + """Test that grad on execution uses the `device.execute_and_gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(dev, "execute_and_gradients") + + a = torch.tensor([0.1, 0.2], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian"}, + interface=interface, + ) + + # two device executions; one for the value, one for the Jacobian + assert dev.num_executions == 2 + spy.assert_called() + + def test_no_grad_on_execution(self, interface, mocker): + """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" + dev = qml.device("default.qubit.legacy", wires=1) + spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") + spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") + + a = torch.tensor([0.1, 0.2], requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute( + [tape], + dev, + gradient_fn="device", + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + interface=interface, + )[0] + + assert dev.num_executions == 1 + spy_execute.assert_called() + spy_gradients.assert_not_called() + + res.backward() + spy_gradients.assert_called() + + +class TestCaching: + """Test for caching behaviour""" + + def test_cache_maxsize(self, mocker): + """Test the cachesize property of the cache""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cachesize): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], dev, gradient_fn=param_shift, cachesize=cachesize, interface="torch" + )[0][0] + + params = torch.tensor([0.1, 0.2], requires_grad=True) + res = cost(params, cachesize=2) + res.backward() + cache = spy.call_args.kwargs["cache"] + + assert cache.maxsize == 2 + assert cache.currsize == 2 + assert len(cache) == 2 + + def test_custom_cache(self, mocker): + """Test the use of a custom cache object""" + dev = qml.device("default.qubit.legacy", wires=1) + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="torch")[0][ + 0 + ] + + custom_cache = {} + params = torch.tensor([0.1, 0.2], requires_grad=True) + res = cost(params, cache=custom_cache) + res.backward() + + cache = spy.call_args.kwargs["cache"] + assert cache is custom_cache + + def test_caching_param_shift(self): + """Test that, with the parameter-shift transform, + Torch always uses the optimum number of evals when computing the Jacobian.""" + dev = qml.device("default.qubit.legacy", wires=1) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="torch")[0][ + 0 + ] + + # Without caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + params = torch.tensor([0.1, 0.2], requires_grad=True) + torch_functional.jacobian(lambda p: cost(p, cache=None), params) + assert dev.num_executions == 5 + + # With caching, 5 evaluations are required to compute + # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) + dev._num_executions = 0 + torch_functional.jacobian(lambda p: cost(p, cache=True), params) + assert dev.num_executions == 5 + + @pytest.mark.parametrize("num_params", [2, 3]) + def test_caching_param_shift_hessian(self, num_params, tol): + """Test that, with the parameter-shift transform, + caching reduces the number of evaluations to their optimum + when computing Hessians.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor(np.arange(1, num_params + 1) / 10, requires_grad=True) + + N = len(params) + + def cost(x, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + + for i in range(2, num_params): + qml.RZ(x[i], wires=[i % 2]) + + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], dev, gradient_fn=param_shift, cache=cache, interface="torch", max_diff=2 + )[0] + + # No caching: number of executions is not ideal + hess1 = torch.autograd.functional.hessian(lambda x: cost(x, cache=None), params) + + if num_params == 2: + # compare to theoretical result + x, y, *_ = params.detach() + expected = torch.tensor( + [ + [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert np.allclose(expected, hess1, atol=tol, rtol=0) + + expected_runs = 1 # forward pass + expected_runs += 2 * N # Jacobian + expected_runs += 4 * N + 1 # Hessian diagonal + expected_runs += 4 * N**2 # Hessian off-diagonal + assert dev.num_executions == expected_runs + + # Use caching: number of executions is ideal + dev._num_executions = 0 + hess2 = torch.autograd.functional.hessian(lambda x: cost(x, cache=True), params) + assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + expected_runs_ideal = 1 # forward pass + expected_runs_ideal += 2 * N # Jacobian + expected_runs_ideal += N + 1 # Hessian diagonal + expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal + assert dev.num_executions == expected_runs_ideal + assert expected_runs_ideal < expected_runs + + def test_caching_adjoint_no_grad_on_execution(self): + """Test that caching reduces the number of adjoint evaluations + when grad_on_execution=False""" + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor([0.1, 0.2, 0.3]) + + def cost(a, cache): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a[0], wires=0) + qml.RX(a[1], wires=0) + qml.RY(a[2], wires=0) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + cache=cache, + grad_on_execution=False, + gradient_kwargs={"method": "adjoint_jacobian"}, + interface="torch", + )[0] + + # Without caching, 2 evaluations are required. + # 1 for the forward pass, and one for the backward pass + torch_functional.jacobian(lambda x: cost(x, cache=None), params) + assert dev.num_executions == 2 + + # With caching, only 2 evaluations are required. One + # for the forward pass, and one for the backward pass. + dev._num_executions = 0 + torch_functional.jacobian(lambda x: cost(x, cache=True), params) + assert dev.num_executions == 2 + + +torch_devices = [None] + +if torch_cuda.is_available(): + torch_devices.append(torch.device("cuda")) + + +execute_kwargs_integration = [ + {"gradient_fn": param_shift, "interface": "torch"}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": False}, + "interface": "torch", + }, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + "interface": "torch", + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + "interface": "torch", + }, + {"gradient_fn": param_shift, "interface": "auto"}, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": False}, + "interface": "auto", + }, + { + "gradient_fn": "device", + "grad_on_execution": True, + "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, + "interface": "auto", + }, + { + "gradient_fn": "device", + "grad_on_execution": False, + "gradient_kwargs": {"method": "adjoint_jacobian"}, + "interface": "auto", + }, +] + + +@pytest.mark.gpu +@pytest.mark.parametrize("torch_device", torch_devices) +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +class TestTorchExecuteIntegration: + """Test the torch interface execute function + integrates well for both forward and backward execution""" + + def test_execution(self, torch_device, execute_kwargs): + """Test that the execute function produces results with the expected shapes""" + dev = qml.device("default.qubit.legacy", wires=1) + a = torch.tensor(0.1, requires_grad=True, device=torch_device) + b = torch.tensor(0.2, requires_grad=False, device=torch_device) + + with qml.queuing.AnnotatedQueue() as q1: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(a, wires=0) + qml.RX(b, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + res = execute([tape1, tape2], dev, **execute_kwargs) + + assert isinstance(res, TensorLike) + assert len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () + + def test_scalar_jacobian(self, torch_device, execute_kwargs, tol): + """Test scalar jacobian calculation by comparing two types of pipelines""" + a = torch.tensor(0.1, requires_grad=True, dtype=torch.float64, device=torch_device) + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + res.backward() + + # compare to backprop gradient + def cost(a): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + dev = qml.device("default.qubit.autograd", wires=2) + return dev.batch_execute([tape])[0] + + expected = qml.grad(cost, argnum=0)(0.1) + assert torch.allclose(a.grad, torch.tensor(expected, device=torch_device), atol=tol, rtol=0) + + def test_jacobian(self, torch_device, execute_kwargs, tol): + """Test jacobian calculation by checking against analytic values""" + a_val = 0.1 + b_val = 0.2 + + a = torch.tensor(a_val, requires_grad=True, device=torch_device) + b = torch.tensor(b_val, requires_grad=True, device=torch_device) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RZ(torch.tensor(0.543, device=torch_device), wires=0) + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [1, 2] + + assert isinstance(res, tuple) + + assert isinstance(res[0], torch.Tensor) + assert res[0].shape == () + + assert isinstance(res[1], torch.Tensor) + assert res[1].shape == () + + expected = torch.tensor( + [np.cos(a_val), -np.cos(a_val) * np.sin(b_val)], device=torch_device + ) + assert torch.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) + assert torch.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + + loss = res[0] + res[1] + + loss.backward() + expected = torch.tensor( + [-np.sin(a_val) + np.sin(a_val) * np.sin(b_val), -np.cos(a_val) * np.cos(b_val)], + dtype=a.dtype, + device=torch_device, + ) + assert torch.allclose(a.grad, expected[0], atol=tol, rtol=0) + assert torch.allclose(b.grad, expected[1], atol=tol, rtol=0) + + def test_tape_no_parameters(self, torch_device, execute_kwargs, tol): + """Test that a tape with no parameters is correctly + ignored during the gradient computation""" + dev = qml.device("default.qubit.legacy", wires=1) + params = torch.tensor([0.1, 0.2], requires_grad=True, device=torch_device) + x, y = params.detach() + + with qml.queuing.AnnotatedQueue() as q1: + qml.Hadamard(0) + qml.expval(qml.PauliX(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RY(0.5, wires=0) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + with qml.queuing.AnnotatedQueue() as q3: + qml.RY(params[0], wires=0) + qml.RX(params[1], wires=0) + qml.expval(qml.PauliZ(0)) + + tape3 = qml.tape.QuantumScript.from_queue(q3) + + res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs)) + expected = torch.tensor(1 + np.cos(0.5), dtype=res.dtype) + torch.cos(x) * torch.cos(y) + expected = expected.to(device=res.device) + + assert torch.allclose(res, expected, atol=tol, rtol=0) + + res.backward() + grad = params.grad.detach() + expected = torch.tensor( + [-torch.cos(y) * torch.sin(x), -torch.cos(x) * torch.sin(y)], + dtype=grad.dtype, + device=grad.device, + ) + assert torch.allclose(grad, expected, atol=tol, rtol=0) + + def test_reusing_quantum_tape(self, torch_device, execute_kwargs, tol): + """Test re-using a quantum tape by passing new parameters""" + a = torch.tensor(0.1, requires_grad=True, device=torch_device) + b = torch.tensor(0.2, requires_grad=True, device=torch_device) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(a, wires=0) + qml.RX(b, wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliY(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + assert tape.trainable_params == [0, 1] + + res = execute([tape], dev, **execute_kwargs)[0] + loss = res[0] + res[1] + loss.backward() + + a_val = 0.54 + b_val = 0.8 + a = torch.tensor(a_val, requires_grad=True, device=torch_device) + b = torch.tensor(b_val, requires_grad=True, device=torch_device) + + tape = tape.bind_new_parameters([2 * a, b], [0, 1]) + res2 = execute([tape], dev, **execute_kwargs)[0] + + expected = torch.tensor( + [np.cos(2 * a_val), -np.cos(2 * a_val) * np.sin(b_val)], + device=torch_device, + dtype=res2[0].dtype, + ) + assert torch.allclose(res2[0].detach(), expected[0], atol=tol, rtol=0) + assert torch.allclose(res2[1].detach(), expected[1], atol=tol, rtol=0) + + loss = res2[0] + res2[1] + loss.backward() + + expected = torch.tensor( + [ + -2 * np.sin(2 * a_val) + 2 * np.sin(2 * a_val) * np.sin(b_val), + -np.cos(2 * a_val) * np.cos(b_val), + ], + dtype=a.dtype, + device=torch_device, + ) + + assert torch.allclose(a.grad, expected[0], atol=tol, rtol=0) + assert torch.allclose(b.grad, expected[1], atol=tol, rtol=0) + + def test_classical_processing(self, torch_device, execute_kwargs, tol): + """Test the classical processing of gate parameters within the quantum tape""" + p_val = [0.1, 0.2] + params = torch.tensor(p_val, requires_grad=True, device=torch_device) + + dev = qml.device("default.qubit.legacy", wires=1) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(params[0] * params[1], wires=0) + qml.RZ(0.2, wires=0) + qml.RX(params[1] + params[1] ** 2 + torch.sin(params[0]), wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + assert tape.trainable_params == [0, 2] + + tape_params = torch.tensor([i.detach() for i in tape.get_parameters()], device=torch_device) + expected = torch.tensor( + [p_val[0] * p_val[1], p_val[1] + p_val[1] ** 2 + np.sin(p_val[0])], + dtype=tape_params.dtype, + device=torch_device, + ) + + assert torch.allclose( + tape_params, + expected, + atol=tol, + rtol=0, + ) + + res.backward() + + assert isinstance(params.grad, torch.Tensor) + assert params.shape == (2,) + + def test_no_trainable_parameters(self, torch_device, execute_kwargs): + """Test evaluation and Jacobian if there are no trainable parameters""" + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RY(0.2, wires=0) + qml.RX(torch.tensor(0.1, device=torch_device), wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.expval(qml.PauliZ(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [] + + assert isinstance(res, tuple) + assert len(res) == 2 + + assert isinstance(res[0], torch.Tensor) + assert res[0].shape == () + + assert isinstance(res[1], torch.Tensor) + assert res[1].shape == () + + with pytest.raises( + RuntimeError, + match="element 0 of tensors does not require grad and does not have a grad_fn", + ): + res[0].backward() + + with pytest.raises( + RuntimeError, + match="element 0 of tensors does not require grad and does not have a grad_fn", + ): + res[1].backward() + + @pytest.mark.parametrize( + "U", [torch.tensor([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]])] + ) + def test_matrix_parameter(self, torch_device, U, execute_kwargs, tol): + """Test that the torch interface works correctly + with a matrix parameter""" + a_val = 0.1 + a = torch.tensor(a_val, requires_grad=True, device=torch_device) + + if isinstance(U, torch.Tensor) and torch_device is not None: + U = U.to(torch_device) + + dev = qml.device("default.qubit.legacy", wires=2) + + with qml.queuing.AnnotatedQueue() as q: + qml.QubitUnitary(U, wires=0) + qml.RY(a, wires=0) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + assert tape.trainable_params == [1] + + expected = torch.tensor(-np.cos(a_val), dtype=res.dtype, device=torch_device) + assert torch.allclose(res.detach(), expected, atol=tol, rtol=0) + + res.backward() + expected = torch.tensor([np.sin(a_val)], dtype=a.grad.dtype, device=torch_device) + assert torch.allclose(a.grad, expected, atol=tol, rtol=0) + + def test_differentiable_expand(self, torch_device, execute_kwargs, tol): + """Test that operation and nested tape expansion + is differentiable""" + + class U3(qml.U3): + def decomposition(self): + theta, phi, lam = self.data + wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] + + dev = qml.device("default.qubit.legacy", wires=1) + a = np.array(0.1) + p_val = [0.1, 0.2, 0.3] + p = torch.tensor(p_val, requires_grad=True, device=torch_device) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(a, wires=0) + U3(p[0], p[1], p[2], wires=0) + qml.expval(qml.PauliX(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + expected = torch.tensor( + np.cos(a) * np.cos(p_val[1]) * np.sin(p_val[0]) + + np.sin(a) + * ( + np.cos(p_val[2]) * np.sin(p_val[1]) + + np.cos(p_val[0]) * np.cos(p_val[1]) * np.sin(p_val[2]) + ), + dtype=res.dtype, + device=torch_device, + ) + assert torch.allclose(res.detach(), expected, atol=tol, rtol=0) + + res.backward() + expected = torch.tensor( + [ + np.cos(p_val[1]) + * (np.cos(a) * np.cos(p_val[0]) - np.sin(a) * np.sin(p_val[0]) * np.sin(p_val[2])), + np.cos(p_val[1]) * np.cos(p_val[2]) * np.sin(a) + - np.sin(p_val[1]) + * (np.cos(a) * np.sin(p_val[0]) + np.cos(p_val[0]) * np.sin(a) * np.sin(p_val[2])), + np.sin(a) + * ( + np.cos(p_val[0]) * np.cos(p_val[1]) * np.cos(p_val[2]) + - np.sin(p_val[1]) * np.sin(p_val[2]) + ), + ], + dtype=p.grad.dtype, + device=torch_device, + ) + assert torch.allclose(p.grad, expected, atol=tol, rtol=0) + + def test_probability_differentiation(self, torch_device, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob outputs""" + + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + x_val = 0.543 + y_val = -0.654 + x = torch.tensor(x_val, requires_grad=True, device=torch_device) + y = torch.tensor(y_val, requires_grad=True, device=torch_device) + + def circuit(x, y): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=[0]) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, **execute_kwargs)[0] + + res = circuit(x, y) + + expected_0 = torch.tensor( + [np.cos(x_val / 2) ** 2, np.sin(x_val / 2) ** 2], + dtype=res[0].dtype, + device=torch_device, + ) + + expected_1 = torch.tensor( + [ + (1 + np.cos(x_val) * np.cos(y_val)) / 2, + (1 - np.cos(x_val) * np.cos(y_val)) / 2, + ], + dtype=res[0].dtype, + device=torch_device, + ) + + assert torch.allclose(res[0], expected_0, atol=tol, rtol=0) + assert torch.allclose(res[1], expected_1, atol=tol, rtol=0) + + jac = torch_functional.jacobian(circuit, (x, y)) + dtype_jac = jac[0][0].dtype + + res_0 = torch.tensor( + [-np.sin(x_val) / 2, np.sin(x_val) / 2], dtype=dtype_jac, device=torch_device + ) + res_1 = torch.tensor([0.0, 0.0], dtype=dtype_jac, device=torch_device) + res_2 = torch.tensor( + [-np.sin(x_val) * np.cos(y_val) / 2, np.cos(y_val) * np.sin(x_val) / 2], + dtype=dtype_jac, + device=torch_device, + ) + res_3 = torch.tensor( + [-np.cos(x_val) * np.sin(y_val) / 2, +np.cos(x_val) * np.sin(y_val) / 2], + dtype=dtype_jac, + device=torch_device, + ) + + assert torch.allclose(jac[0][0], res_0, atol=tol, rtol=0) + assert torch.allclose(jac[0][1], res_1, atol=tol, rtol=0) + assert torch.allclose(jac[1][0], res_2, atol=tol, rtol=0) + assert torch.allclose(jac[1][1], res_3, atol=tol, rtol=0) + + def test_ragged_differentiation(self, torch_device, execute_kwargs, tol): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + if execute_kwargs["gradient_fn"] == "device": + pytest.skip("Adjoint differentiation does not yet support probabilities") + + dev = qml.device("default.qubit.legacy", wires=2) + x_val = 0.543 + y_val = -0.654 + x = torch.tensor(x_val, requires_grad=True, device=torch_device) + y = torch.tensor(y_val, requires_grad=True, device=torch_device) + + def circuit(x, y): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + qml.probs(wires=[1]) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute([tape], dev, **execute_kwargs)[0] + + res = circuit(x, y) + + res_0 = torch.tensor(np.cos(x_val), dtype=res[0].dtype, device=torch_device) + res_1 = torch.tensor( + [(1 + np.cos(x_val) * np.cos(y_val)) / 2, (1 - np.cos(x_val) * np.cos(y_val)) / 2], + dtype=res[0].dtype, + device=torch_device, + ) + + assert isinstance(res, tuple) + assert len(res) == 2 + + assert torch.allclose(res[0], res_0, atol=tol, rtol=0) + assert torch.allclose(res[1], res_1, atol=tol, rtol=0) + + jac = torch_functional.jacobian(circuit, (x, y)) + dtype_jac = jac[0][0].dtype + + res_0 = torch.tensor( + -np.sin(x_val), + dtype=dtype_jac, + device=torch_device, + ) + res_1 = torch.tensor( + 0.0, + dtype=dtype_jac, + device=torch_device, + ) + res_2 = torch.tensor( + [-np.sin(x_val) * np.cos(y_val) / 2, np.cos(y_val) * np.sin(x_val) / 2], + dtype=dtype_jac, + device=torch_device, + ) + res_3 = torch.tensor( + [-np.cos(x_val) * np.sin(y_val) / 2, +np.cos(x_val) * np.sin(y_val) / 2], + dtype=dtype_jac, + device=torch_device, + ) + + assert torch.allclose(jac[0][0], res_0, atol=tol, rtol=0) + assert torch.allclose(jac[0][1], res_1, atol=tol, rtol=0) + assert torch.allclose(jac[1][0], res_2, atol=tol, rtol=0) + assert torch.allclose(jac[1][1], res_3, atol=tol, rtol=0) + + def test_sampling(self, torch_device, execute_kwargs): + """Test sampling works as expected""" + # pylint: disable=unused-argument + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + if execute_kwargs["interface"] == "auto": + pytest.skip("Can't detect interface without a parametrized gate in the tape") + + dev = qml.device("default.qubit.legacy", wires=2, shots=10) + + with qml.queuing.AnnotatedQueue() as q: + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.sample(qml.PauliZ(0)) + qml.sample(qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + assert isinstance(res, tuple) + assert len(res) == 2 + + assert isinstance(res[0], torch.Tensor) + assert res[0].shape == (10,) + + assert isinstance(res[1], torch.Tensor) + assert res[1].shape == (10,) + + def test_sampling_expval(self, torch_device, execute_kwargs): + """Test sampling works as expected if combined with expectation values""" + # pylint: disable=unused-argument + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + if execute_kwargs["interface"] == "auto": + pytest.skip("Can't detect interface without a parametrized gate in the tape") + + dev = qml.device("default.qubit.legacy", wires=2, shots=10) + + with qml.queuing.AnnotatedQueue() as q: + qml.Hadamard(wires=[0]) + qml.CNOT(wires=[0, 1]) + qml.sample(qml.PauliZ(0)) + qml.expval(qml.PauliX(1)) + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + assert len(res) == 2 + assert isinstance(res, tuple) + assert res[0].shape == (10,) + assert res[1].shape == () + assert isinstance(res[0], torch.Tensor) + assert isinstance(res[1], torch.Tensor) + + def test_sampling_gradient_error(self, torch_device, execute_kwargs): + """Test differentiating a tape with sampling results in an error""" + # pylint: disable=unused-argument + if ( + execute_kwargs["gradient_fn"] == "device" + and execute_kwargs["grad_on_execution"] is True + ): + pytest.skip("Adjoint differentiation does not support samples") + + dev = qml.device("default.qubit.legacy", wires=1, shots=10) + + x = torch.tensor(0.65, requires_grad=True) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x, wires=[0]) + qml.sample() + + tape = qml.tape.QuantumScript.from_queue(q) + + res = execute([tape], dev, **execute_kwargs)[0] + + with pytest.raises( + RuntimeError, + match="element 0 of tensors does not require grad and does not have a grad_fn", + ): + res.backward() + + def test_repeated_application_after_expand(self, torch_device, execute_kwargs): + """Test that the Torch interface continues to work after + tape expansions""" + # pylint: disable=unused-argument + n_qubits = 2 + dev = qml.device("default.qubit.legacy", wires=n_qubits) + + weights = torch.ones((3,)) + + with qml.queuing.AnnotatedQueue() as q: + qml.U3(*weights, wires=0) + qml.expval(qml.PauliZ(wires=0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + tape = tape.expand() + execute([tape], dev, **execute_kwargs) + + +@pytest.mark.parametrize("torch_device", torch_devices) +class TestHigherOrderDerivatives: + """Test that the torch execute function can be differentiated""" + + @pytest.mark.parametrize( + "params", + [ + torch.tensor([0.543, -0.654], requires_grad=True), + torch.tensor([0, -0.654], requires_grad=True), + torch.tensor([-2.0, 0], requires_grad=True), + ], + ) + def test_parameter_shift_hessian(self, torch_device, params, tol): + """Tests that the output of the parameter-shift transform + can be differentiated using torch, yielding second derivatives.""" + # pylint: disable=unused-argument + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[0], wires=0) + qml.RY(x[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + result = execute( + [tape1, tape2], dev, gradient_fn=param_shift, interface="torch", max_diff=2 + ) + return result[0] + result[1][0] + + res = cost_fn(params) + x, y = params.detach() + expected = torch.as_tensor(0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))) + assert torch.allclose(res, expected, atol=tol, rtol=0) + + res.backward() + expected = torch.tensor( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert torch.allclose(params.grad.detach(), expected, atol=tol, rtol=0) + + res = torch.autograd.functional.hessian(cost_fn, params) + expected = torch.tensor( + [ + [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], + [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + ] + ) + assert torch.allclose(res, expected, atol=tol, rtol=0) + + def test_hessian_vector_valued(self, torch_device, tol): + """Test hessian calculation of a vector valued tape""" + dev = qml.device("default.qubit.legacy", wires=1) + + def circuit(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RY(x[0], wires=0) + qml.RX(x[1], wires=0) + qml.probs(wires=0) + + tape = qml.tape.QuantumScript.from_queue(q) + + return torch.stack( + execute([tape], dev, gradient_fn=param_shift, interface="torch", max_diff=2) + ) + + x = torch.tensor([1.0, 2.0], requires_grad=True, device=torch_device) + res = circuit(x) + + if torch_device is not None: + a, b = x.detach().cpu().numpy() + else: + a, b = x.detach().numpy() + + expected_res = torch.tensor( + [ + 0.5 + 0.5 * np.cos(a) * np.cos(b), + 0.5 - 0.5 * np.cos(a) * np.cos(b), + ], + dtype=res.dtype, + device=torch_device, + ) + assert torch.allclose(res.detach(), expected_res, atol=tol, rtol=0) + + def jac_fn(x): + return torch_functional.jacobian(circuit, x, create_graph=True) + + g = jac_fn(x) + + hess = torch_functional.jacobian(jac_fn, x) + + expected_g = torch.tensor( + [ + [-0.5 * np.sin(a) * np.cos(b), -0.5 * np.cos(a) * np.sin(b)], + [0.5 * np.sin(a) * np.cos(b), 0.5 * np.cos(a) * np.sin(b)], + ], + dtype=g.dtype, + device=torch_device, + ) + assert torch.allclose(g.detach(), expected_g, atol=tol, rtol=0) + + expected_hess = torch.tensor( + [ + [ + [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.sin(a) * np.sin(b)], + [0.5 * np.sin(a) * np.sin(b), -0.5 * np.cos(a) * np.cos(b)], + ], + [ + [0.5 * np.cos(a) * np.cos(b), -0.5 * np.sin(a) * np.sin(b)], + [-0.5 * np.sin(a) * np.sin(b), 0.5 * np.cos(a) * np.cos(b)], + ], + ], + dtype=hess.dtype, + device=torch_device, + ) + assert torch.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) + + def test_adjoint_hessian(self, torch_device, tol): + """Since the adjoint hessian is not a differentiable transform, + higher-order derivatives are not supported.""" + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor( + [0.543, -0.654], requires_grad=True, dtype=torch.float64, device=torch_device + ) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape = qml.tape.QuantumScript.from_queue(q) + + return execute( + [tape], + dev, + gradient_fn="device", + gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, + interface="torch", + )[0] + + res = torch.autograd.functional.hessian(cost_fn, params) + expected = torch.zeros([2, 2], dtype=torch.float64, device=torch_device) + assert torch.allclose(res, expected, atol=tol, rtol=0) + + def test_max_diff(self, torch_device, tol): + """Test that setting the max_diff parameter blocks higher-order + derivatives""" + dev = qml.device("default.qubit.legacy", wires=2) + params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) + + def cost_fn(x): + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) + qml.CNOT(wires=[0, 1]) + qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(x[0], wires=0) + qml.RY(x[0], wires=1) + qml.CNOT(wires=[0, 1]) + qml.probs(wires=1) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + + result = execute( + [tape1, tape2], dev, gradient_fn=param_shift, max_diff=1, interface="torch" + ) + return result[0] + result[1][0] + + res = cost_fn(params) + x, y = params.detach() + expected = torch.as_tensor(0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))) + assert torch.allclose(res.to(torch_device), expected.to(torch_device), atol=tol, rtol=0) + + res.backward() + expected = torch.tensor( + [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + ) + assert torch.allclose( + params.grad.detach().to(torch_device), expected.to(torch_device), atol=tol, rtol=0 + ) + + res = torch.autograd.functional.hessian(cost_fn, params) + expected = torch.zeros([2, 2], dtype=torch.float64) + assert torch.allclose(res.to(torch_device), expected.to(torch_device), atol=tol, rtol=0) + + +execute_kwargs_hamiltonian = [ + {"gradient_fn": param_shift, "interface": "torch"}, + {"gradient_fn": finite_diff, "interface": "torch"}, +] + + +@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +class TestHamiltonianWorkflows: + """Test that tapes ending with expectations + of Hamiltonians provide correct results and gradients""" + + @pytest.fixture + def cost_fn(self, execute_kwargs): + """Cost function for gradient tests""" + + def _cost_fn(weights, coeffs1, coeffs2, dev=None): + obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] + H1 = qml.Hamiltonian(coeffs1, obs1) + + obs2 = [qml.PauliZ(0)] + H2 = qml.Hamiltonian(coeffs2, obs2) + + with qml.queuing.AnnotatedQueue() as q: + qml.RX(weights[0], wires=0) + qml.RY(weights[1], wires=1) + qml.CNOT(wires=[0, 1]) + qml.expval(H1) + qml.expval(H2) + + tape = qml.tape.QuantumScript.from_queue(q) + + return torch.hstack(execute([tape], dev, **execute_kwargs)[0]) + + return _cost_fn + + @staticmethod + def cost_fn_expected(weights, coeffs1, coeffs2): + """Analytic value of cost_fn above""" + a, b, c = coeffs1.detach().numpy() + d = coeffs2.detach().numpy()[0] + x, y = weights.detach().numpy() + return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + + @staticmethod + def cost_fn_jacobian(weights, coeffs1, coeffs2): + """Analytic jacobian of cost_fn above""" + a, b, c = coeffs1.detach().numpy() + d = coeffs2.detach().numpy()[0] + x, y = weights.detach().numpy() + return np.array( + [ + [ + -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), + b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), + np.cos(x), + np.cos(x) * np.sin(y), + -(np.sin(x) * np.sin(y)), + 0, + ], + [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + ] + ) + + def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): + # pylint: disable=unused-argument + coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False, dtype=torch.float64) + coeffs2 = torch.tensor([0.7], requires_grad=False, dtype=torch.float64) + weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) + assert np.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + + res = torch.hstack( + torch_functional.jacobian(lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)) + ) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + assert np.allclose(res.detach(), expected, atol=tol, rtol=0) + + def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): + # pylint: disable=unused-argument + coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True, dtype=torch.float64) + coeffs2 = torch.tensor([0.7], requires_grad=True, dtype=torch.float64) + weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64) + dev = qml.device("default.qubit.legacy", wires=2) + + res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + assert np.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) + assert np.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + + res = torch.hstack( + torch_functional.jacobian(lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)) + ) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + assert np.allclose(res.detach(), expected, atol=tol, rtol=0) diff --git a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py b/tests/interfaces/legacy_devices_integration/test_torch_qnode_legacy.py similarity index 74% rename from tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py rename to tests/interfaces/legacy_devices_integration/test_torch_qnode_legacy.py index 2f1c85f6275..3965255a4a4 100644 --- a/tests/interfaces/default_qubit_2_integration/test_torch_qnode_default_qubit_2.py +++ b/tests/interfaces/legacy_devices_integration/test_torch_qnode_legacy.py @@ -1,4 +1,4 @@ -# Copyright 2023 Xanadu Quantum Technologies Inc. +# Copyright 2022 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. @@ -12,15 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the Torch interface with a QNode""" -# pylint: disable=too-many-arguments,unexpected-keyword-arg,no-member,comparison-with-callable, no-name-in-module -# pylint: disable=use-implicit-booleaness-not-comparison, unnecessary-lambda-assignment, use-dict-literal +# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods, +# pylint: disable=use-dict-literal, use-implicit-booleaness-not-comparison, unnecessary-lambda-assignment import numpy as np import pytest import pennylane as qml from pennylane import qnode -from pennylane.devices import DefaultQubit -from tests.param_shift_dev import ParamShiftDerivativesDevice pytestmark = pytest.mark.torch @@ -28,25 +26,14 @@ jacobian = torch.autograd.functional.jacobian hessian = torch.autograd.functional.hessian -# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - [DefaultQubit(), "finite-diff", False, False], - [DefaultQubit(), "parameter-shift", False, False], - [DefaultQubit(), "backprop", True, False], - [DefaultQubit(), "adjoint", True, False], - [DefaultQubit(), "adjoint", False, False], - [DefaultQubit(), "adjoint", True, True], - [DefaultQubit(), "adjoint", False, True], - [DefaultQubit(), "spsa", False, False], - [DefaultQubit(), "hadamard", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", False, True], - [qml.device("lightning.qubit", wires=5), "adjoint", True, True], - [qml.device("lightning.qubit", wires=5), "adjoint", False, False], - [qml.device("lightning.qubit", wires=5), "adjoint", True, False], - [ParamShiftDerivativesDevice(), "parameter-shift", False, False], - [ParamShiftDerivativesDevice(), "best", False, False], - [ParamShiftDerivativesDevice(), "parameter-shift", True, False], - [ParamShiftDerivativesDevice(), "parameter-shift", False, True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "spsa", False], + ["default.qubit.legacy", "hadamard", False], ] interface_and_qubit_device_and_diff_method = [ @@ -59,25 +46,25 @@ @pytest.mark.parametrize( - "interface, dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface, dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQNode: """Test that using the QNode with Torch integrates with the PennyLane stack""" - def test_execution_with_interface( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): + def test_execution_with_interface(self, interface, dev_name, diff_method, grad_on_execution): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(a): qml.RY(a, wires=0) @@ -92,7 +79,7 @@ def circuit(a): # with the interface, the tape returns torch tensors assert isinstance(res, torch.Tensor) - assert res.shape == () + assert res.shape == tuple() # the tape is able to deduce trainable parameters assert circuit.qtape.trainable_params == [0] @@ -102,18 +89,20 @@ def circuit(a): grad = a.grad assert isinstance(grad, torch.Tensor) - assert grad.shape == () + assert grad.shape == tuple() - def test_interface_swap(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_interface_swap(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the Torch interface can be applied to a QNode with a pre-existing interface""" + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - diff_method=diff_method, - interface="autograd", - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface="autograd", grad_on_execution=grad_on_execution ) def circuit(a): qml.RY(a, wires=0) @@ -139,19 +128,22 @@ def circuit(a): assert np.allclose(res1, res2.detach().numpy(), atol=tol, rtol=0) assert np.allclose(grad1, grad2, atol=tol, rtol=0) - def test_drawing(self, interface, dev, diff_method, grad_on_execution, device_vjp): + def test_drawing(self, interface, dev_name, diff_method, grad_on_execution): """Test circuit drawing when using the torch interface""" x = torch.tensor(0.1, requires_grad=True) y = torch.tensor([0.2, 0.3], requires_grad=True) z = torch.tensor(0.4, requires_grad=True) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(p1, p2=y, **kwargs): qml.RX(p1, wires=0) @@ -167,17 +159,14 @@ def circuit(p1, p2=y, **kwargs): assert result == expected - def test_jacobian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test jacobian calculation""" kwargs = dict( - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA a_val = 0.1 @@ -186,6 +175,13 @@ def test_jacobian(self, interface, dev, diff_method, grad_on_execution, device_v a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -221,29 +217,19 @@ def circuit(a, b): assert np.allclose(b.grad, expected[1], atol=tol, rtol=0) # TODO: fix this behavior with float: already present before return type. - def test_jacobian_dtype( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + @pytest.mark.xfail + def test_jacobian_dtype(self, interface, dev_name, diff_method, grad_on_execution): """Test calculating the jacobian with a different datatype""" - if not "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("Failing unless lightning.qubit") if diff_method == "backprop": pytest.skip("Test does not support backprop") a = torch.tensor(0.1, dtype=torch.float32, requires_grad=True) b = torch.tensor(0.2, dtype=torch.float32, requires_grad=True) + dev = qml.device(dev_name, wires=2) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -267,20 +253,15 @@ def circuit(a, b): assert a.grad.dtype is torch.float32 assert b.grad.dtype is torch.float32 - def test_jacobian_options( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execution): """Test setting jacobian options""" if diff_method not in {"finite-diff", "spsa"}: pytest.skip("Test only works with finite-diff and spsa") a = torch.tensor([0.1, 0.2], requires_grad=True) + dev = qml.device(dev_name, wires=1) + @qnode( dev, diff_method=diff_method, @@ -288,7 +269,6 @@ def test_jacobian_options( interface=interface, h=1e-8, approx_order=2, - device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -298,9 +278,7 @@ def circuit(a): res = circuit(a) res.backward() - def test_changing_trainability( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -312,12 +290,10 @@ def test_changing_trainability( a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True) + dev = qml.device(dev_name, wires=2) + @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit(a, b): qml.RY(a, wires=0) @@ -366,25 +342,21 @@ def circuit(a, b): expected = -np.sin(a_val) + np.sin(a_val) * np.sin(b_val) assert np.allclose(a.grad, expected, atol=tol, rtol=0) - def test_classical_processing( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + def test_classical_processing(self, interface, dev_name, diff_method, grad_on_execution): """Test classical processing within the quantum tape""" a = torch.tensor(0.1, dtype=torch.float64, requires_grad=True) b = torch.tensor(0.2, dtype=torch.float64, requires_grad=False) c = torch.tensor(0.3, dtype=torch.float64, requires_grad=True) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -404,22 +376,17 @@ def circuit(a, b, c): assert b.grad is None assert isinstance(c.grad, torch.Tensor) - def test_no_trainable_parameters( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + def test_no_trainable_parameters(self, interface, dev_name, diff_method, grad_on_execution): """Test evaluation and Jacobian if there are no trainable parameters""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(a, b): qml.RY(a, wires=0) @@ -463,20 +430,21 @@ def circuit(a, b): np.array([[0, 1], [1, 0]]), ], ) - def test_matrix_parameter( - self, interface, dev, diff_method, grad_on_execution, device_vjp, U, tol - ): + def test_matrix_parameter(self, interface, dev_name, diff_method, grad_on_execution, U, tol): """Test that the Torch interface works correctly with a matrix parameter""" a_val = 0.1 a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -493,23 +461,18 @@ def circuit(U, a): res.backward() assert np.allclose(a.grad, np.sin(a_val), atol=tol, rtol=0) - def test_differentiable_expand( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_differentiable_expand(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that operation and nested tapes expansion is differentiable""" kwargs = dict( - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA - class U3(qml.U3): # pylint:disable=too-few-public-methods + class U3(qml.U3): def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -518,6 +481,12 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) a = np.array(0.1) p_val = [0.1, 0.2, 0.3] p = torch.tensor(p_val, dtype=torch.float64, requires_grad=True) @@ -558,9 +527,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and differentiates it.""" - def test_changing_shots(self): + def test_changing_shots(self, mocker, tol): """Test that changing shots works on execution""" - dev = DefaultQubit() + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) @@ -568,22 +537,33 @@ def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.sample(wires=(0, 1)) + return qml.expval(qml.PauliY(1)) + + spy = mocker.spy(dev, "sample") # execute with device default shots (None) - with pytest.raises(qml.DeviceError): - circuit(a, b) + res = circuit(a, b) + assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0) + spy.assert_not_called() # execute with shots=100 - res = circuit(a, b, shots=100) - assert res.shape == (100, 2) + res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg + spy.assert_called_once() + assert spy.spy_return.shape == (100,) + + # device state has been unaffected + assert dev.shots is None + res = circuit(a, b) + assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0) + spy.assert_called_once() # only same call as above # TODO: add this test after shot vectors addition @pytest.mark.xfail def test_gradient_integration(self): """Test that temporarily setting the shots works for gradient computations""" - dev = DefaultQubit() + # pylint: disable=unexpected-keyword-arg + dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = torch.tensor([0.543, -0.654], requires_grad=True) @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) @@ -605,10 +585,11 @@ def test_multiple_gradient_integration(self, tol): """Test that temporarily setting the shots works for gradient computations, even if the QNode has been re-evaluated with a different number of shots in the meantime.""" + dev = qml.device("default.qubit.legacy", wires=2, shots=None) weights = torch.tensor([0.543, -0.654], requires_grad=True) a, b = weights - @qnode(DefaultQubit(), interface="torch", diff_method=qml.gradients.param_shift) + @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -616,7 +597,7 @@ def circuit(a, b): return qml.expval(qml.PauliY(1)) res1 = circuit(*weights) - assert qml.math.shape(res1) == () + assert qml.math.shape(res1) == tuple() res2 = circuit(*weights, shots=[(1, 1000)]) assert len(res2) == 1000 @@ -628,47 +609,120 @@ def circuit(a, b): def test_update_diff_method(self, mocker): """Test that temporarily setting the shots updates the diff method""" + dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = torch.tensor([0.543, -0.654], requires_grad=True) spy = mocker.spy(qml, "execute") - @qnode(DefaultQubit(), interface="torch") + @qnode(dev, interface="torch") def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn == qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift - # if we use the default shots value of None, backprop can now be used cost_fn(a, b) + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + + # if we set the shots to None, backprop can now be used + cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" + assert cost_fn.gradient_fn == "backprop" + + # original QNode settings are unaffected + + cost_fn(a, b) + assert cost_fn.gradient_fn is qml.gradients.param_shift + assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + + +class TestAdjoint: + """Specific integration tests for the adjoint method""" + + def test_reuse_state(self, mocker): + """Tests that the Torch interface reuses the device state for adjoint differentiation""" + dev = qml.device("default.qubit.legacy", wires=2) + + @qnode(dev, diff_method="adjoint", interface="torch") + def circ(x): + qml.RX(x[0], wires=0) + qml.RY(x[1], wires=1) + qml.CNOT(wires=(0, 1)) + return qml.expval(qml.PauliZ(0)) + + expected_grad = lambda x: torch.tensor([-torch.sin(x[0]), torch.cos(x[1])]) + + spy = mocker.spy(dev, "adjoint_jacobian") + + x1 = torch.tensor([1.0, 1.0], requires_grad=True) + res = circ(x1) + res.backward() + + assert np.allclose(x1.grad[0], expected_grad(x1)[0]) + assert circ.device.num_executions == 1 + spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) + + def test_resuse_state_multiple_evals(self, mocker, tol): + """Tests that the Torch interface reuses the device state for adjoint differentiation, + even where there are intermediate evaluations.""" + dev = qml.device("default.qubit.legacy", wires=2) + + x_val = 0.543 + y_val = -0.654 + x = torch.tensor(x_val, requires_grad=True) + y = torch.tensor(y_val, requires_grad=True) + + @qnode(dev, diff_method="adjoint", interface="torch") + def circuit(x, y): + qml.RX(x, wires=[0]) + qml.RY(y, wires=[1]) + qml.CNOT(wires=[0, 1]) + return qml.expval(qml.PauliZ(0)) + + spy = mocker.spy(dev, "adjoint_jacobian") + + res1 = circuit(x, y) + assert np.allclose(res1.detach(), np.cos(x_val), atol=tol, rtol=0) + + # intermediate evaluation with different values + circuit(torch.tan(x), torch.cosh(y)) + + # the adjoint method will continue to compute the correct derivative + res1.backward() + assert np.allclose(x.grad.detach(), -np.sin(x_val), atol=tol, rtol=0) + assert dev.num_executions == 2 + spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) @pytest.mark.parametrize( - "interface,dev,diff_method,grad_on_execution, device_vjp", - interface_and_qubit_device_and_diff_method, + "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" def test_probability_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measureing probabilities with adjoint.") kwargs = {} - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + x_val = 0.543 y_val = -0.654 x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64) @@ -679,7 +733,6 @@ def test_probability_differentiation( diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface, - device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -715,24 +768,25 @@ def circuit(x, y): assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0) assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0) - def test_ragged_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_ragged_differentiation(self, interface, dev_name, diff_method, grad_on_execution, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("lightning does not support measureing probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, - grad_on_execution=grad_on_execution, - interface=interface, - device_vjp=device_vjp, + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) - if diff_method == "spsa": + if diff_method == "adjoint": + pytest.skip("The adjoint method does not currently support returning probabilities") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + x_val = 0.543 y_val = -0.654 x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64) @@ -770,22 +824,17 @@ def circuit(x, y): assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0) assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0) - def test_chained_qnodes( - self, - interface, - dev, - diff_method, - grad_on_execution, - device_vjp, - ): + def test_chained_qnodes(self, interface, dev_name, diff_method, grad_on_execution): """Test that the gradient of chained QNodes works without error""" + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution ) def circuit1(weights): qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1]) @@ -817,11 +866,18 @@ def cost(weights): loss = cost(weights) loss.backward() - def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a scalar valued QNode""" - if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": + if diff_method in {"adjoint", "spsa"}: pytest.skip("Adjoint and SPSA do not support second derivative.") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -832,7 +888,6 @@ def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vj grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, **options, ) def circuit(x): @@ -868,13 +923,18 @@ def circuit(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a vector valued QNode""" - if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": + if diff_method in {"adjoint", "spsa"}: pytest.skip("Adjoint and SPSA do not support second derivative.") + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -885,7 +945,6 @@ def test_hessian_vector_valued( grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, **options, ) def circuit(x): @@ -932,11 +991,18 @@ def circuit(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test hessian calculation of a ragged QNode""" - if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": + if diff_method in {"adjoint", "spsa"}: pytest.skip("Adjoint and SPSA do not support second derivative.") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires) + options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -947,7 +1013,6 @@ def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, de grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, **options, ) def circuit(x): @@ -1006,23 +1071,29 @@ def circuit_stack(x): assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + self, interface, dev_name, diff_method, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" - if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": + if diff_method in {"adjoint", "spsa"}: pytest.skip("Adjoint and SPSA do not support second derivative.") options = {} if diff_method == "finite-diff": options = {"h": 1e-6} + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) + @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, - device_vjp=device_vjp, **options, ) def circuit(x): @@ -1067,23 +1138,23 @@ def cost_fn(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): + def test_state(self, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the state can be returned and differentiated""" + if diff_method == "adjoint": + pytest.skip("Adjoint does not support states") + + num_wires = 2 - if dev.name == "param_shift.qubit": - pytest.skip("parameter shift does not support measuring the state.") - if "lightning" in getattr(dev, "name", "").lower(): - pytest.xfail("Lightning devices do not support state with adjoint diff.") + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires) x = torch.tensor(0.543, requires_grad=True) y = torch.tensor(-0.654, requires_grad=True) @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1093,7 +1164,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert torch.is_complex(res) + assert res.dtype is torch.complex128 probs = torch.abs(res) ** 2 return probs[0] + probs[2] @@ -1110,25 +1181,20 @@ def cost_fn(x, y): assert torch.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector( - self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol - ): + def test_projector(self, state, interface, dev_name, diff_method, grad_on_execution, tol): """Test that the variance of a projector is correctly returned""" kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface ) if diff_method == "adjoint": - pytest.skip("adjoint supports either all expvals or all diagonal measurements") - if diff_method == "spsa": + pytest.skip("Adjoint does not support projectors") + elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard does not support variances.") + dev = qml.device(dev_name, wires=2) P = torch.tensor(state, requires_grad=False) x, y = 0.765, -0.654 @@ -1156,64 +1222,99 @@ def circuit(x, y): ) assert np.allclose(weights.grad.detach(), expected, atol=tol, rtol=0) - def test_postselection_differentiation( - self, interface, dev, diff_method, grad_on_execution, device_vjp - ): - """Test that when postselecting with default.qubit, differentiation works correctly.""" - if diff_method in ["adjoint", "spsa", "hadamard"]: - pytest.skip("Diff method does not support postselection.") +@pytest.mark.parametrize( + "diff_method,kwargs", + [ + ["finite-diff", {}], + ["spsa", {"num_directions": 100, "h": 0.05}], + ("parameter-shift", {}), + ("parameter-shift", {"force_order2": True}), + ], +) +class TestCV: + """Tests for CV integration""" - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) - def circuit(phi, theta): - qml.RX(phi, wires=0) - qml.CNOT([0, 1]) - qml.measure(wires=0, postselect=1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) - - @qml.qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, + def test_first_order_observable(self, diff_method, kwargs, tol): + """Test variance of a first order CV observable""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA + + r = torch.tensor(0.543, dtype=torch.float64, requires_grad=True) + phi = torch.tensor(-0.654, dtype=torch.float64, requires_grad=True) + + @qnode(dev, interface="torch", diff_method=diff_method, **kwargs) + def circuit(r, phi): + qml.Squeezing(r, 0, wires=0) + qml.Rotation(phi, wires=0) + return qml.var(qml.QuadX(0)) + + res = circuit(r, phi) + expected = torch.exp(2 * r) * torch.sin(phi) ** 2 + torch.exp(-2 * r) * torch.cos(phi) ** 2 + assert torch.allclose(res, expected, atol=tol, rtol=0) + + # circuit jacobians + res.backward() + res = torch.tensor([r.grad, phi.grad]) + expected = torch.tensor( + [ + [ + 2 * torch.exp(2 * r) * torch.sin(phi) ** 2 + - 2 * torch.exp(-2 * r) * torch.cos(phi) ** 2, + 2 * torch.sinh(2 * r) * torch.sin(2 * phi), + ] + ] ) - def expected_circuit(theta): - qml.PauliX(1) - qml.RX(theta, wires=1) - return qml.expval(qml.PauliZ(1)) + assert torch.allclose(res, expected, atol=tol, rtol=0) - phi = torch.tensor(1.23, requires_grad=True) - theta = torch.tensor(4.56, requires_grad=True) + def test_second_order_observable(self, diff_method, kwargs, tol): + """Test variance of a second order CV expectation value""" + dev = qml.device("default.gaussian", wires=1) + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA - assert qml.math.allclose(circuit(phi, theta), expected_circuit(theta)) + n = torch.tensor(0.12, dtype=torch.float64, requires_grad=True) + a = torch.tensor(0.765, dtype=torch.float64, requires_grad=True) - gradient = torch.autograd.grad(circuit(phi, theta), [phi, theta]) - exp_theta_grad = torch.autograd.grad(expected_circuit(theta), theta)[0] - assert qml.math.allclose(gradient, [0.0, exp_theta_grad]) + @qnode(dev, interface="torch", diff_method=diff_method, **kwargs) + def circuit(n, a): + qml.ThermalState(n, wires=0) + qml.Displacement(a, 0, wires=0) + return qml.var(qml.NumberOperator(0)) + res = circuit(n, a) + expected = n**2 + n + torch.abs(a) ** 2 * (1 + 2 * n) + assert torch.allclose(res, expected, atol=tol, rtol=0) -@pytest.mark.parametrize( - "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method -) + # circuit jacobians + res.backward() + res = torch.tensor([n.grad, a.grad]) + expected = torch.tensor([[2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]]) + assert torch.allclose(res, expected, atol=tol, rtol=0) + + +@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the Torch interface""" - def test_gradient_expansion(self, dev, diff_method, grad_on_execution, device_vjp): + def test_gradient_expansion(self, dev_name, diff_method, grad_on_execution): """Test that a *supported* operation with no gradient recipe is expanded for both parameter-shift and finite-differences, but not for execution.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + class PhaseShift(qml.PhaseShift): grad_method = None def decomposition(self): @@ -1224,7 +1325,6 @@ def decomposition(self): diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, - device_vjp=device_vjp, interface="torch", ) def circuit(x): @@ -1235,32 +1335,33 @@ def circuit(x): x = torch.tensor(0.5, requires_grad=True, dtype=torch.float64) loss = circuit(x) - loss.backward() res = x.grad assert torch.allclose(res, -3 * torch.sin(3 * x)) - if diff_method == "parameter-shift" and dev.name != "param_shift.qubit": + if diff_method == "parameter-shift": # test second order derivatives res = torch.autograd.functional.hessian(circuit, x) assert torch.allclose(res, -9 * torch.cos(3 * x)) @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, - dev, - diff_method, - grad_on_execution, - device_vjp, - max_diff, + self, dev_name, diff_method, grad_on_execution, max_diff ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires) + + class PhaseShift(qml.PhaseShift): grad_method = None def decomposition(self): @@ -1272,7 +1373,6 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", - device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1288,7 +1388,7 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev, diff_method, grad_on_execution, max_diff, device_vjp, tol + self, dev_name, diff_method, grad_on_execution, max_diff, tol ): """Test that if there are non-commuting groups and the number of shots is None @@ -1298,17 +1398,17 @@ def test_hamiltonian_expansion_analytic( grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", - device_vjp=device_vjp, ) if diff_method == "adjoint": pytest.skip("The adjoint method does not yet support Hamiltonians") elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - kwargs["num_directions"] = 20 + spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) + kwargs = {**kwargs, **spsa_kwargs} tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("The hadamard method does not yet support Hamiltonians") + dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1348,11 +1448,7 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if ( - diff_method in ("parameter-shift", "backprop") - and max_diff == 2 - and dev.name != "param_shift.qubit" - ): + if diff_method in ("parameter-shift", "backprop") and max_diff == 2: hessians = torch.autograd.functional.hessian(circuit, (d, w, c)) grad2_c = hessians[2][2] @@ -1372,14 +1468,14 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_finite_shots( - self, dev, diff_method, device_vjp, grad_on_execution, max_diff + def test_hamiltonian_expansion_finite_shots( + self, dev_name, diff_method, grad_on_execution, max_diff, mocker ): - """Test that the Hamiltonian is correctly measured if there + """Test that the Hamiltonian is expanded if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.1 + tol = 0.3 if diff_method in ("adjoint", "backprop"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": @@ -1391,10 +1487,11 @@ def test_hamiltonian_finite_shots( tol = TOL_FOR_SPSA elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} - tol = 0.15 elif diff_method == "hadamard": pytest.skip("The hadamard method does not yet support Hamiltonians") + dev = qml.device(dev_name, wires=3, shots=50000) + spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1403,7 +1500,6 @@ def test_hamiltonian_finite_shots( grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", - device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1418,17 +1514,15 @@ def circuit(data, weights, coeffs): c = torch.tensor([-0.6543, 0.24, 0.54], requires_grad=True, dtype=torch.float64) # test output - res = circuit(d, w, c, shots=50000) + res = circuit(d, w, c) expected = c[2] * torch.cos(d[1] + w[1]) - c[1] * torch.sin(d[0] + w[0]) * torch.sin( d[1] + w[1] ) assert torch.allclose(res, expected, atol=tol) + spy.assert_called() # test gradients - if diff_method in ["finite-diff", "spsa"]: - pytest.skip(f"{diff_method} not compatible") - res.backward() grad = (w.grad, c.grad) @@ -1447,10 +1541,8 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2 and dev.name != "param_shift.qubit": - hessians = torch.autograd.functional.hessian( - lambda _d, _w, _c: circuit(_d, _w, _c, shots=50000), (d, w, c) - ) + if diff_method == "parameter-shift" and max_diff == 2: + hessians = torch.autograd.functional.hessian(circuit, (d, w, c)) grad2_c = hessians[2][2] assert torch.allclose(grad2_c, torch.zeros([3, 3], dtype=torch.float64), atol=tol) @@ -1474,14 +1566,15 @@ class TestSample: def test_sample_dimension(self): """Test sampling works as expected""" + dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert isinstance(res, tuple) assert len(res) == 2 @@ -1494,33 +1587,37 @@ def circuit(): def test_sampling_expval(self): """Test sampling works as expected if combined with expectation values""" + shots = 10 + dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert len(res) == 2 assert isinstance(res, tuple) assert isinstance(res[0], torch.Tensor) - assert res[0].shape == (10,) + assert res[0].shape == (shots,) assert isinstance(res[1], torch.Tensor) assert res[1].shape == () def test_counts_expval(self): """Test counts works as expected if combined with expectation values""" + shots = 10 + dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - @qnode(qml.device("default.qubit"), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.counts(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit(shots=10) + res = circuit() assert len(res) == 2 assert isinstance(res, tuple) @@ -1531,19 +1628,22 @@ def circuit(): def test_sample_combination(self): """Test the output of combining expval, var and sample""" + n_sample = 10 + + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit(shots=10) + result = circuit() assert isinstance(result, tuple) assert len(result) == 3 - assert np.array_equal(result[0].shape, (10,)) + assert np.array_equal(result[0].shape, (n_sample,)) assert result[1].shape == () assert isinstance(result[1], torch.Tensor) assert result[2].shape == () @@ -1552,73 +1652,73 @@ def circuit(): def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" + n_sample = 10 - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) + + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit(shots=10) + result = circuit() assert isinstance(result, torch.Tensor) - assert np.array_equal(result.shape, (10,)) + assert np.array_equal(result.shape, (n_sample,)) def test_multi_wire_sample_regular_shape(self): """Test the return type and shape of sampling multiple wires where a rectangular array is expected""" + n_sample = 10 + + dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") + @qnode(dev, diff_method="parameter-shift", interface="torch") def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit(shots=10) + result = circuit() # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) - assert tuple(result[0].shape) == (10,) - assert tuple(result[1].shape) == (10,) - assert tuple(result[2].shape) == (10,) + assert tuple(result[0].shape) == (n_sample,) + assert tuple(result[1].shape) == (n_sample,) + assert tuple(result[2].shape) == (n_sample,) assert result[0].dtype == torch.float64 assert result[1].dtype == torch.float64 assert result[2].dtype == torch.float64 qubit_device_and_diff_method_and_grad_on_execution = [ - [DefaultQubit(), "backprop", True, False], - [DefaultQubit(), "finite-diff", False, False], - [DefaultQubit(), "parameter-shift", False, False], - [DefaultQubit(), "adjoint", True, False], - [DefaultQubit(), "adjoint", False, False], - [DefaultQubit(), "adjoint", True, True], - [DefaultQubit(), "adjoint", False, True], - [DefaultQubit(), "hadamard", False, False], + ["default.qubit.legacy", "backprop", True], + ["default.qubit.legacy", "finite-diff", False], + ["default.qubit.legacy", "parameter-shift", False], + ["default.qubit.legacy", "adjoint", True], + ["default.qubit.legacy", "adjoint", False], + ["default.qubit.legacy", "hadamard", False], ] @pytest.mark.parametrize( - "dev,diff_method,grad_on_execution, device_vjp", - qubit_device_and_diff_method_and_grad_on_execution, + "dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method_and_grad_on_execution ) @pytest.mark.parametrize("shots", [None, 10000]) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" - # pylint:disable=too-many-public-methods - - def test_grad_single_measurement_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots - ): + def test_grad_single_measurement_param(self, dev_name, diff_method, grad_on_execution, shots): """For one measurement and one param, the gradient is a float.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1626,7 +1726,7 @@ def circuit(a): a = torch.tensor(0.1, requires_grad=True) - res = circuit(a, shots=shots) + res = circuit(a) assert isinstance(res, torch.Tensor) assert res.shape == () @@ -1638,19 +1738,20 @@ def circuit(a): assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, shots, device_vjp + self, dev_name, diff_method, grad_on_execution, shots ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1659,7 +1760,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - res = circuit(a, b, shots=shots) + res = circuit(a, b) # gradient res.backward() @@ -1670,19 +1771,20 @@ def circuit(a, b): assert grad_b.shape == () def test_grad_single_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """For one measurement and multiple param as a single array params, the gradient is an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 1 + + if diff_method == "hadamard": + num_wires = 2 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1690,7 +1792,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - res = circuit(a, shots=shots) + res = circuit(a) # gradient res.backward() @@ -1700,7 +1802,7 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" @@ -1710,13 +1812,14 @@ def test_jacobian_single_measurement_param_probs( if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1730,7 +1833,7 @@ def circuit(a): assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1739,13 +1842,14 @@ def test_jacobian_single_measurement_probs_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1754,7 +1858,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - jac = jacobian(lambda _a, _b: circuit(_a, _b, shots=shots), (a, b)) + jac = jacobian(circuit, (a, b)) assert isinstance(jac, tuple) @@ -1765,7 +1869,7 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1774,31 +1878,39 @@ def test_jacobian_single_measurement_probs_multiple_param_single_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, torch.Tensor) assert jac.shape == (4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + par_0 = torch.tensor(0.1, requires_grad=True) par_1 = torch.tensor(0.2, requires_grad=True) @@ -1807,7 +1919,6 @@ def test_jacobian_expval_expval_multiple_params( interface="torch", diff_method=diff_method, max_diff=1, - device_vjp=device_vjp, grad_on_execution=grad_on_execution, ) def circuit(x, y): @@ -1816,7 +1927,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jacobian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) + jac = jacobian(circuit, (par_0, par_1)) assert isinstance(jac, tuple) @@ -1835,19 +1946,20 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_expval_expval_multiple_params_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1855,7 +1967,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1867,7 +1979,7 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_var_var_multiple_params( - self, dev, diff_method, device_vjp, grad_on_execution, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -1877,6 +1989,8 @@ def test_jacobian_var_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + dev = qml.device(dev_name, wires=2, shots=shots) + par_0 = torch.tensor(0.1, requires_grad=True) par_1 = torch.tensor(0.2, requires_grad=True) @@ -1884,7 +1998,6 @@ def test_jacobian_var_var_multiple_params( dev, interface="torch", diff_method=diff_method, - device_vjp=device_vjp, max_diff=1, grad_on_execution=grad_on_execution, ) @@ -1894,7 +2007,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jacobian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) + jac = jacobian(circuit, (par_0, par_1)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1914,7 +2027,7 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_var_var_multiple_params_array( - self, dev, diff_method, device_vjp, grad_on_execution, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -1924,13 +2037,9 @@ def test_jacobian_var_var_multiple_params_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + dev = qml.device(dev_name, wires=2, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1938,7 +2047,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1950,22 +2059,22 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_multiple_measurement_single_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a single params return an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1973,7 +2082,7 @@ def circuit(a): a = torch.tensor(0.1, requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1985,7 +2094,7 @@ def circuit(a): assert jac[1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -1993,13 +2102,14 @@ def test_jacobian_multiple_measurement_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - device_vjp=device_vjp, - grad_on_execution=grad_on_execution, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -2008,7 +2118,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - jac = jacobian(lambda _a, _b: circuit(_a, _b, shots=shots), (a, b)) + jac = jacobian(circuit, (a, b)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2028,7 +2138,7 @@ def circuit(a, b): assert jac[1][1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2036,13 +2146,14 @@ def test_jacobian_multiple_measurement_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - @qnode( - dev, - interface="torch", - diff_method=diff_method, - grad_on_execution=grad_on_execution, - device_vjp=device_vjp, - ) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 3 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2050,7 +2161,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(lambda _a: circuit(_a, shots=shots), a) + jac = jacobian(circuit, a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2061,18 +2172,23 @@ def circuit(a): assert isinstance(jac[1], torch.Tensor) assert jac[1].shape == (4, 2) - def test_hessian_expval_multiple_params( - self, dev, diff_method, grad_on_execution, shots, device_vjp - ): + def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_execution, shots): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") - par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) - par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) + par_0 = torch.tensor(0.1, requires_grad=True) + par_1 = torch.tensor(0.2, requires_grad=True) @qnode( dev, @@ -2080,7 +2196,6 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2088,7 +2203,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) + hess = hessian(circuit, (par_0, par_1)) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2108,7 +2223,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2116,6 +2231,13 @@ def test_hessian_expval_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + params = torch.tensor([0.1, 0.2], requires_grad=True) @qnode( @@ -2124,7 +2246,6 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2132,14 +2253,12 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(lambda _a: circuit(_a, shots=shots), params) + hess = hessian(circuit, params) assert isinstance(hess, torch.Tensor) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots - ): + def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution, shots): """The hessian of a single measurement with multiple params returns a tuple of arrays.""" if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2148,8 +2267,15 @@ def test_hessian_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) - par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + par_0 = torch.tensor(0.1, requires_grad=True) + par_1 = torch.tensor(0.2, requires_grad=True) @qnode( dev, @@ -2157,7 +2283,6 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2165,7 +2290,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(lambda _a, _b: circuit(_a, _b, shots=shots), (par_0, par_1)) + hess = hessian(circuit, (par_0, par_1)) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2185,7 +2310,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2195,7 +2320,9 @@ def test_hessian_var_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - params = torch.tensor([0.1, 0.2], requires_grad=True, dtype=torch.float64) + dev = qml.device(dev_name, wires=2, shots=shots) + + params = torch.tensor([0.1, 0.2], requires_grad=True) @qnode( dev, @@ -2203,7 +2330,6 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2211,13 +2337,13 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(lambda _a: circuit(_a, shots=shots), params) + hess = hessian(circuit, params) assert isinstance(hess, torch.Tensor) assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2228,8 +2354,15 @@ def test_hessian_probs_expval_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) - par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) + num_wires = 2 + + if diff_method == "hadamard": + num_wires = 4 + + dev = qml.device(dev_name, wires=num_wires, shots=shots) + + par_0 = torch.tensor(0.1, requires_grad=True) + par_1 = torch.tensor(0.2, requires_grad=True) @qnode( dev, @@ -2237,7 +2370,6 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2246,7 +2378,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x, y): - return torch.hstack(circuit(x, y, shots=shots)) + return torch.hstack(circuit(x, y)) jac_fn = lambda x, y: jacobian(circuit_stack, (x, y), create_graph=True) @@ -2270,9 +2402,10 @@ def circuit_stack(x, y): assert tuple(hess[1][1].shape) == (3,) def test_hessian_expval_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with a multiple param array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2281,7 +2414,7 @@ def test_hessian_expval_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - par = torch.tensor([0.1, 0.2], requires_grad=True, dtype=torch.float64) + par = torch.tensor([0.1, 0.2], requires_grad=True) @qnode( dev, @@ -2289,7 +2422,6 @@ def test_hessian_expval_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2298,7 +2430,7 @@ def circuit(x): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x): - return torch.hstack(circuit(x, shots=shots)) + return torch.hstack(circuit(x)) jac_fn = lambda x: jacobian(circuit_stack, x, create_graph=True) @@ -2308,9 +2440,10 @@ def circuit_stack(x): assert tuple(hess.shape) == (3, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" + dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2328,7 +2461,6 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2337,7 +2469,7 @@ def circuit(x, y): return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x, y): - return torch.hstack(circuit(x, y, shots=shots)) + return torch.hstack(circuit(x, y)) jac_fn = lambda x, y: jacobian(circuit_stack, (x, y), create_graph=True) @@ -2361,9 +2493,10 @@ def circuit_stack(x, y): assert tuple(hess[1][1].shape) == (3,) def test_hessian_var_probs_multiple_param_array( - self, dev, diff_method, grad_on_execution, device_vjp, shots + self, dev_name, diff_method, grad_on_execution, shots ): """The hessian of multiple measurements with a multiple param array return a single array.""" + dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2380,7 +2513,6 @@ def test_hessian_var_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, - device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2389,7 +2521,7 @@ def circuit(x): return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x): - return torch.hstack(circuit(x, shots=shots)) + return torch.hstack(circuit(x)) jac_fn = lambda x: jacobian(circuit_stack, x, create_graph=True) @@ -2399,11 +2531,14 @@ def circuit_stack(x): assert tuple(hess.shape) == (3, 2, 2) -def test_no_ops(): +@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) +def test_no_ops(dev_name): """Test that the return value of the QNode matches in the interface even if there are no ops""" - @qml.qnode(DefaultQubit(), interface="torch") + dev = qml.device(dev_name, wires=1) + + @qml.qnode(dev, interface="torch") def circuit(): qml.Hadamard(wires=0) return qml.state() diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py index 311c2f32b21..2a6ee306508 100644 --- a/tests/interfaces/test_autograd.py +++ b/tests/interfaces/test_autograd.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 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. @@ -11,325 +11,31 @@ # 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. -"""Unit tests for the autograd interface""" -# pylint: disable=protected-access,too-few-public-methods -import sys - +"""Autograd specific tests for execute and default qubit 2.""" import autograd import pytest +from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml +from pennylane import execute from pennylane import numpy as np -from pennylane.devices import DefaultQubitLegacy -from pennylane.gradients import finite_diff, param_shift -from pennylane.operation import AnyWires, Observable +from pennylane.devices import DefaultQubit +from pennylane.gradients import param_shift +from pennylane.measurements import Shots pytestmark = pytest.mark.autograd -class TestAutogradExecuteUnitTests: - """Unit tests for autograd execution""" - - def test_import_error(self, mocker): - """Test that an exception is caught on import error""" - - mock = mocker.patch.object(autograd.extend, "defvjp") - mock.side_effect = ImportError() - - try: - del sys.modules["pennylane.workflow.interfaces.autograd"] - except KeyError: - pass - - dev = qml.device("default.qubit.legacy", wires=2, shots=None) - - with qml.queuing.AnnotatedQueue() as q: - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - with pytest.raises( - qml.QuantumFunctionError, - match="autograd not found. Please install the latest version " - "of autograd to enable the 'autograd' interface", - ): - qml.execute([tape], dev, gradient_fn=param_shift, interface="autograd") - - def test_jacobian_options(self, mocker): - """Test setting jacobian options""" - spy = mocker.spy(qml.gradients, "param_shift") - - a = np.array([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], - device, - gradient_fn=param_shift, - gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, - )[0] - - qml.jacobian(cost)(a, device=dev) - - for args in spy.call_args_list: - assert args[1]["shifts"] == [(np.pi / 4,)] * 2 - - def test_incorrect_grad_on_execution(self): - """Test that an error is raised if a gradient transform - is used with grad_on_execution=True""" - a = np.array([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, gradient_fn=param_shift, grad_on_execution=True)[0] - - with pytest.raises( - ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" - ): - qml.jacobian(cost)(a, device=dev) - - def test_unknown_interface(self): - """Test that an error is raised if the interface is unknown""" - a = np.array([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, gradient_fn=param_shift, interface="None")[0] - - with pytest.raises(ValueError, match="interface must be in"): - cost(a, device=dev) - - def test_grad_on_execution(self, mocker): - """Test that grad on execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(dev, "execute_and_gradients") - - def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - )[0] - - a = np.array([0.1, 0.2], requires_grad=True) - cost(a) - - # adjoint method only performs a single device execution, but gets both result and gradient - assert dev.num_executions == 1 - spy.assert_called() - - def test_no_gradients_on_execution(self, mocker): - """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") - - def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], - dev, - gradient_fn="device", - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - )[0] - - a = np.array([0.1, 0.2], requires_grad=True) - cost(a) - - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() - - qml.jacobian(cost)(a) - spy_gradients.assert_called() - - -class TestBatchTransformExecution: - """Tests to ensure batch transforms can be correctly executed - via qml.execute and batch_transform""" - - def test_no_batch_transform(self, mocker): - """Test that batch transforms can be disabled and enabled""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100000) - - H = qml.PauliZ(0) @ qml.PauliZ(1) - qml.PauliX(0) - x = 0.6 - y = 0.2 - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=0) - qml.RY(y, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(H) - - tape = qml.tape.QuantumScript.from_queue(q) - spy = mocker.spy(dev, "batch_transform") - - if not qml.operation.active_new_opmath(): - with pytest.raises(AssertionError, match="Hamiltonian must be used with shots=None"): - with pytest.warns( - qml.PennyLaneDeprecationWarning, - match="The device_batch_transform argument is deprecated", - ): - _ = qml.execute([tape], dev, None, device_batch_transform=False) - else: - with pytest.warns( - qml.PennyLaneDeprecationWarning, - match="The device_batch_transform argument is deprecated", - ): - res = qml.execute([tape], dev, None, device_batch_transform=False) - assert np.allclose(res[0], np.cos(y), atol=0.1) - - spy.assert_not_called() - - with pytest.warns( - qml.PennyLaneDeprecationWarning, - match="The device_batch_transform argument is deprecated", - ): - res = qml.execute([tape], dev, None, device_batch_transform=True) - spy.assert_called() - - assert qml.math.shape(res[0]) == () - assert np.allclose(res[0], np.cos(y), rtol=0.05) - - def test_batch_transform_dynamic_shots(self): - """Tests that the batch transform considers the number of shots for the execution, not those - statically on the device.""" - dev = qml.device("default.qubit.legacy", wires=1) - H = 2.0 * qml.PauliZ(0) - qscript = qml.tape.QuantumScript(measurements=[qml.expval(H)]) - with pytest.warns( - qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" - ): - res = qml.execute([qscript], dev, interface=None, override_shots=10) - assert res == (2.0,) - - +# pylint: disable=too-few-public-methods class TestCaching: - """Test for caching behaviour""" - - def test_cache_maxsize(self, mocker): - """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cachesize): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, gradient_fn=param_shift, cachesize=cachesize)[0] - - params = np.array([0.1, 0.2]) - qml.jacobian(cost)(params, cachesize=2) - cache = spy.call_args.kwargs["cache"] - - assert cache.maxsize == 2 - assert cache.currsize == 2 - assert len(cache) == 2 - - def test_custom_cache(self, mocker): - """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache)[0] - - custom_cache = {} - params = np.array([0.1, 0.2]) - qml.jacobian(cost)(params, cache=custom_cache) - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache - - def test_caching_param_shift(self, tol): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum.""" - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache)[0] - - # Without caching, 5 jacobians should still be performed - params = np.array([0.1, 0.2]) - qml.jacobian(cost)(params, cache=None) - assert dev.num_executions == 5 - - # With caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - jac_fn = qml.jacobian(cost) - grad1 = jac_fn(params, cache=True) - assert dev.num_executions == 5 - - # Check that calling the cost function again - # continues to evaluate the device (that is, the cache - # is emptied between calls) - grad2 = jac_fn(params, cache=True) - assert dev.num_executions == 10 - assert np.allclose(grad1, grad2, atol=tol, rtol=0) - - # Check that calling the cost function again - # with different parameters produces a different Jacobian - grad2 = jac_fn(2 * params, cache=True) - assert dev.num_executions == 15 - assert not np.allclose(grad1, grad2, atol=tol, rtol=0) + """Tests for caching behaviour""" @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params, tol): + def test_caching_param_shift_hessian(self, num_params): """Test that, when using parameter-shift transform, caching reduces the number of evaluations to their optimum when computing Hessians.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = DefaultQubit() params = np.arange(1, num_params + 1) / 10 N = len(params) @@ -343,13 +49,16 @@ def cost(x, cache): qml.RZ(x[i], wires=[i % 2]) qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, gradient_fn=param_shift, cache=cache, max_diff=2)[0] + return qml.execute( + [tape], dev, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 + )[0] # No caching: number of executions is not ideal - hess1 = qml.jacobian(qml.grad(cost))(params, cache=False) + with qml.Tracker(dev) as tracker: + hess1 = qml.jacobian(qml.grad(cost))(params, cache=False) if num_params == 2: # compare to theoretical result @@ -360,7 +69,7 @@ def cost(x, cache): [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], ] ) - assert np.allclose(expected, hess1, atol=tol, rtol=0) + assert np.allclose(expected, hess1) expected_runs = 1 # forward pass @@ -375,62 +84,26 @@ def cost(x, cache): # Each tape used to compute the Jacobian is then shifted again expected_runs += runs_for_jacobian * num_shifted_evals - assert dev.num_executions == expected_runs + assert tracker.totals["executions"] == expected_runs # Use caching: number of executions is ideal - dev._num_executions = 0 - hess2 = qml.jacobian(qml.grad(cost))(params, cache=True) - assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + with qml.Tracker(dev) as tracker2: + hess2 = qml.jacobian(qml.grad(cost))(params, cache=True) + assert np.allclose(hess1, hess2) expected_runs_ideal = 1 # forward pass expected_runs_ideal += 2 * N # Jacobian expected_runs_ideal += N + 1 # Hessian diagonal expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert dev.num_executions == expected_runs_ideal + assert tracker2.totals["executions"] == expected_runs_ideal assert expected_runs_ideal < expected_runs - def test_caching_adjoint_no_grad_on_execution(self): - """Test that caching reduces the number of adjoint evaluations - when the grads is not on execution.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = np.array([0.1, 0.2, 0.3]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack( - qml.execute( - [tape], - dev, - gradient_fn="device", - cache=cache, - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - )[0] - ) - - # no caching, but jac for each batch still stored. - qml.jacobian(cost)(params, cache=None) - assert dev.num_executions == 2 - - # With caching, only 2 evaluations are required. One - # for the forward pass, and one for the backward pass. - dev._num_executions = 0 - jac_fn = qml.jacobian(cost) - jac_fn(params, cache=True) - assert dev.num_executions == 2 - def test_single_backward_pass_batch(self): """Tests that the backward pass is one single batch, not a bunch of batches, when parameter shift is requested for multiple tapes.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit") def f(x): tape1 = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.probs(wires=0)]) @@ -444,241 +117,293 @@ def f(x): out = qml.jacobian(f)(x) assert dev.tracker.totals["batches"] == 2 - assert dev.tracker.history["batch_len"] == [2, 4] + assert dev.tracker.history["simulations"] == [1, 1, 1, 1, 1, 1] expected = [-2 * np.cos(x / 2) * np.sin(x / 2), 2 * np.sin(x / 2) * np.cos(x / 2)] assert qml.math.allclose(out, expected) - def test_single_backward_pass_split_hamiltonian(self): - """Tests that the backward pass is one single batch, not a bunch of batches, when parameter - shift derivatives are requested for a tape that the device split into batches.""" - - dev = qml.device("default.qubit.legacy", wires=2, shots=50000) - H = qml.Hamiltonian([1, 1], [qml.PauliY(0), qml.PauliZ(0)], grouping_type="qwc") - - def f(x): - tape = qml.tape.QuantumScript([qml.RX(x, wires=0)], [qml.expval(H)]) - return qml.execute([tape], dev, gradient_fn=qml.gradients.param_shift)[0] +# add tests for lightning 2 when possible +# set rng for device when possible +test_matrix = [ + ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), + ({"gradient_fn": param_shift}, Shots((100000, 100000)), DefaultQubit(seed=42)), + ({"gradient_fn": param_shift}, Shots(None), DefaultQubit()), + ({"gradient_fn": "backprop"}, Shots(None), DefaultQubit()), + ( + {"gradient_fn": "adjoint", "grad_on_execution": True, "device_vjp": False}, + Shots(None), + DefaultQubit(), + ), + ( + { + "gradient_fn": "adjoint", + "grad_on_execution": False, + "device_vjp": False, + }, + Shots(None), + DefaultQubit(), + ), + ({"gradient_fn": "adjoint", "device_vjp": True}, Shots(None), DefaultQubit()), + ( + {"gradient_fn": "device", "device_vjp": False}, + Shots((100000, 100000)), + ParamShiftDerivativesDevice(seed=904747894), + ), + ( + {"gradient_fn": "device", "device_vjp": True}, + Shots((100000, 100000)), + ParamShiftDerivativesDevice(seed=10490244), + ), +] - x = qml.numpy.array(0.1) - with dev.tracker: - out = qml.grad(f)(x) - assert dev.tracker.totals["batches"] == 2 - assert dev.tracker.history["batch_len"] == [2, 4] - - assert qml.math.allclose(out, -np.cos(x) - np.sin(x), atol=0.05) - - -execute_kwargs_integration = [ - {"gradient_fn": param_shift}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - }, -] +def atol_for_shots(shots): + """Return higher tolerance if finite shots.""" + return 1e-2 if shots else 1e-6 -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) class TestAutogradExecuteIntegration: """Test the autograd interface execute function integrates well for both forward and backward execution""" - def test_execution(self, execute_kwargs): + def test_execution(self, execute_kwargs, shots, device): """Test execution""" - dev = qml.device("default.qubit.legacy", wires=1) def cost(a, b): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] + tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] + tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - tape2 = qml.tape.QuantumScript.from_queue(q2) - return qml.execute([tape1, tape2], dev, **execute_kwargs) + return execute([tape1, tape2], device, **execute_kwargs) a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=False) - res = cost(a, b) + with device.tracker: + res = cost(a, b) + + if execute_kwargs.get("grad_on_execution", False): + assert device.tracker.totals["execute_and_derivative_batches"] == 1 + else: + assert device.tracker.totals["batches"] == 1 + assert device.tracker.totals["executions"] == 2 # different wires so different hashes assert len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () - def test_scalar_jacobian(self, execute_kwargs, tol): + assert qml.math.allclose(res[0], np.cos(a) * np.cos(b), atol=atol_for_shots(shots)) + assert qml.math.allclose(res[1], np.cos(a) * np.cos(b), atol=atol_for_shots(shots)) + + def test_scalar_jacobian(self, execute_kwargs, shots, device): """Test scalar jacobian calculation""" a = np.array(0.1, requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], dev, **execute_kwargs)[0] + tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - res = qml.jacobian(cost)(a) - assert res.shape == () + if shots.has_partitioned_shots: + res = qml.jacobian(lambda x: qml.math.hstack(cost(x)))(a) + else: + res = qml.jacobian(cost)(a) + assert res.shape == () # pylint: disable=no-member # compare to standard tape jacobian - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) + tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) tape.trainable_params = [0] tapes, fn = param_shift(tape) - expected = fn(dev.batch_execute(tapes)) + expected = fn(device.execute(tapes)) assert expected.shape == () - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res, -np.sin(a), atol=atol_for_shots(shots)) - def test_jacobian(self, execute_kwargs, tol): + def test_jacobian(self, execute_kwargs, shots, device): """Test jacobian calculation""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - def cost(a, b, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) - - dev = qml.device("default.qubit.legacy", wires=2) + def cost(a, b): + ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - res = cost(a, b, device=dev) + res = cost(a, b) expected = [np.cos(a), -np.cos(a) * np.sin(b)] - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = qml.jacobian(cost)(a, b, device=dev) + res = qml.jacobian(cost)(a, b) assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (2,) - assert res[1].shape == (2,) + if shots.has_partitioned_shots: + assert res[0].shape == (4,) + assert res[1].shape == (4,) + + expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) + for _r, _e in zip(res, expected): + assert np.allclose(_r[:2], _e, atol=atol_for_shots(shots)) + assert np.allclose(_r[2:], _e, atol=atol_for_shots(shots)) + else: + assert res[0].shape == (2,) + assert res[1].shape == (2,) - expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) - assert all(np.allclose(_r, _e, atol=tol, rtol=0) for _r, _e in zip(res, expected)) + expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) + for _r, _e in zip(res, expected): + assert np.allclose(_r, _e, atol=atol_for_shots(shots)) @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") - def test_tape_no_parameters(self, execute_kwargs, tol): + def test_tape_no_parameters(self, execute_kwargs, shots, device): """Test that a tape with no parameters is correctly ignored during the gradient computation""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots + ) + + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5, requires_grad=False), wires=0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - dev = qml.device("default.qubit.legacy", wires=2) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - def cost(params): - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.expval(qml.PauliX(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(np.array(0.5, requires_grad=False), wires=0) - qml.expval(qml.PauliZ(0)) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - with qml.queuing.AnnotatedQueue() as q3: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape3 = qml.tape.QuantumScript.from_queue(q3) - with qml.queuing.AnnotatedQueue() as q4: - qml.RY(np.array(0.5, requires_grad=False), wires=0) - qml.probs(wires=[0, 1]) - - tape4 = qml.tape.QuantumScript.from_queue(q4) - return sum( - autograd.numpy.hstack( - qml.execute([tape1, tape2, tape3, tape4], dev, **execute_kwargs) - ) + tape4 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5, requires_grad=False), 0)], + [qml.probs(wires=[0, 1])], + shots=shots, ) + res = qml.execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) + if shots.has_partitioned_shots: + res = tuple(i for r in res for i in r) + return sum(autograd.numpy.hstack(res)) params = np.array([0.1, 0.2], requires_grad=True) x, y = params res = cost(params) expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + expected = shots.num_copies * expected + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) grad = qml.grad(cost)(params) - expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)] - assert np.allclose(grad, expected, atol=tol, rtol=0) + expected = np.array([-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)]) + if shots.has_partitioned_shots: + expected = shots.num_copies * expected + assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) @pytest.mark.filterwarnings("ignore:Attempted to compute the gradient") - def test_tapes_with_different_return_size(self, execute_kwargs): + def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): """Test that tapes wit different can be executed and differentiated.""" - dev = qml.device("default.qubit.legacy", wires=2) + + if execute_kwargs["gradient_fn"] == "backprop": + pytest.xfail("backprop is not compatible with something about this situation.") def cost(params): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(np.array(0.5, requires_grad=False), wires=0) - qml.expval(qml.PauliZ(0)) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - with qml.queuing.AnnotatedQueue() as q3: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape3 = qml.tape.QuantumScript.from_queue(q3) - return autograd.numpy.hstack(qml.execute([tape1, tape2, tape3], dev, **execute_kwargs)) + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + shots=shots, + ) + + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5, requires_grad=False), 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + res = execute([tape1, tape2, tape3], device, **execute_kwargs) + if shots.has_partitioned_shots: + res = tuple(i for r in res for i in r) + return autograd.numpy.hstack(res) params = np.array([0.1, 0.2], requires_grad=True) + x, y = params res = cost(params) assert isinstance(res, np.ndarray) - assert res.shape == (4,) - + if not shots: + assert res.shape == (4,) + + if shots.has_partitioned_shots: + for i in (0, 1): + assert np.allclose(res[2 * i], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[2 * i + 1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[4 + i], np.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[6 + i], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) + else: + assert np.allclose(res[0], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[2], np.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[3], np.cos(x) * np.cos(y), atol=atol_for_shots(shots)) + + if shots.has_partitioned_shots: + pytest.xfail("autograd jacobians do not work with ragged results and shot vectors.") + # TODO: autograd jacobians with ragged results and shot vectors jac = qml.jacobian(cost)(params) assert isinstance(jac, np.ndarray) - assert jac.shape == (4, 2) + if not shots.has_partitioned_shots: + assert jac.shape == (4, 2) # pylint: disable=no-member - def test_reusing_quantum_tape(self, execute_kwargs, tol): + d1 = -np.sin(x) * np.cos(y) + d2 = -np.cos(x) * np.sin(y) + + if shots.has_partitioned_shots: + assert np.allclose(jac[1], 0, atol=atol_for_shots(shots)) + assert np.allclose(jac[3:4], 0, atol=atol_for_shots(shots)) + + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[2, 0], d1, atol=atol_for_shots(shots)) + + assert np.allclose(jac[6, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[7, 0], d1, atol=atol_for_shots(shots)) + + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[6, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[7, 1], d2, atol=atol_for_shots(shots)) + + else: + assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) + + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) + + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) + + def test_reusing_quantum_tape(self, execute_kwargs, shots, device): """Test re-using a quantum tape by passing new parameters""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) + tape = qml.tape.QuantumScript( + [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], + ) assert tape.trainable_params == [0, 1] def cost(a, b): new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return autograd.numpy.hstack(qml.execute([new_tape], dev, **execute_kwargs)[0]) + return autograd.numpy.hstack(execute([new_tape], device, **execute_kwargs)[0]) jac_fn = qml.jacobian(cost) jac = jac_fn(a, b) @@ -690,7 +415,7 @@ def cost(a, b): # values of the parameters for subsequent calls res2 = cost(2 * a, b) expected = [np.cos(2 * a), -np.cos(2 * a) * np.sin(b)] - assert np.allclose(res2, expected, atol=tol, rtol=0) + assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) jac_fn = qml.jacobian(lambda a, b: cost(2 * a, b)) jac = jac_fn(a, b) @@ -699,92 +424,89 @@ def cost(a, b): [0, -np.cos(2 * a) * np.cos(b)], ) assert isinstance(jac, tuple) and len(jac) == 2 - assert all(np.allclose(_j, _e, atol=tol, rtol=0) for _j, _e in zip(jac, expected)) + for _j, _e in zip(jac, expected): + assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - def test_classical_processing(self, execute_kwargs): + def test_classical_processing(self, execute_kwargs, device, shots): """Test classical processing within the quantum tape""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=False) c = np.array(0.3, requires_grad=True) - def cost(a, b, c, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a * c, wires=0) - qml.RZ(b, wires=0) - qml.RX(c + c**2 + np.sin(a), wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, b, c): + ops = [ + qml.RY(a * c, wires=0), + qml.RZ(b, wires=0), + qml.RX(c + c**2 + np.sin(a), wires=0), + ] - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, **execute_kwargs)[0] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) + if shots.has_partitioned_shots: + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0]) + return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) - res = qml.jacobian(cost)(a, b, c, device=dev) + res = qml.jacobian(cost)(a, b, c) # Only two arguments are trainable assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () - def test_no_trainable_parameters(self, execute_kwargs): + # I tried getting analytic results for this circuit but I kept being wrong and am giving up + + def test_no_trainable_parameters(self, execute_kwargs, shots, device): """Test evaluation and Jacobian if there are no trainable parameters""" a = np.array(0.1, requires_grad=False) b = np.array(0.2, requires_grad=False) - def cost(a, b, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + def cost(a, b): + ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - dev = qml.device("default.qubit.legacy", wires=2) - res = cost(a, b, device=dev) - assert res.shape == (2,) + res = cost(a, b) + assert res.shape == (2 * shots.num_copies,) if shots else (2,) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - res = qml.jacobian(cost)(a, b, device=dev) + res = qml.jacobian(cost)(a, b) assert len(res) == 0 def loss(a, b): - return np.sum(cost(a, b, device=dev)) + return np.sum(cost(a, b)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): res = qml.grad(loss)(a, b) assert np.allclose(res, 0) - def test_matrix_parameter(self, execute_kwargs, tol): + def test_matrix_parameter(self, execute_kwargs, device, shots): """Test that the autograd interface works correctly with a matrix parameter""" U = np.array([[0, 1], [1, 0]], requires_grad=False) a = np.array(0.1, requires_grad=True) - def cost(a, U, device): - with qml.queuing.AnnotatedQueue() as q: - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, **execute_kwargs)[0] + def cost(a, U): + ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) + return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) - res = cost(a, U, device=dev) - assert isinstance(res, np.ndarray) - assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + res = cost(a, U) + assert np.allclose(res, -np.cos(a), atol=atol_for_shots(shots)) jac_fn = qml.jacobian(cost) - jac = jac_fn(a, U, device=dev) + jac = jac_fn(a, U) assert isinstance(jac, np.ndarray) - assert np.allclose(jac, np.sin(a), atol=tol, rtol=0) + assert np.allclose(jac, np.sin(a), atol=atol_for_shots(shots), rtol=0) - def test_differentiable_expand(self, execute_kwargs, tol): + def test_differentiable_expand(self, execute_kwargs, device, shots): """Test that operation and nested tapes expansion is differentiable""" class U3(qml.U3): + """Dummy operator.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -793,28 +515,37 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - def cost_fn(a, p, device): - with qml.queuing.AnnotatedQueue() as q_tape: - qml.RX(a, wires=0) - U3(*p, wires=0) - qml.expval(qml.PauliX(0)) - - tape = qml.tape.QuantumScript.from_queue(q_tape) - tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name)) - return qml.execute([tape], device, **execute_kwargs)[0] + def cost_fn(a, p): + tape = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] + ) + gradient_fn = execute_kwargs["gradient_fn"] + + if gradient_fn is None: + _gradient_method = None + elif isinstance(gradient_fn, str): + _gradient_method = gradient_fn + else: + _gradient_method = "gradient-transform" + config = qml.devices.ExecutionConfig( + interface="autograd", + gradient_method=_gradient_method, + grad_on_execution=execute_kwargs.get("grad_on_execution", None), + ) + program, _ = device.preprocess(execution_config=config) + return execute([tape], device, **execute_kwargs, transform_program=program)[0] a = np.array(0.1, requires_grad=False) p = np.array([0.1, 0.2, 0.3], requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=1) - res = cost_fn(a, p, device=dev) + res = cost_fn(a, p) expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) jac_fn = qml.jacobian(cost_fn) - res = jac_fn(a, p, device=dev) + res = jac_fn(a, p) expected = np.array( [ np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), @@ -825,31 +556,22 @@ def cost_fn(a, p, device): * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_probability_differentiation(self, execute_kwargs, tol): + def test_probability_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") + def cost(x, y): + ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.probs(wires=0), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - def cost(x, y, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=[0]) - qml.probs(wires=[1]) - - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) - - dev = qml.device("default.qubit.legacy", wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) - res = cost(x, y, device=dev) + res = cost(x, y) expected = np.array( [ [ @@ -860,10 +582,10 @@ def cost(x, y, device): ], ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) jac_fn = qml.jacobian(cost) - res = jac_fn(x, y, device=dev) + res = jac_fn(x, y) assert isinstance(res, tuple) and len(res) == 2 assert res[0].shape == (4,) assert res[1].shape == (4,) @@ -886,38 +608,30 @@ def cost(x, y, device): ), ) - assert np.allclose(res[0], expected[0], atol=tol, rtol=0) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - def test_ragged_differentiation(self, execute_kwargs, tol): + def test_ragged_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") - def cost(x, y, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.probs(wires=[1]) - - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], device, **execute_kwargs)[0]) + def cost(x, y): + ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) - dev = qml.device("default.qubit.legacy", wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) - res = cost(x, y, device=dev) + res = cost(x, y) expected = np.array( [np.cos(x), (1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) jac_fn = qml.jacobian(cost) - res = jac_fn(x, y, device=dev) + res = jac_fn(x, y) assert isinstance(res, tuple) and len(res) == 2 assert res[0].shape == (3,) assert res[1].shape == (3,) @@ -926,34 +640,8 @@ def cost(x, y, device): np.array([-np.sin(x), -np.sin(x) * np.cos(y) / 2, np.sin(x) * np.cos(y) / 2]), np.array([0, -np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]), ) - assert np.allclose(res[0], expected[0], atol=tol, rtol=0) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - def test_sampling(self, execute_kwargs): - """Test sampling works as expected""" - if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True - ): - pytest.skip("Adjoint differentiation does not support samples") - - def cost(device): - with qml.queuing.AnnotatedQueue() as q: - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.sample(qml.PauliZ(0)) - qml.sample(qml.PauliX(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute([tape], device, **execute_kwargs)[0] - - shots = 10 - dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - res = cost(device=dev) - assert isinstance(res, tuple) - assert len(res) == 2 - assert res[0].shape == (shots,) - assert res[1].shape == (shots,) + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) class TestHigherOrderDerivatives: @@ -970,24 +658,15 @@ class TestHigherOrderDerivatives: def test_parameter_shift_hessian(self, params, tol): """Tests that the output of the parameter-shift transform can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = DefaultQubit() def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[0], wires=0) - qml.RY(x[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = qml.execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) return result[0] + result[1][0] res = cost_fn(params) @@ -1010,83 +689,20 @@ def cost_fn(x): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_adjoint_hessian_one_param(self, tol): - """Since the adjoint hessian is not a differentiable transform, - higher-order derivatives are not supported.""" - dev = qml.device("default.qubit.autograd", wires=2) - params = np.array([0.543, -0.654], requires_grad=True) - - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return qml.execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - )[0] - - with pytest.warns(UserWarning, match="Output seems independent"): - res = qml.jacobian(qml.grad(cost_fn))(params) - - assert np.allclose(res, np.zeros([2, 2]), atol=tol, rtol=0) - - def test_adjoint_hessian_multiple_params(self, tol): - """Since the adjoint hessian is not a differentiable transform, - higher-order derivatives are not supported.""" - dev = qml.device("default.qubit.autograd", wires=2) - params = np.array([0.543, -0.654], requires_grad=True) - - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack( - qml.execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - )[0] - ) - - with pytest.warns(UserWarning, match="Output seems independent"): - res = qml.jacobian(qml.jacobian(cost_fn))(params) - - assert np.allclose(res, np.zeros([2, 2, 2]), atol=tol, rtol=0) - def test_max_diff(self, tol): """Test that setting the max_diff parameter blocks higher-order derivatives""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = DefaultQubit() params = np.array([0.543, -0.654], requires_grad=True) def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[0], wires=0) - qml.RY(x[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = qml.execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) return result[0] + result[1][0] res = cost_fn(params) @@ -1107,165 +723,26 @@ def cost_fn(x): assert np.allclose(res, expected, atol=tol, rtol=0) -class TestOverridingShots: - """Test overriding shots on execution""" - - def test_changing_shots(self, mocker, tol): - """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) - a, b = np.array([0.543, -0.654], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - spy = mocker.spy(dev, "sample") - - # execute with device default shots (None) - res = qml.execute([tape], dev, gradient_fn=param_shift) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() - - # execute with shots=100 - with pytest.warns( - qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" - ): - res = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100) - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = qml.execute([tape], dev, gradient_fn=param_shift) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # same single call from above, no additional calls - - def test_overriding_shots_with_same_value(self, mocker): - """Overriding shots with the same value as the device will have no effect""" - dev = qml.device("default.qubit.legacy", wires=2, shots=123) - a, b = np.array([0.543, -0.654], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - spy = mocker.Mock(wraps=qml.Device.shots.fset) - # pylint:disable=assignment-from-no-return,too-many-function-args - mock_property = qml.Device.shots.setter(spy) - mocker.patch.object(qml.Device, "shots", mock_property) - - with pytest.warns( - qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" - ): - qml.execute([tape], dev, gradient_fn=param_shift, override_shots=123) - # overriden shots is the same, no change - spy.assert_not_called() - - with pytest.warns( - qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" - ): - qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100) - # overriden shots is not the same, shots were changed - spy.assert_called() - - # shots were temporarily set to the overriden value - assert spy.call_args_list[0][0] == (dev, 100) - # shots were then returned to the built-in value - assert spy.call_args_list[1][0] == (dev, 123) - - def test_overriding_device_with_shot_vector(self): - """Overriding a device that has a batch of shots set - results in original shots being returned after execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=[10, (1, 3), 5]) - - assert dev.shots == 18 - assert dev._shot_vector == [(10, 1), (1, 3), (5, 1)] - - a, b = np.array([0.543, -0.654], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - with pytest.warns( - qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" - ): - res = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=100)[0] - - assert isinstance(res, np.ndarray) - assert res.shape == () - - # device is unchanged - assert dev.shots == 18 - assert dev._shot_vector == [(10, 1), (1, 3), (5, 1)] - - res = qml.execute([tape], dev, gradient_fn=param_shift)[0] - assert len(res) == 5 - - @pytest.mark.xfail(reason="Shots vector must be adapted for new return types.") - def test_gradient_integration(self): - """Test that temporarily setting the shots works - for gradient computations""" - # TODO: Update here when shot vectors are supported - dev = qml.device("default.qubit.legacy", wires=2, shots=None) - a, b = np.array([0.543, -0.654], requires_grad=True) - - def cost_fn(a, b, shots): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - with pytest.warns( - qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" - ): - result = qml.execute([tape], dev, gradient_fn=param_shift, override_shots=shots) - return result[0] - - res = qml.jacobian(cost_fn)(a, b, shots=[10000, 10000, 10000]) - assert dev.shots is None - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == (3,) - assert res[1].shape == (3,) - - expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] - assert all( - np.allclose(np.mean(r, axis=0), e, atol=0.1, rtol=0) for r, e in zip(res, expected) - ) - - -execute_kwargs_hamiltonian = [ - {"gradient_fn": param_shift}, - {"gradient_fn": finite_diff}, -] - - -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +@pytest.mark.usefixtures("use_legacy_and_new_opmath") class TestHamiltonianWorkflows: """Test that tapes ending with expectations of Hamiltonians provide correct results and gradients""" @pytest.fixture - def cost_fn(self, execute_kwargs): + def cost_fn(self, execute_kwargs, shots, device): """Cost function for gradient tests""" - def _cost_fn(weights, coeffs1, coeffs2, dev=None): + def _cost_fn(weights, coeffs1, coeffs2): obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] H1 = qml.Hamiltonian(coeffs1, obs1) + if qml.operation.active_new_opmath(): + H1 = qml.pauli.pauli_sentence(H1).operation() obs2 = [qml.PauliZ(0)] H2 = qml.Hamiltonian(coeffs2, obs2) + if qml.operation.active_new_opmath(): + H2 = qml.pauli.pauli_sentence(H2).operation() with qml.queuing.AnnotatedQueue() as q: qml.RX(weights[0], wires=0) @@ -1274,8 +751,8 @@ def _cost_fn(weights, coeffs1, coeffs2, dev=None): qml.expval(H1) qml.expval(H2) - tape = qml.tape.QuantumScript.from_queue(q) - return autograd.numpy.hstack(qml.execute([tape], dev, **execute_kwargs)[0]) + tape = qml.tape.QuantumScript.from_queue(q, shots=shots) + return autograd.numpy.hstack(execute([tape], device, **execute_kwargs)[0]) return _cost_fn @@ -1307,246 +784,56 @@ def cost_fn_jacobian(weights, coeffs1, coeffs2): ] ) - def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): + def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): """Test hamiltonian with no trainable parameters.""" - # pylint: disable=unused-argument + + if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): + pytest.skip("adjoint differentiation does not suppport hamiltonians.") + coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=False) coeffs2 = np.array([0.7], requires_grad=False) weights = np.array([0.4, 0.5], requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev) + res = qml.jacobian(cost_fn)(weights, coeffs1, coeffs2) expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert np.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): + def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): """Test hamiltonian with trainable parameters.""" - # pylint: disable=unused-argument + if execute_kwargs["gradient_fn"] == "adjoint": + pytest.skip("trainable hamiltonians not supported with adjoint") + if qml.operation.active_new_opmath(): + pytest.skip("parameter shift derivatives do not yet support sums.") + coeffs1 = np.array([0.1, 0.2, 0.3], requires_grad=True) coeffs2 = np.array([0.7], requires_grad=True) weights = np.array([0.4, 0.5], requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2, dev=dev)) + res = np.hstack(qml.jacobian(cost_fn)(weights, coeffs1, coeffs2)) expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -class TestCustomJacobian: - """Test for custom Jacobian.""" - - def test_custom_jacobians(self): - """Test custom Jacobian device methood""" - - class CustomJacobianDevice(DefaultQubitLegacy): - @classmethod - def capabilities(cls): - capabilities = super().capabilities() - capabilities["provides_jacobian"] = True - return capabilities - - def jacobian(self, tape): - # pylint: disable=unused-argument - return np.array([1.0, 2.0, 3.0, 4.0]) - - dev = CustomJacobianDevice(wires=2) - - @qml.qnode(dev, diff_method="device") - def circuit(v): - qml.RX(v, wires=0) - return qml.probs(wires=[0, 1]) - - d_circuit = qml.jacobian(circuit, argnum=0) - - params = np.array(1.0, requires_grad=True) - - d_out = d_circuit(params) - assert np.allclose(d_out, np.array([1.0, 2.0, 3.0, 4.0])) - - def test_custom_jacobians_param_shift(self): - """Test computing the gradient using the parameter-shift - rule with a device that provides a jacobian""" - - class MyQubit(DefaultQubitLegacy): - @classmethod - def capabilities(cls): - capabilities = super().capabilities().copy() - capabilities.update( - provides_jacobian=True, - ) - return capabilities - - def jacobian(self, *args, **kwargs): - raise NotImplementedError() - - dev = MyQubit(wires=2) - - @qml.qnode(dev, diff_method="parameter-shift", grad_on_execution=False) - def qnode(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - def cost(a, b): - return autograd.numpy.hstack(qnode(a, b)) - - res = qml.jacobian(cost)(a, b) - expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) - - assert np.allclose(res[0], expected[0]) - assert np.allclose(res[1], expected[1]) - - -class SpecialObject: - """SpecialObject - A special object that conveniently encapsulates the return value of - a special observable supported by a special device and which supports - multiplication with scalars and addition. - """ - - def __init__(self, val): - self.val = val - - def __mul__(self, other): - new = SpecialObject(self.val) - new *= other - return new - - def __imul__(self, other): - self.val *= other - return self - - def __rmul__(self, other): - return self * other - - def __iadd__(self, other): - self.val += other.val if isinstance(other, self.__class__) else other - return self - - def __add__(self, other): - new = SpecialObject(self.val) - new += other.val if isinstance(other, self.__class__) else other - return new - - def __radd__(self, other): - return self + other - - -class SpecialObservable(Observable): - """SpecialObservable""" - - num_wires = AnyWires - num_params = 0 - par_domain = None - - def diagonalizing_gates(self): - """Diagonalizing gates""" - return [] - - -class DeviceSupportingSpecialObservable(DefaultQubitLegacy): - name = "Device supporting SpecialObservable" - short_name = "default.qubit.specialobservable" - observables = DefaultQubitLegacy.observables.union({"SpecialObservable"}) - - @staticmethod - def _asarray(arr, dtype=None): - # pylint: disable=unused-argument - return arr - - @classmethod - def capabilities(cls): - capabilities = super().capabilities().copy() - capabilities.update( - provides_jacobian=True, - ) - return capabilities - - def expval(self, observable, **kwargs): - if self.analytic and isinstance(observable, SpecialObservable): - val = super().expval(qml.PauliZ(wires=0), **kwargs) - return np.array(SpecialObject(val)) - - return super().expval(observable, **kwargs) - - def jacobian(self, tape): - # we actually let pennylane do the work of computing the - # jacobian for us but return it as a device jacobian - gradient_tapes, fn = qml.gradients.param_shift(tape) - tape_jacobian = fn(qml.execute(gradient_tapes, self, None)) - return tape_jacobian - - -@pytest.mark.autograd -class TestObservableWithObjectReturnType: - """Unit tests for qnode returning a custom object""" - - def test_custom_return_type(self): - """Test custom return values for a qnode""" - - dev = DeviceSupportingSpecialObservable(wires=1, shots=None) - - # force diff_method='parameter-shift' because otherwise - # PennyLane swaps out dev for default.qubit.autograd - @qml.qnode(dev, diff_method="parameter-shift") - def qnode(x): - qml.RY(x, wires=0) - return qml.expval(SpecialObservable(wires=0)) - - @qml.qnode(dev, diff_method="parameter-shift") - def reference_qnode(x): - qml.RY(x, wires=0) - return qml.expval(qml.PauliZ(wires=0)) - - out = qnode(0.2) - assert isinstance(out, np.ndarray) - assert isinstance(out.item(), SpecialObject) - assert np.isclose(out.item().val, reference_qnode(0.2)) - - def test_jacobian_with_custom_return_type(self): - """Test differentiation of a QNode on a device supporting a - special observable that returns an object rather than a number.""" - - dev = DeviceSupportingSpecialObservable(wires=1, shots=None) - - # force diff_method='parameter-shift' because otherwise - # PennyLane swaps out dev for default.qubit.autograd - @qml.qnode(dev, diff_method="parameter-shift") - def qnode(x): - qml.RY(x, wires=0) - return qml.expval(SpecialObservable(wires=0)) - - @qml.qnode(dev, diff_method="parameter-shift") - def reference_qnode(x): - qml.RY(x, wires=0) - return qml.expval(qml.PauliZ(wires=0)) - - reference_jac = (qml.jacobian(reference_qnode)(np.array(0.2, requires_grad=True)),) - - assert np.isclose( - reference_jac, - qml.jacobian(qnode)(np.array(0.2, requires_grad=True)).item().val, - ) - - # now check that also the device jacobian works with a custom return type - @qml.qnode(dev, diff_method="device") - def device_gradient_qnode(x): - qml.RY(x, wires=0) - return qml.expval(SpecialObservable(wires=0)) - - assert np.isclose( - reference_jac, - qml.jacobian(device_gradient_qnode)(np.array(0.2, requires_grad=True)).item().val, - ) + if shots.has_partitioned_shots: + pytest.xfail( + "multiple hamiltonians with shot vectors does not seem to be differentiable." + ) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/test_autograd_qnode.py b/tests/interfaces/test_autograd_qnode.py index 818d7d58ff5..a0d1e0f2e66 100644 --- a/tests/interfaces/test_autograd_qnode.py +++ b/tests/interfaces/test_autograd_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2018-2020 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 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. @@ -12,8 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the autograd interface with a QNode""" -# pylint: disable=too-many-arguments,too-few-public-methods, use-dict-literal, use-implicit-booleaness-not-comparison, -# pylint: disable=unnecessary-lambda-assignment +# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg, use-dict-literal, no-name-in-module + import autograd import autograd.numpy as anp import pytest @@ -21,32 +21,47 @@ import pennylane as qml from pennylane import numpy as np from pennylane import qnode +from pennylane.devices import DefaultQubit +from tests.param_shift_dev import ParamShiftDerivativesDevice +# dev, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], + [qml.device("default.qubit"), "finite-diff", False, False], + [qml.device("default.qubit"), "parameter-shift", False, False], + [qml.device("default.qubit"), "backprop", True, False], + [qml.device("default.qubit"), "adjoint", True, False], + [qml.device("default.qubit"), "adjoint", False, False], + [qml.device("default.qubit"), "spsa", False, False], + [qml.device("default.qubit"), "hadamard", False, False], + [ParamShiftDerivativesDevice(), "parameter-shift", False, False], + [ParamShiftDerivativesDevice(), "best", False, False], + [ParamShiftDerivativesDevice(), "parameter-shift", True, False], + [ParamShiftDerivativesDevice(), "parameter-shift", False, True], ] interface_qubit_device_and_diff_method = [ - ["autograd", "default.qubit.legacy", "finite-diff", False], - ["autograd", "default.qubit.legacy", "parameter-shift", False], - ["autograd", "default.qubit.legacy", "backprop", True], - ["autograd", "default.qubit.legacy", "adjoint", True], - ["autograd", "default.qubit.legacy", "adjoint", False], - ["autograd", "default.qubit.legacy", "spsa", False], - ["autograd", "default.qubit.legacy", "hadamard", False], - ["auto", "default.qubit.legacy", "finite-diff", False], - ["auto", "default.qubit.legacy", "parameter-shift", False], - ["auto", "default.qubit.legacy", "backprop", True], - ["auto", "default.qubit.legacy", "adjoint", True], - ["auto", "default.qubit.legacy", "adjoint", False], - ["auto", "default.qubit.legacy", "spsa", False], - ["auto", "default.qubit.legacy", "hadamard", False], + ["autograd", DefaultQubit(), "finite-diff", False, False], + ["autograd", DefaultQubit(), "parameter-shift", False, False], + ["autograd", DefaultQubit(), "backprop", True, False], + ["autograd", DefaultQubit(), "adjoint", True, False], + ["autograd", DefaultQubit(), "adjoint", False, False], + ["autograd", DefaultQubit(), "adjoint", True, True], + ["autograd", DefaultQubit(), "adjoint", False, True], + ["autograd", DefaultQubit(), "spsa", False, False], + ["autograd", DefaultQubit(), "hadamard", False, False], + ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", False, True], + ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", True, True], + ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", False, False], + ["autograd", qml.device("lightning.qubit", wires=5), "adjoint", True, False], + ["auto", DefaultQubit(), "finite-diff", False, False], + ["auto", DefaultQubit(), "parameter-shift", False, False], + ["auto", DefaultQubit(), "backprop", True, False], + ["auto", DefaultQubit(), "adjoint", True, False], + ["auto", DefaultQubit(), "adjoint", False, False], + ["auto", DefaultQubit(), "spsa", False, False], + ["auto", DefaultQubit(), "hadamard", False, False], + ["auto", qml.device("lightning.qubit", wires=5), "adjoint", False, False], + ["auto", qml.device("lightning.qubit", wires=5), "adjoint", True, False], ] pytestmark = pytest.mark.autograd @@ -57,78 +72,21 @@ @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_qubit_device_and_diff_method, ) class TestQNode: """Test that using the QNode with Autograd integrates with the PennyLane stack""" - # pylint: disable=unused-argument - - def test_nondiff_param_unwrapping( - self, interface, dev_name, diff_method, grad_on_execution, mocker + # pylint:disable=unused-argument + def test_execution_no_interface( + self, interface, dev, diff_method, grad_on_execution, device_vjp ): - """Test that non-differentiable parameters are correctly unwrapped - to NumPy ndarrays or floats (if 0-dimensional)""" - if diff_method != "parameter-shift": - pytest.skip("Test only supports parameter-shift") - - dev = qml.device("default.qubit.legacy", wires=1) - - @qnode(dev, interface=interface, diff_method=diff_method) - def circuit(x, y): - qml.RX(x[0], wires=0) - qml.Rot(*x[1:], wires=0) - qml.RY(y[0], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = np.array([0.1, 0.2, 0.3, 0.4], requires_grad=False) - y = np.array([0.5], requires_grad=True) - - param_data = [] - - def mock_apply(*args, **kwargs): - for op in args[0]: - param_data.extend(op.data) - - mocker.patch.object(dev, "apply", side_effect=mock_apply) - circuit(x, y) - assert param_data == [0.1, 0.2, 0.3, 0.4, 0.5] - assert not any(isinstance(p, np.tensor) for p in param_data) - - # test the jacobian works correctly - param_data = [] - qml.grad(circuit)(x, y) - assert param_data == [ - 0.1, - 0.2, - 0.3, - 0.4, - 0.5, - 0.1, - 0.2, - 0.3, - 0.4, - 0.5 + np.pi / 2, - 0.1, - 0.2, - 0.3, - 0.4, - 0.5 - np.pi / 2, - ] - assert not any(isinstance(p, np.tensor) for p in param_data) - - def test_execution_no_interface(self, interface, dev_name, diff_method, grad_on_execution): """Test execution works without an interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface=None, diff_method=diff_method) + @qnode(dev, interface=None) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -138,26 +96,24 @@ def circuit(a): res = circuit(a) - # without the interface, the QNode simply returns a scalar array - assert isinstance(res, np.ndarray) - assert res.shape == tuple() - - # gradients should cause an error - with pytest.raises(TypeError, match="must be real number, not ArrayBox"): - qml.grad(circuit)(a) + # without the interface, the QNode simply returns a scalar array or float + assert isinstance(res, (np.ndarray, float)) + assert qml.math.shape(res) == tuple() # pylint: disable=comparison-with-callable - def test_execution_with_interface(self, interface, dev_name, diff_method, grad_on_execution): + def test_execution_with_interface( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -165,31 +121,27 @@ def circuit(a): a = np.array(0.1, requires_grad=True) assert circuit.interface == interface - # gradients should work grad = qml.grad(circuit)(a) - assert isinstance(grad, float) assert grad.shape == tuple() - def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_jacobian(self, interface, dev, diff_method, grad_on_execution, tol, device_vjp): """Test jacobian calculation""" - num_wires = 2 kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) + if diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - elif diff_method == "hadamard": - num_wires = 3 a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -220,24 +172,23 @@ def cost(x, y): assert res[1].shape == (2,) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - def test_jacobian_no_evaluate(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_jacobian_no_evaluate( + self, interface, dev, diff_method, grad_on_execution, tol, device_vjp + ): """Test jacobian calculation when no prior circuit evaluation has been performed""" - num_wires = 2 kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - elif diff_method == "hadamard": - num_wires = 3 a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -263,36 +214,31 @@ def cost(x, y): assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execution): + def test_jacobian_options(self, interface, dev, diff_method, grad_on_execution, device_vjp): """Test setting jacobian options""" - wires = [0] - if diff_method in ["backprop", "adjoint"]: - pytest.skip("Test does not support backprop or adjoint method") - elif diff_method == "finite-diff": - kwargs = {"h": 1e-8, "approx_order": 2} - elif diff_method == "parameter-shift": - kwargs = {"shifts": [(0.1,), (0.2,)]} - elif diff_method == "hadamard": - wires = [0, "aux"] - kwargs = {"aux_wire": qml.wires.Wires("aux"), "device_wires": wires} - else: - kwargs = {} + if diff_method != "finite-diff": + pytest.skip("Test only supports finite diff.") a = np.array([0.1, 0.2], requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=wires) - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) + @qnode( + dev, + interface=interface, + h=1e-8, + order=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) - circuit(a) - qml.jacobian(circuit)(a) - def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_changing_trainability( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -301,9 +247,13 @@ def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_e a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - dev = qml.device("default.qubit.legacy", wires=2) - - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, + interface=interface, + diff_method="parameter-shift", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -340,21 +290,18 @@ def loss(a, b): circuit(a, b) assert circuit.qtape.trainable_params == [1] - def test_classical_processing(self, interface, dev_name, diff_method, grad_on_execution): + def test_classical_processing(self, interface, dev, diff_method, grad_on_execution, device_vjp): """Test classical processing within the quantum tape""" a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=False) c = np.array(0.3, requires_grad=True) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -372,12 +319,17 @@ def circuit(a, b, c): assert res[0].shape == () assert res[1].shape == () - def test_no_trainable_parameters(self, interface, dev_name, diff_method, grad_on_execution): + def test_no_trainable_parameters( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): """Test evaluation and Jacobian if there are no trainable parameters""" - dev = qml.device(dev_name, wires=2) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -396,35 +348,34 @@ def circuit(a, b): assert len(res) == 2 assert isinstance(res, tuple) - def cost0(x, y): + def cost(x, y): return autograd.numpy.hstack(circuit(x, y)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - assert not qml.jacobian(cost0)(a, b) + assert not qml.jacobian(cost)(a, b) - def cost1(a, b): + def cost2(a, b): return np.sum(circuit(a, b)) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - grad = qml.grad(cost1)(a, b) + grad = qml.grad(cost2)(a, b) assert grad == tuple() - def test_matrix_parameter(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_matrix_parameter( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that the autograd interface works correctly with a matrix parameter""" U = np.array([[0, 1], [1, 0]], requires_grad=False) a = np.array(0.1, requires_grad=True) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -440,13 +391,18 @@ def circuit(U, a): assert np.allclose(res, np.sin(a), atol=tol, rtol=0) def test_gradient_non_differentiable_exception( - self, interface, dev_name, diff_method, grad_on_execution + self, interface, dev, diff_method, grad_on_execution, device_vjp ): """Test that an exception is raised if non-differentiable data is differentiated""" - dev = qml.device(dev_name, wires=2) - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(data1): qml.templates.AmplitudeEmbedding(data1, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) @@ -457,18 +413,27 @@ def circuit(data1): with pytest.raises(qml.numpy.NonDifferentiableError, match="is non-differentiable"): grad_fn(data1) - def test_differentiable_expand(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_differentiable_expand( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that operation and nested tape expansion is differentiable""" kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) + if diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=20) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 10 tol = TOL_FOR_SPSA + # pylint: disable=too-few-public-methods class U3(qml.U3): + """Custom U3.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -477,7 +442,6 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - dev = qml.device(dev_name, wires=2) a = np.array(0.1, requires_grad=False) p = np.array([0.1, 0.2, 0.3], requires_grad=True) @@ -491,14 +455,14 @@ def circuit(a, p): expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) ) - assert isinstance(res, np.ndarray) - assert res.shape == () + # assert isinstance(res, np.ndarray) + # assert res.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) res = qml.grad(circuit)(a, p) - assert isinstance(res, np.ndarray) - assert len(res) == 3 + # assert isinstance(res, np.ndarray) + # assert len(res) == 3 expected = np.array( [ @@ -517,9 +481,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and remains differentiable.""" - def test_changing_shots(self, mocker, tol): + def test_changing_shots(self): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = np.array([0.543, -0.654], requires_grad=True) @qnode(dev, diff_method=qml.gradients.param_shift) @@ -527,31 +491,21 @@ def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + res = circuit(a, b) # execute with shots=100 - res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # same single call performed above + res = circuit(a, b, shots=100) + assert res.shape == (100, 2) # pylint: disable=comparison-with-callable @pytest.mark.xfail(reason="Param shift and shot vectors.") def test_gradient_integration(self): """Test that temporarily setting the shots works for gradient computations""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = np.array([0.543, -0.654], requires_grad=True) @qnode(dev, diff_method=qml.gradients.param_shift) @@ -574,62 +528,57 @@ def cost_fn(a, b): def test_update_diff_method(self, mocker): """Test that temporarily setting the shots updates the diff method""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = np.array([0.543, -0.654], requires_grad=True) spy = mocker.spy(qml, "execute") - @qnode(dev) + @qnode(DefaultQubit()) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - cost_fn(a, b) + assert cost_fn.gradient_fn == "backprop" # gets restored to default + + cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn is qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn is qml.gradients.param_shift - # if we set the shots to None, backprop can now be used - cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg - assert spy.call_args[1]["gradient_fn"] == "backprop" - assert cost_fn.gradient_fn == "backprop" - + # if we use the default shots value of None, backprop can now be used cost_fn(a, b) - assert cost_fn.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] == "backprop" @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_qubit_device_and_diff_method, ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - # pylint: disable=unused-argument - def test_probability_differentiation( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Tests correct output shape and evaluation for a tape with a single prob output""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measuring probabilities with adjoint.") + kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -650,25 +599,23 @@ def circuit(x, y): assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) def test_multiple_probability_differentiation( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measuring probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -717,22 +664,24 @@ def cost(x, y): ) assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) - def test_ragged_differentiation(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_ragged_differentiation( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measuring probabilities with adjoint.") + kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - num_wires = 2 - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - elif diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -768,22 +717,25 @@ def cost(x, y): assert np.allclose(res[1], expected[1], atol=tol, rtol=0) def test_ragged_differentiation_variance( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measuring probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") - dev = qml.device(dev_name, wires=2) x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @@ -804,12 +756,12 @@ def circuit(x, y): assert isinstance(res, tuple) assert len(res) == 2 - assert isinstance(res[0], np.ndarray) - assert res[0].shape == () + # assert isinstance(res[0], np.ndarray) + # assert res[0].shape == () assert np.allclose(res[0], expected_var, atol=tol, rtol=0) - assert isinstance(res[1], np.ndarray) - assert res[1].shape == (2,) + # assert isinstance(res[1], np.ndarray) + # assert res[1].shape == (2,) assert np.allclose(res[1], expected_probs, atol=tol, rtol=0) def cost(x, y): @@ -826,33 +778,38 @@ def cost(x, y): assert isinstance(jac, tuple) assert len(jac) == 2 - assert isinstance(jac[0], np.ndarray) - assert jac[0].shape == (3,) + # assert isinstance(jac[0], np.ndarray) + # assert jac[0].shape == (3,) assert np.allclose(jac[0], expected[0], atol=tol, rtol=0) - assert isinstance(jac[1], np.ndarray) - assert jac[1].shape == (3,) + # assert isinstance(jac[1], np.ndarray) + # assert jac[1].shape == (3,) assert np.allclose(jac[1], expected[1], atol=tol, rtol=0) - def test_chained_qnodes(self, interface, dev_name, diff_method, grad_on_execution): + def test_chained_qnodes(self, interface, dev, diff_method, grad_on_execution, device_vjp): """Test that the gradient of chained QNodes works without error""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + # pylint: disable=too-few-public-methods class Template(qml.templates.StronglyEntanglingLayers): + """Custom template.""" + def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + ) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) - @qnode(dev, interface=interface, diff_method=diff_method) + @qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit2(data, weights): qml.templates.AngleEmbedding(data, wires=[0, 1]) Template(weights, wires=[0, 1]) @@ -876,19 +833,22 @@ def cost(w1, w2): assert len(res) == 2 - def test_chained_gradient_value(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_chained_gradient_value( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that the returned gradient value for two chained qubit QNodes is correct.""" - kwargs = dict(interface=interface, diff_method=diff_method) + kwargs = dict( + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 3 - - if diff_method == "hadamard": - num_wires = 4 - - dev1 = qml.device(dev_name, wires=num_wires) + dev1 = qml.device("default.qubit") @qnode(dev1, **kwargs) def circuit1(a, b, c): @@ -899,9 +859,11 @@ def circuit1(a, b, c): qml.CNOT(wires=[1, 2]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(2)) - dev2 = qml.device("default.qubit.legacy", wires=num_wires) + dev2 = dev - @qnode(dev2, interface=interface, diff_method=diff_method) + @qnode( + dev2, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + ) def circuit2(data, weights): qml.RX(data[0], wires=0) qml.RX(data[1], wires=1) @@ -971,19 +933,20 @@ def cost(a, b, c, weights): # to the first parameter of circuit1. assert circuit1.qtape.trainable_params == [1, 2] - def test_second_derivative(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_second_derivative( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test second derivative calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, max_diff=2, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -1013,18 +976,17 @@ def circuit(x): assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1039,8 +1001,8 @@ def circuit(x): expected_res = np.cos(a) * np.cos(b) - assert isinstance(res, np.ndarray) - assert res.shape == () + # assert isinstance(res, np.ndarray) + # assert res.shape == () assert np.allclose(res, expected_res, atol=tol, rtol=0) grad_fn = qml.grad(circuit) @@ -1068,19 +1030,18 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_unused_parameter( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1109,19 +1070,20 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian_vector_valued( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test hessian calculation of a vector valued QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1137,7 +1099,7 @@ def circuit(x): expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] assert isinstance(res, np.ndarray) - assert res.shape == (2,) + assert res.shape == (2,) # pylint: disable=comparison-with-callable assert np.allclose(res, expected_res, atol=tol, rtol=0) jac_fn = qml.jacobian(circuit) @@ -1174,19 +1136,18 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1229,19 +1190,18 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(a, b): @@ -1255,7 +1215,7 @@ def circuit(a, b): expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] assert isinstance(res, np.ndarray) - assert res.shape == (2,) + assert res.shape == (2,) # pylint: disable=comparison-with-callable assert np.allclose(res, expected_res, atol=tol, rtol=0) jac_fn = qml.jacobian(circuit) @@ -1272,8 +1232,12 @@ def circuit(a, b): assert g[1].shape == (2,) assert np.allclose(g[1], expected_g[1], atol=tol, rtol=0) - jac_fn_a = lambda *args: jac_fn(*args)[0] - jac_fn_b = lambda *args: jac_fn(*args)[1] + def jac_fn_a(*args): + return jac_fn(*args)[0] + + def jac_fn_b(*args): + return jac_fn(*args)[1] + hess_a = qml.jacobian(jac_fn_a)(a, b) hess_b = qml.jacobian(jac_fn_b)(a, b) assert isinstance(hess_a, tuple) and len(hess_a) == 2 @@ -1295,18 +1259,17 @@ def circuit(a, b): assert np.allclose(hess[0], exp_hess[0], atol=tol, rtol=0) assert np.allclose(hess[1], exp_hess[1], atol=tol, rtol=0) - def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test hessian calculation of a ragged QNode""" if diff_method not in {"parameter-shift", "backprop"}: pytest.skip("Test only supports parameter-shift or backprop") - dev = qml.device(dev_name, wires=2) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, ) def circuit(x): @@ -1320,14 +1283,18 @@ def circuit(x): a, b = x - cos_prod = np.cos(a) * np.cos(b) - expected_res = (cos_prod, [0.5 + 0.5 * cos_prod, 0.5 - 0.5 * cos_prod]) - res = circuit(x) - assert all(qml.math.allclose(r, e) for r, e in zip(res, expected_res)) + expected_res = [ + np.cos(a) * np.cos(b), + 0.5 + 0.5 * np.cos(a) * np.cos(b), + 0.5 - 0.5 * np.cos(a) * np.cos(b), + ] def cost_fn(x): return autograd.numpy.hstack(circuit(x)) + res = cost_fn(x) + assert qml.math.allclose(res, expected_res) + jac_fn = qml.jacobian(cost_fn) hess = qml.jacobian(jac_fn)(x) @@ -1351,18 +1318,21 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - dev = qml.device(dev_name, wires=2) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("Lightning does not support state adjoint diff.") x = np.array(0.543, requires_grad=True) y = np.array(-0.654, requires_grad=True) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1393,20 +1363,25 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector(self, state, interface, dev_name, diff_method, grad_on_execution, tol): + def test_projector( + self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that the variance of a projector is correctly returned""" + if diff_method == "adjoint": + pytest.skip("adjoint supports either expvals or diagonal measurements.") kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support projectors") - elif diff_method == "spsa": + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") - dev = qml.device(dev_name, wires=2) P = np.array(state, requires_grad=False) x, y = np.array([0.765, -0.654], requires_grad=True) @@ -1419,8 +1394,8 @@ def circuit(x, y): res = circuit(x, y) expected = 0.25 * np.sin(x / 2) ** 2 * (3 + np.cos(2 * y) + 2 * np.cos(x) * np.sin(y) ** 2) - assert isinstance(res, np.ndarray) - assert res.shape == () + # assert isinstance(res, np.ndarray) + # assert res.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) jac = qml.jacobian(circuit)(x, y) @@ -1444,106 +1419,60 @@ def circuit(x, y): assert np.allclose(jac, expected, atol=tol, rtol=0) + def test_postselection_differentiation( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): + """Test that when postselecting with default.qubit, differentiation works correctly.""" -@pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": 0.05}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - elif diff_method == "hadamard": - pytest.skip("Hadamard gradient does not support variances.") - - r = np.array(0.543, requires_grad=True) - phi = np.array(-0.654, requires_grad=True) - - @qnode(dev, diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - res = circuit(r, phi) - expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") - # circuit jacobians - res = qml.jacobian(circuit)(r, phi) - expected = np.array( - [ - [ - 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, - 2 * np.sinh(2 * r) * np.sin(2 * phi), - ] - ] + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_second_order_observable(self, diff_method, kwargs, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - tol = TOL_FOR_SPSA - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - elif diff_method == "hadamard": - pytest.skip("Hadamard gradient does not support variances.") - - n = np.array(0.12, requires_grad=True) - a = np.array(0.765, requires_grad=True) - - @qnode(dev, diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) - - res = circuit(n, a) - expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = qml.jacobian(circuit)(n, a) - expected = np.array([[2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -def test_adjoint_reuse_device_state(mocker): - """Tests that the autograd interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=1) - - @qnode(dev, diff_method="adjoint") - def circ(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) - spy = mocker.spy(dev, "adjoint_jacobian") + phi = np.array(1.23, requires_grad=True) + theta = np.array(4.56, requires_grad=True) - qml.grad(circ, argnum=0)(1.0) - assert circ.device.num_executions == 1 + assert np.allclose(circuit(phi, theta), expected_circuit(theta)) - spy.assert_called_with(mocker.ANY, use_device_state=True) + gradient = qml.grad(circuit)(phi, theta) + exp_theta_grad = qml.grad(expected_circuit)(theta) + assert np.allclose(gradient, [0.0, exp_theta_grad]) -@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) +@pytest.mark.parametrize( + "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method +) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the Autograd interface""" @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff + self, dev, diff_method, grad_on_execution, max_diff, device_vjp ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" @@ -1552,20 +1481,22 @@ def test_gradient_expansion_trainable_only( if max_diff == 2 and diff_method == "hadamard": pytest.skip("Max diff > 1 not supported for Hadamard gradient.") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - + # pylint: disable=too-few-public-methods class PhaseShift(qml.PhaseShift): + """dummy phase shift.""" + grad_method = None def decomposition(self): return [qml.RY(3 * self.data[0], wires=self.wires)] - @qnode(dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff) + @qnode( + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + max_diff=max_diff, + device_vjp=device_vjp, + ) def circuit(x, y): qml.Hadamard(wires=0) PhaseShift(x, wires=0) @@ -1576,11 +1507,11 @@ def circuit(x, y): y = np.array(0.7, requires_grad=False) circuit(x, y) - qml.grad(circuit)(x, y) + _ = qml.grad(circuit)(x, y) @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, tol + self, dev, diff_method, grad_on_execution, max_diff, tol, device_vjp ): """Test that if there are non-commuting groups and the number of shots is None the first and second order gradients are correctly evaluated""" @@ -1588,15 +1519,16 @@ def test_hamiltonian_expansion_analytic( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, + device_vjp=device_vjp, ) + if diff_method in ["adjoint", "hadamard"]: pytest.skip("The diff method requested does not yet support Hamiltonians") elif diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 10 tol = TOL_FOR_SPSA - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} - dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1626,7 +1558,11 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method in ("parameter-shift", "backprop") and max_diff == 2: + if ( + diff_method in ("parameter-shift", "backprop") + and max_diff == 2 + and dev.name != "param_shift.qubit" + ): if diff_method == "backprop": with pytest.warns(UserWarning, match=r"Output seems independent of input."): grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c) @@ -1644,28 +1580,26 @@ def circuit(data, weights, coeffs): @pytest.mark.slow @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, max_diff, mocker + def test_hamiltonian_finite_shots( + self, dev, diff_method, grad_on_execution, max_diff, device_vjp ): - """Test that the Hamiltonian is expanded if there + """Test that the Hamiltonian is correctly measured if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.3 + tol = 0.1 if diff_method in ("adjoint", "backprop", "hadamard"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": - gradient_kwargs = dict( - h=H_FOR_SPSA, - sampler_rng=np.random.default_rng(SEED_FOR_SPSA), - num_directions=20, - ) + gradient_kwargs = { + "h": H_FOR_SPSA, + "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), + "num_directions": 10, + } tol = TOL_FOR_SPSA elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} - dev = qml.device(dev_name, wires=3, shots=50000) - spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1673,6 +1607,7 @@ def test_hamiltonian_expansion_finite_shots( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1688,13 +1623,15 @@ def circuit(data, weights, coeffs): c = np.array([-0.6543, 0.24, 0.54], requires_grad=True) # test output - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_called() # test gradients - grad = qml.grad(circuit)(d, w, c) + if diff_method in ["finite-diff", "spsa"]: + pytest.skip(f"{diff_method} not compatible") + + grad = qml.grad(circuit)(d, w, c, shots=50000) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1704,12 +1641,11 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2: - with pytest.warns(UserWarning, match=r"Output seems independent of input."): - grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c) + if diff_method == "parameter-shift" and max_diff == 2 and dev.name != "param_shift.qubit": + grad2_c = qml.jacobian(qml.grad(circuit, argnum=2), argnum=2)(d, w, c, shots=50000) assert np.allclose(grad2_c, 0, atol=tol) - grad2_w_c = qml.jacobian(qml.grad(circuit, argnum=1), argnum=2)(d, w, c) + grad2_w_c = qml.jacobian(qml.grad(circuit, argnum=1), argnum=2)(d, w, c, shots=50000) expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ 0, -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), @@ -1723,36 +1659,38 @@ class TestSample: def test_backprop_error(self): """Test that sampling in backpropagation grad_on_execution raises an error""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = DefaultQubit() @qnode(dev, diff_method="backprop") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - with pytest.raises(qml.QuantumFunctionError, match="only supported when shots=None"): - circuit(shots=10) # pylint: disable=unexpected-keyword-arg + with pytest.raises( + qml.QuantumFunctionError, match="does not support backprop with requested" + ): + circuit(shots=10) def test_sample_dimension(self): """Test that the sample function outputs samples of the right size""" n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=2, shots=n_sample) + dev = DefaultQubit() - @qnode(dev) + @qnode(dev, diff_method=None) def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=n_sample) assert isinstance(res, tuple) assert len(res) == 2 - assert res[0].shape == (10,) + assert res[0].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(res[0], np.ndarray) - assert res[1].shape == (10,) + assert res[1].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(res[1], np.ndarray) def test_sample_combination(self): @@ -1760,7 +1698,7 @@ def test_sample_combination(self): n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) + dev = DefaultQubit() @qnode(dev, diff_method="parameter-shift") def circuit(): @@ -1768,29 +1706,29 @@ def circuit(): return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit() + result = circuit(shots=n_sample) assert isinstance(result, tuple) assert len(result) == 3 assert np.array_equal(result[0].shape, (n_sample,)) - assert isinstance(result[1], np.ndarray) - assert isinstance(result[2], np.ndarray) + assert isinstance(result[1], (float, np.ndarray)) + assert isinstance(result[2], (float, np.ndarray)) assert result[0].dtype == np.dtype("float") def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) + dev = DefaultQubit() - @qnode(dev) + @qnode(dev, diff_method=None) def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit() + result = circuit(shots=n_sample) assert isinstance(result, np.ndarray) assert np.array_equal(result.shape, (n_sample,)) @@ -1800,44 +1738,44 @@ def test_multi_wire_sample_regular_shape(self): where a rectangular array is expected""" n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) + dev = DefaultQubit() - @qnode(dev) + @qnode(dev, diff_method=None) def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit() + result = circuit(shots=n_sample) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) assert len(result) == 3 - assert result[0].shape == (10,) + assert result[0].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(result[0], np.ndarray) - assert result[1].shape == (10,) + assert result[1].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(result[1], np.ndarray) - assert result[2].shape == (10,) + assert result[2].shape == (10,) # pylint: disable=comparison-with-callable assert isinstance(result[2], np.ndarray) -@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) +@pytest.mark.parametrize( + "dev,diff_method,grad_on_execution,device_vjp", qubit_device_and_diff_method +) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" - # pylint: disable=unused-argument - - def test_grad_single_measurement_param(self, dev_name, diff_method, grad_on_execution): + def test_grad_single_measurement_param(self, dev, diff_method, grad_on_execution, device_vjp): """For one measurement and one param, the gradient is a float.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1847,25 +1785,20 @@ def circuit(a): grad = qml.grad(circuit)(a) - import sys + assert isinstance(grad, np.tensor if diff_method == "backprop" else float) - python_version = sys.version_info.minor - if diff_method == "backprop" and python_version > 7: - # Since numpy 1.23.0 - assert isinstance(grad, np.ndarray) - else: - assert isinstance(grad, float) - - def test_grad_single_measurement_multiple_param(self, dev_name, diff_method, grad_on_execution): + def test_grad_single_measurement_multiple_param( + self, dev, diff_method, grad_on_execution, device_vjp + ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1882,17 +1815,17 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1907,21 +1840,18 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1935,21 +1865,18 @@ def circuit(a): assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1969,20 +1896,17 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """For a multi dimensional measurement (probs), check that a single array is returned.""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1995,20 +1919,17 @@ def circuit(a): assert jac.shape == (4, 2) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """The jacobian of multiple measurements with a single params return an array.""" - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -2025,21 +1946,17 @@ def cost(x): assert jac.shape == (5,) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -2063,21 +1980,17 @@ def cost(x, y): assert jac[1].shape == (5,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - @qnode(dev, interface="autograd", diff_method=diff_method) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2093,14 +2006,8 @@ def cost(x): assert isinstance(jac, np.ndarray) assert jac.shape == (5, 2) - def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_execution): + def test_hessian_expval_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2108,7 +2015,14 @@ def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_exe par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2129,22 +2043,24 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (2,) - def test_hessian_expval_multiple_param_array(self, dev_name, diff_method, grad_on_execution): + def test_hessian_expval_multiple_param_array( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of single measurement with a multiple params array return a single array.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2156,9 +2072,8 @@ def circuit(x): assert isinstance(hess, np.ndarray) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution): + def test_hessian_var_multiple_params(self, dev, diff_method, grad_on_execution, device_vjp): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2168,7 +2083,14 @@ def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execut par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2189,18 +2111,25 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (2,) - def test_hessian_var_multiple_param_array(self, dev_name, diff_method, grad_on_execution): + def test_hessian_var_multiple_param_array( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") - dev = qml.device(dev_name, wires=2) - params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2212,19 +2141,24 @@ def circuit(x): assert isinstance(hess, np.ndarray) assert hess.shape == (2, 2) - def test_hessian_probs_expval_multiple_params(self, dev_name, diff_method, grad_on_execution): + def test_hessian_probs_expval_multiple_params( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint method does not currently support second-order diff.") par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2249,18 +2183,23 @@ def cost(x, y): assert hess[1].shape == (6,) def test_hessian_expval_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution + self, dev, diff_method, grad_on_execution, device_vjp ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint method does not currently support second-order diff.") - dev = qml.device(dev_name, wires=2) - params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2273,11 +2212,12 @@ def cost(x): hess = qml.jacobian(qml.jacobian(cost))(params) assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) + assert hess.shape == (3, 2, 2) # pylint: disable=no-member - def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, grad_on_execution): + def test_hessian_probs_var_multiple_params( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") @@ -2287,7 +2227,14 @@ def test_hessian_probs_var_multiple_params(self, dev_name, diff_method, grad_on_ par_0 = qml.numpy.array(0.1, requires_grad=True) par_1 = qml.numpy.array(0.2, requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -2311,18 +2258,25 @@ def cost(x, y): assert isinstance(hess[1], np.ndarray) assert hess[1].shape == (6,) - def test_hessian_var_probs_multiple_param_array(self, dev_name, diff_method, grad_on_execution): + def test_hessian_var_multiple_param_array2( + self, dev, diff_method, grad_on_execution, device_vjp + ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": pytest.skip("The adjoint method does not currently support second-order diff.") elif diff_method == "hadamard": pytest.skip("Hadamard gradient does not support variances.") - dev = qml.device(dev_name, wires=2) - params = qml.numpy.array([0.1, 0.2], requires_grad=True) - @qnode(dev, interface="autograd", diff_method=diff_method, max_diff=2) + @qnode( + dev, + interface="autograd", + diff_method=diff_method, + max_diff=2, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(x): qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -2335,15 +2289,14 @@ def cost(x): hess = qml.jacobian(qml.jacobian(cost))(params) assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) + assert hess.shape == (3, 2, 2) # pylint: disable=no-member -@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) -def test_no_ops(dev_name): +def test_no_ops(): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = qml.device(dev_name, wires=1) + dev = DefaultQubit() @qml.qnode(dev, interface="autograd") def circuit(): diff --git a/tests/interfaces/test_autograd_qnode_shot_vector.py b/tests/interfaces/test_autograd_qnode_shot_vector.py index 6aee767d01e..87e534fdb52 100644 --- a/tests/interfaces/test_autograd_qnode_shot_vector.py +++ b/tests/interfaces/test_autograd_qnode_shot_vector.py @@ -29,9 +29,9 @@ spsa_kwargs = {"h": 0.05, "num_directions": 20, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 0.05}], - ["default.qubit.legacy", "parameter-shift", {}], - ["default.qubit.legacy", "spsa", spsa_kwargs], + ["default.qubit", "finite-diff", {"h": 0.05}], + ["default.qubit", "parameter-shift", {}], + ["default.qubit", "spsa", spsa_kwargs], ] TOLS = { diff --git a/tests/interfaces/test_execute.py b/tests/interfaces/test_execute.py index e12c5d8ff55..87201a27618 100644 --- a/tests/interfaces/test_execute.py +++ b/tests/interfaces/test_execute.py @@ -1,4 +1,4 @@ -# Copyright 2023 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 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. @@ -11,18 +11,219 @@ # 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. -""" -Interface independent tests for qml.execute -""" - +"""Tests for exeuction with default qubit 2 independent of any interface.""" import pytest import pennylane as qml +from pennylane import numpy as np +from pennylane.devices import DefaultQubit +from pennylane.workflow.execution import _preprocess_expand_fn + + +class TestPreprocessExpandFn: + """Tests the _preprocess_expand_fn helper function.""" + + def test_provided_is_callable(self): + """Test that if the expand_fn is not "device", it is simply returned.""" + + dev = DefaultQubit() + + def f(tape): + return tape + + out = _preprocess_expand_fn(f, dev, 10) + assert out is f + + def test_new_device_blank_expand_fn(self): + """Test that the expand_fn is blank if is new device.""" + + dev = DefaultQubit() + + out = _preprocess_expand_fn("device", dev, 10) + + x = [1] + assert out(x) is x + + +class TestBatchTransformHelper: + """Unit tests for the _batch_transform helper function.""" + + def test_warns_if_requested_off(self): + """Test that a warning is raised if the the batch transform is requested to not be used.""" + + # pylint: disable=too-few-public-methods + class CustomOp(qml.operation.Operator): + """Dummy operator.""" + + def decomposition(self): + return [qml.PauliX(self.wires[0])] + + dev = DefaultQubit() + + qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"): + program, _ = dev.preprocess() + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + qml.execute( + (qs, qs), device=dev, device_batch_transform=False, transform_program=program + ) + + def test_split_and_expand_performed(self): + """Test that preprocess returns the correct tapes when splitting and expanding + is needed.""" + + class NoMatOp(qml.operation.Operation): + """Dummy operation for expanding circuit.""" + + # pylint: disable=missing-function-docstring + num_wires = 1 + + # pylint: disable=arguments-renamed, invalid-overridden-method + @property + def has_matrix(self): + return False + + def decomposition(self): + return [qml.PauliX(self.wires), qml.PauliY(self.wires)] + + ops = [qml.Hadamard(0), NoMatOp(1), qml.RX([np.pi, np.pi / 2], wires=1)] + # Need to specify grouping type to transform tape + measurements = [qml.expval(qml.PauliX(0)), qml.expval(qml.PauliZ(1))] + tapes = [ + qml.tape.QuantumScript(ops=ops, measurements=[measurements[0]]), + qml.tape.QuantumScript(ops=ops, measurements=[measurements[1]]), + ] + + dev = DefaultQubit() + config = qml.devices.ExecutionConfig(gradient_method="adjoint") + + program, new_config = dev.preprocess(config) + res_tapes, batch_fn = program(tapes) + expected_ops = [ + [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi, wires=1)], + [qml.Hadamard(0), qml.PauliX(1), qml.PauliY(1), qml.RX(np.pi / 2, wires=1)], + ] + + assert len(res_tapes) == 4 + for i, t in enumerate(res_tapes): + for op, expected_op in zip(t.operations, expected_ops[i % 2]): + qml.assert_equal(op, expected_op) + assert len(t.measurements) == 1 + if i < 2: + qml.assert_equal(t.measurements[0], measurements[0]) + else: + qml.assert_equal(t.measurements[0], measurements[1]) + + input = ([[1, 2]], [[3, 4]], [[5, 6]], [[7, 8]]) + assert np.array_equal(batch_fn(input), np.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])) + + assert new_config.grad_on_execution + assert new_config.use_device_gradient + + +def test_warning_if_not_device_batch_transform(): + """Test that a warning is raised if the users requests to not run device batch transform.""" + + # pylint: disable=too-few-public-methods + class CustomOp(qml.operation.Operator): + """Dummy operator.""" + + def decomposition(self): + return [qml.PauliX(self.wires[0])] + + dev = DefaultQubit() + + qs = qml.tape.QuantumScript([CustomOp(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns(UserWarning, match="device batch transforms cannot be turned off"): + program, _ = dev.preprocess() + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + results = qml.execute( + [qs], dev, device_batch_transform=False, transform_program=program + ) + + assert len(results) == 1 + assert qml.math.allclose(results[0], -1) + + +@pytest.mark.parametrize("gradient_fn", (None, "backprop", qml.gradients.param_shift)) +def test_caching(gradient_fn): + """Test that cache execute returns the cached result if the same script is executed + multiple times, both in multiple times in a batch and in separate batches.""" + dev = DefaultQubit() + + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) + + cache = {} + + with qml.Tracker(dev) as tracker: + results = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn) + results2 = qml.execute([qs, qs], dev, cache=cache, gradient_fn=gradient_fn) + + assert len(cache) == 1 + assert cache[qs.hash] == -1.0 + + assert list(results) == [-1.0, -1.0] + assert list(results2) == [-1.0, -1.0] + + assert tracker.totals["batches"] == 1 + assert tracker.totals["executions"] == 1 + assert cache[qs.hash] == -1.0 + + +class TestExecuteDeprecations: + """Class to test deprecation warnings in qml.execute. Warnings should be raised even if the default value is used.""" + + @pytest.mark.parametrize("expand_fn", (None, lambda qs: qs, "device")) + def test_expand_fn_is_deprecated(self, expand_fn): + """Test that expand_fn argument of qml.execute is deprecated.""" + dev = DefaultQubit() + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The expand_fn argument is deprecated" + ): + # None is a value used for expand_fn in the QNode + qml.execute([qs], dev, expand_fn=expand_fn) + + def test_max_expansion_is_deprecated(self): + """Test that max_expansion argument of qml.execute is deprecated.""" + dev = DefaultQubit() + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The max_expansion argument is deprecated" + ): + qml.execute([qs], dev, max_expansion=10) + + @pytest.mark.parametrize("override_shots", (False, 10)) + def test_override_shots_is_deprecated(self, override_shots): + """Test that override_shots argument of qml.execute is deprecated.""" + dev = DefaultQubit() + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, match="The override_shots argument is deprecated" + ): + qml.execute([qs], dev, override_shots=override_shots) + + @pytest.mark.parametrize("device_batch_transform", (False, True)) + def test_device_batch_transform_is_deprecated(self, device_batch_transform): + """Test that device_batch_transform argument of qml.execute is deprecated.""" + # Need to use legacy device, otherwise another warning would be raised due to new Device interface + dev = qml.device("default.qubit.legacy", wires=1) + qs = qml.tape.QuantumScript([qml.PauliX(0)], [qml.expval(qml.PauliZ(0))]) -def test_old_interface_no_device_jacobian_products(): - """Test that an error is always raised for the old device interface if device jacobian products are requested.""" - dev = qml.device("default.qubit.legacy", wires=2) - tape = qml.tape.QuantumScript([qml.RX(1.0, wires=0)], [qml.expval(qml.PauliZ(0))]) - with pytest.raises(qml.QuantumFunctionError): - qml.execute((tape,), dev, device_vjp=True) + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="The device_batch_transform argument is deprecated", + ): + qml.execute([qs], dev, device_batch_transform=device_batch_transform) diff --git a/tests/interfaces/test_jacobian_products.py b/tests/interfaces/test_jacobian_products.py index 07f3ccb60ab..44510e71879 100644 --- a/tests/interfaces/test_jacobian_products.py +++ b/tests/interfaces/test_jacobian_products.py @@ -31,7 +31,8 @@ ) dev = qml.device("default.qubit") -dev_old = qml.device("default.qubit.legacy", wires=5) +with pytest.warns(qml.PennyLaneDeprecationWarning): + dev_old = qml.device("default.qubit.legacy", wires=5) dev_lightning = qml.device("lightning.qubit", wires=5) adjoint_config = qml.devices.ExecutionConfig(gradient_method="adjoint") dev_ps = ParamShiftDerivativesDevice() diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py index e9b7fdd8904..519c0daa028 100644 --- a/tests/interfaces/test_jax.py +++ b/tests/interfaces/test_jax.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 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. @@ -11,606 +11,456 @@ # 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. -"""Unit tests for the JAX-Python interface""" +"""Jax specific tests for execute and default qubit 2.""" import numpy as np - -# pylint: disable=protected-access,too-few-public-methods import pytest +from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml from pennylane import execute +from pennylane.devices import DefaultQubit from pennylane.gradients import param_shift -from pennylane.typing import TensorLike - -pytestmark = pytest.mark.jax +from pennylane.measurements import Shots jax = pytest.importorskip("jax") +jnp = pytest.importorskip("jax.numpy") jax.config.update("jax_enable_x64", True) +pytestmark = pytest.mark.jax -class TestJaxExecuteUnitTests: - """Unit tests for jax execution""" - - def test_jacobian_options(self, mocker): - """Test setting jacobian options""" - spy = mocker.spy(qml.gradients, "param_shift") - - a = jax.numpy.array([0.1, 0.2]) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - device, - gradient_fn=param_shift, - gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, - )[0] - - jax.grad(cost)(a, device=dev) - - for args in spy.call_args_list: - assert args[1]["shifts"] == [(np.pi / 4,)] * 2 - - def test_incorrect_grad_on_execution(self): - """Test that an error is raised if an gradient transform - is used with grad_on_execution=True""" - a = jax.numpy.array([0.1, 0.2]) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - device, - gradient_fn=param_shift, - grad_on_execution=True, - )[0] - - with pytest.raises( - ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" - ): - jax.grad(cost)(a, device=dev) - - def test_unknown_interface(self): - """Test that an error is raised if the interface is unknown""" - a = jax.numpy.array([0.1, 0.2]) - - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - device, - gradient_fn=param_shift, - interface="None", - )[0] - - with pytest.raises(ValueError, match="Unknown interface"): - cost(a, device=dev) - - def test_grad_on_execution(self, mocker): - """Test that grad_on_execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=2) - spy = mocker.spy(dev, "execute_and_gradients") - - def cost(params): - tape1 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], - ) - - tape2 = qml.tape.QuantumScript( - [qml.RY(np.array(0.5), 0)], - [qml.expval(qml.PauliZ(0))], - ) - - tape3 = qml.tape.QuantumScript( - [qml.RY(params[0], 0), qml.RX(params[1], 0)], - [qml.expval(qml.PauliZ(0))], - ) - return execute( - [tape1, tape2, tape3], - dev, - gradient_fn="device", - gradient_kwargs={ - "method": "adjoint_jacobian", - "use_device_state": True, - }, - ) - - a = jax.numpy.array([0.1, 0.2]) - res = cost(a) - - x, y = a - assert np.allclose(res[0][0], np.cos(x) * np.cos(y)) - assert np.allclose(res[0][1], 1) - assert np.allclose(res[1], np.cos(0.5)) - assert np.allclose(res[2], np.cos(x) * np.cos(y)) - - # adjoint method only performs a single device execution per tape, but gets both result and gradient - assert dev.num_executions == 3 - spy.assert_not_called() - - g = jax.jacobian(cost)(a) - spy.assert_called() - expected_g = (-np.sin(x) * np.cos(y), -np.cos(x) * np.sin(y)) - assert qml.math.allclose(g[0][0], expected_g) - assert qml.math.allclose(g[0][1], np.zeros(2)) - assert qml.math.allclose(g[1], np.zeros(2)) - assert qml.math.allclose(g[2], expected_g) - - def test_no_grad_on_execution(self, mocker): - """Test that `grad_on_execution=False` uses the `device.execute_and_gradients`.""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "execute_and_gradients") - - def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn="device", - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - )[0] - a = jax.numpy.array([0.1, 0.2]) - cost(a) +def test_jit_execution(): + """Test that qml.execute can be directly jitted.""" + dev = qml.device("default.qubit") - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() + tape = qml.tape.QuantumScript( + [qml.RX(jax.numpy.array(0.1), 0)], [qml.expval(qml.s_prod(2.0, qml.PauliZ(0)))] + ) - jax.grad(cost)(a) - spy_gradients.assert_called() + out = jax.jit(qml.execute, static_argnames=("device", "gradient_fn"))( + (tape,), device=dev, gradient_fn=qml.gradients.param_shift + ) + expected = 2.0 * jax.numpy.cos(jax.numpy.array(0.1)) + assert qml.math.allclose(out[0], expected) +# pylint: disable=too-few-public-methods class TestCaching: - """Test for caching behaviour""" - - def test_cache_maxsize(self, mocker): - """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cachesize): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) + """Tests for caching behaviour""" - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn=param_shift, - cachesize=cachesize, - )[0] - - params = jax.numpy.array([0.1, 0.2]) - jax.grad(cost)(params, cachesize=2) - cache = spy.call_args.kwargs["cache"] - - assert cache.maxsize == 2 - assert cache.currsize == 2 - assert len(cache) == 2 + @pytest.mark.skip("caching is not implemented for jax") + @pytest.mark.parametrize("num_params", [2, 3]) + def test_caching_param_shift_hessian(self, num_params): + """Test that, when using parameter-shift transform, + caching reduces the number of evaluations to their optimum + when computing Hessians.""" + device = DefaultQubit() + params = jnp.arange(1, num_params + 1) / 10 - def test_custom_cache(self, mocker): - """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") + N = len(params) - def cost(a, cache): + def cost(x, cache): with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn=param_shift, - cache=cache, - )[0] - - custom_cache = {} - params = jax.numpy.array([0.1, 0.2]) - jax.grad(cost)(params, cache=custom_cache) - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache - - def test_custom_cache_multiple(self, mocker): - """Test the use of a custom cache object with multiple tapes""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) - - def cost(a, b, cache): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - res = execute( - [tape1, tape2], - dev, - gradient_fn=param_shift, - cache=cache, - ) - return res[0] - - custom_cache = {} - jax.grad(cost)(a, b, cache=custom_cache) - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache + qml.RX(x[0], wires=[0]) + qml.RY(x[1], wires=[1]) - def test_caching_param_shift(self, tol): - """Test that, when using parameter-shift transform, - caching produces the optimum number of evaluations.""" - dev = qml.device("default.qubit.legacy", wires=1) + for i in range(2, num_params): + qml.RZ(x[i], wires=[i % 2]) - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) + qml.CNOT(wires=[0, 1]) + qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn=param_shift, - cache=cache, + return qml.execute( + [tape], device, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 )[0] - # Without caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) - params = jax.numpy.array([0.1, 0.2]) - jax.grad(cost)(params, cache=None) - assert dev.num_executions == 5 - - # With caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - jac_fn = jax.grad(cost) - grad1 = jac_fn(params, cache=True) - assert dev.num_executions == 5 - - # Check that calling the cost function again - # continues to evaluate the device (that is, the cache - # is emptied between calls) - grad2 = jac_fn(params, cache=True) - assert dev.num_executions == 10 - assert np.allclose(grad1, grad2, atol=tol, rtol=0) - - # Check that calling the cost function again - # with different parameters produces a different Jacobian - grad2 = jac_fn(2 * params, cache=True) - assert dev.num_executions == 15 - assert not np.allclose(grad1, grad2, atol=tol, rtol=0) - - def test_caching_adjoint_backward(self): - """Test that caching produces the optimum number of adjoint evaluations - when no grad on execution.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1, 0.2, 0.3]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) + # No caching: number of executions is not ideal + with qml.Tracker(device) as tracker: + hess1 = jax.jacobian(jax.grad(cost))(params, cache=False) + + if num_params == 2: + # compare to theoretical result + x, y, *_ = params + expected = jnp.array( + [ + [2 * jnp.cos(2 * x) * jnp.sin(y) ** 2, jnp.sin(2 * x) * jnp.sin(2 * y)], + [jnp.sin(2 * x) * jnp.sin(2 * y), -2 * jnp.cos(x) ** 2 * jnp.cos(2 * y)], + ] + ) + assert np.allclose(expected, hess1) + + expected_runs = 1 # forward pass + + # Jacobian of an involutory observable: + # ------------------------------------ + # + # 2 * N execs: evaluate the analytic derivative of + # 1 execs: Get , the expectation value of the tape with unshifted parameters. + num_shifted_evals = 2 * N + runs_for_jacobian = num_shifted_evals + 1 + expected_runs += runs_for_jacobian + + # Each tape used to compute the Jacobian is then shifted again + expected_runs += runs_for_jacobian * num_shifted_evals + assert tracker.totals["executions"] == expected_runs + + # Use caching: number of executions is ideal + + with qml.Tracker(device) as tracker2: + hess2 = jax.jacobian(jax.grad(cost))(params, cache=True) + assert np.allclose(hess1, hess2) + + expected_runs_ideal = 1 # forward pass + expected_runs_ideal += 2 * N # Jacobian + expected_runs_ideal += N + 1 # Hessian diagonal + expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal + assert tracker2.totals["executions"] == expected_runs_ideal + assert expected_runs_ideal < expected_runs + + +# add tests for lightning 2 when possible +# set rng for device when possible +no_shots = Shots(None) +shots_2_10k = Shots((10000, 10000)) +dev_def = DefaultQubit() +dev_ps = ParamShiftDerivativesDevice(seed=54353453) +test_matrix = [ + ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), # 0 + ({"gradient_fn": param_shift}, no_shots, dev_def), # 1 + ({"gradient_fn": "backprop"}, no_shots, dev_def), # 2 + ({"gradient_fn": "adjoint"}, no_shots, dev_def), # 3 + ({"gradient_fn": "adjoint", "device_vjp": True}, no_shots, dev_def), # 4 + ({"gradient_fn": "device"}, shots_2_10k, dev_ps), # 5 +] - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn="device", - cache=cache, - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - )[0] - # Without caching, 2 evaluations are required. - # 1 for the forward pass, and one per output dimension - # on the backward pass. - jax.grad(cost)(params, cache=None) - assert dev.num_executions == 2 - - # With caching, also 2 evaluations are required. One - # for the forward pass, and one for the backward pass. - dev._num_executions = 0 - jac_fn = jax.grad(cost) - jac_fn(params, cache=True) - assert dev.num_executions == 2 - - -execute_kwargs_integration = [ - {"gradient_fn": param_shift}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - }, -] +def atol_for_shots(shots): + """Return higher tolerance if finite shots.""" + return 3e-2 if shots else 1e-6 -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) class TestJaxExecuteIntegration: """Test the jax interface execute function integrates well for both forward and backward execution""" - def test_execution(self, execute_kwargs): + def test_execution(self, execute_kwargs, shots, device): """Test execution""" - dev = qml.device("default.qubit.legacy", wires=1) def cost(a, b): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] + tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] + tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - tape2 = qml.tape.QuantumScript.from_queue(q2) + return execute([tape1, tape2], device, **execute_kwargs) - return execute([tape1, tape2], dev, **execute_kwargs) + a = jnp.array(0.1) + b = np.array(0.2) + with device.tracker: + res = cost(a, b) - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) - res = cost(a, b) + if execute_kwargs.get("gradient_fn", None) == "adjoint": + assert device.tracker.totals.get("execute_and_derivative_batches", 0) == 0 + else: + assert device.tracker.totals["batches"] == 1 + assert device.tracker.totals["executions"] == 2 # different wires so different hashes assert len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () - def test_scalar_jacobian(self, execute_kwargs, tol): + assert qml.math.allclose(res[0], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots)) + assert qml.math.allclose(res[1], jnp.cos(a) * jnp.cos(b), atol=atol_for_shots(shots)) + + def test_scalar_jacobian(self, execute_kwargs, shots, device): """Test scalar jacobian calculation""" - a = jax.numpy.array(0.1) - dev = qml.device("default.qubit.legacy", wires=2) + a = jnp.array(0.1) def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, **execute_kwargs)[0] + tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - res = jax.grad(cost)(a) - assert res.shape == () + res = jax.jacobian(cost)(a) + if not shots.has_partitioned_shots: + assert res.shape == () # pylint: disable=no-member # compare to standard tape jacobian - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) + tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) tape.trainable_params = [0] tapes, fn = param_shift(tape) - expected = fn(dev.batch_execute(tapes)) + expected = fn(device.execute(tapes)) assert expected.shape == () - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res, -jnp.sin(a), atol=atol_for_shots(shots)) - def test_reusing_quantum_tape(self, execute_kwargs, tol): - """Test re-using a quantum tape by passing new parameters""" - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) + def test_jacobian(self, execute_kwargs, shots, device): + """Test jacobian calculation""" - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - assert tape.trainable_params == [0, 1] + a = jnp.array(0.1) + b = jnp.array(0.2) def cost(a, b): - # An explicit call to _update() is required here to update the - # trainable parameters in between tape executions. - # This is different from how the autograd interface works. - # Unless the update is issued, the validation check related to the - # number of provided parameters fails in the tape: (len(params) != - # required_length) and the tape produces incorrect results. - tape._update() - new_tape = tape.bind_new_parameters([a, b], [0, 1]) - return execute([new_tape], dev, **execute_kwargs)[0] + ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return execute([tape], device, **execute_kwargs)[0] - jac_fn = jax.grad(cost) - jac = jac_fn(a, b) + res = cost(a, b) + expected = [jnp.cos(a), -jnp.cos(a) * jnp.sin(b)] + if shots.has_partitioned_shots: + assert np.allclose(res[0], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected, atol=2 * atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + g = jax.jacobian(cost, argnums=[0, 1])(a, b) + assert isinstance(g, tuple) and len(g) == 2 + + expected = ([-jnp.sin(a), jnp.sin(a) * jnp.sin(b)], [0, -jnp.cos(a) * jnp.cos(b)]) + + if shots.has_partitioned_shots: + for i in (0, 1): + assert np.allclose(g[i][0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[i][1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[i][0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[i][1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(g[0][0], expected[0][0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[1][0], expected[0][1], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[0][1], expected[1][0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(g[1][1], expected[1][1], atol=atol_for_shots(shots), rtol=0) + + def test_tape_no_parameters(self, execute_kwargs, shots, device): + """Test that a tape with no parameters is correctly + ignored during the gradient computation""" - a = jax.numpy.array(0.54) - b = jax.numpy.array(0.8) + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots + ) - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - res2 = cost(2 * a, b) - expected = [np.cos(2 * a)] - assert np.allclose(res2, expected, atol=tol, rtol=0) + tape2 = qml.tape.QuantumScript( + [qml.RY(jnp.array(0.5), wires=0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - jac_fn = jax.grad(lambda a, b: cost(2 * a, b)) - jac = jac_fn(a, b) - expected = -2 * np.sin(2 * a) - assert np.allclose(jac, expected, atol=tol, rtol=0) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - def test_grad_with_different_grad_on_execution(self, execute_kwargs): - """Test jax grad for adjoint diff method with different execution kwargs.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1, 0.2, 0.3]) - expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705]) + tape4 = qml.tape.QuantumScript( + [qml.RY(jnp.array(0.5), 0)], [qml.probs(wires=[0, 1])], shots=shots + ) + res = execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) + res = jax.tree_util.tree_leaves(res) + out = sum(jnp.hstack(res)) + if shots.has_partitioned_shots: + out = out / shots.num_copies + return out - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) + params = jnp.array([0.1, 0.2]) - tape = qml.tape.QuantumScript.from_queue(q) - res = qml.execute([tape], dev, cache=cache, **execute_kwargs)[0] - return res + x, y = params - results = jax.grad(cost)(params, cache=None) - for r, e in zip(results, expected_results): - assert jax.numpy.allclose(r, e, atol=1e-7) + res = cost(params) + expected = 2 + jnp.cos(0.5) + jnp.cos(x) * jnp.cos(y) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_classical_processing_single_tape(self, execute_kwargs): - """Test classical processing within the quantum tape for a single tape""" - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) - c = jax.numpy.array(0.3) + grad = jax.grad(cost)(params) + expected = [-jnp.cos(y) * jnp.sin(x), -jnp.cos(x) * jnp.sin(y)] + assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - def cost(a, b, c, device): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a * c, wires=0) - qml.RZ(b, wires=0) - qml.RX(c + c**2 + jax.numpy.sin(a), wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) + # pylint: disable=too-many-statements + def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): + """Test that tapes wit different can be executed and differentiated.""" - return execute([tape], device, **execute_kwargs)[0] + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + shots=shots, + ) - dev = qml.device("default.qubit.legacy", wires=2) - res = jax.grad(cost, argnums=(0, 1, 2))(a, b, c, device=dev) - assert len(res) == 3 + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), 0)], [qml.expval(qml.PauliZ(0))], shots=shots + ) - def test_classical_processing_multiple_tapes(self, execute_kwargs): - """Test classical processing within the quantum tape for multiple - tapes""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.3, 0.2]) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + res = qml.execute([tape1, tape2, tape3], device, **execute_kwargs) + leaves = jax.tree_util.tree_leaves(res) + return jnp.hstack(leaves) + + params = jnp.array([0.1, 0.2]) + x, y = params + + res = cost(params) + assert isinstance(res, jax.Array) + assert res.shape == (4 * shots.num_copies,) if shots.has_partitioned_shots else (4,) + + if shots.has_partitioned_shots: + assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[2], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[3], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[4], jnp.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[5], jnp.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[6], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[7], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + else: + assert np.allclose(res[0], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[2], jnp.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[3], jnp.cos(x) * jnp.cos(y), atol=atol_for_shots(shots)) + + jac = jax.jacobian(cost)(params) + assert isinstance(jac, jnp.ndarray) + assert ( + jac.shape == (8, 2) if shots.has_partitioned_shots else (4, 2) + ) # pylint: disable=no-member + + if shots.has_partitioned_shots: + assert np.allclose(jac[1], 0, atol=atol_for_shots(shots)) + assert np.allclose(jac[3:5], 0, atol=atol_for_shots(shots)) + else: + assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) + + d1 = -jnp.sin(x) * jnp.cos(y) + if shots.has_partitioned_shots: + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[2, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[6, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[7, 0], d1, atol=atol_for_shots(shots)) + else: + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) + + d2 = -jnp.cos(x) * jnp.sin(y) + if shots.has_partitioned_shots: + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[2, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[6, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[7, 1], d2, atol=atol_for_shots(shots)) + else: + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) + + def test_reusing_quantum_tape(self, execute_kwargs, shots, device): + """Test re-using a quantum tape by passing new parameters""" + if execute_kwargs["gradient_fn"] == param_shift: + pytest.skip("Basic QNode execution wipes out trainable params with param-shift") - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.RY(x[0], wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) + a = jnp.array(0.1) + b = jnp.array(0.2) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.Hadamard(0) - qml.CRX(2 * x[0] * x[1], wires=[0, 1]) - qml.RX(2 * x[1], wires=[1]) - qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript( + [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], + shots=shots, + ) + assert tape.trainable_params == [0, 1] - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) - return result[0] + result[1] - 7 * result[1] + def cost(a, b): + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return jnp.hstack(execute([new_tape], device, **execute_kwargs)[0]) - res = jax.grad(cost_fn)(params) - assert res.shape == (2,) + jac_fn = jax.jacobian(cost, argnums=[0, 1]) + jac = jac_fn(a, b) - def test_multiple_tapes_output(self, execute_kwargs): - """Test the output types for the execution of multiple quantum tapes""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.3, 0.2]) + a = jnp.array(0.54) + b = jnp.array(0.8) - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.RY(x[0], wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = [jnp.cos(2 * a), -jnp.cos(2 * a) * jnp.sin(b)] + if shots.has_partitioned_shots: + assert np.allclose(res2[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res2[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) + + jac_fn = jax.jacobian(lambda a, b: cost(2 * a, b), argnums=[0, 1]) + jac = jac_fn(a, b) + expected = ( + [-2 * jnp.sin(2 * a), 2 * jnp.sin(2 * a) * jnp.sin(b)], + [0, -jnp.cos(2 * a) * jnp.cos(b)], + ) + assert isinstance(jac, tuple) and len(jac) == 2 + if shots.has_partitioned_shots: + for offset in (0, 2): + assert np.allclose(jac[0][0 + offset], expected[0][0], atol=atol_for_shots(shots)) + assert np.allclose(jac[0][1 + offset], expected[0][1], atol=atol_for_shots(shots)) + assert np.allclose(jac[1][0 + offset], expected[1][0], atol=atol_for_shots(shots)) + assert np.allclose(jac[1][1 + offset], expected[1][1], atol=atol_for_shots(shots)) + else: + for _j, _e in zip(jac, expected): + assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) + + def test_classical_processing(self, execute_kwargs, shots, device): + """Test classical processing within the quantum tape""" + a = jnp.array(0.1) + b = jnp.array(0.2) + c = jnp.array(0.3) + + def cost(a, b, c): + ops = [ + qml.RY(a * c, wires=0), + qml.RZ(b, wires=0), + qml.RX(c + c**2 + jnp.sin(a), wires=0), + ] - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.Hadamard(0) - qml.CRX(2 * x[0] * x[1], wires=[0, 1]) - qml.RX(2 * x[1], wires=[1]) - qml.expval(qml.PauliZ(0)) + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - tape2 = qml.tape.QuantumScript.from_queue(q2) + res = jax.jacobian(cost, argnums=[0, 2])(a, b, c) - return execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + # Only two arguments are trainable + assert isinstance(res, tuple) and len(res) == 2 + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () - res = cost_fn(params) - assert isinstance(res, TensorLike) - assert all(isinstance(r, jax.numpy.ndarray) for r in res) - assert all(r.shape == () for r in res) + # I tried getting analytic results for this circuit but I kept being wrong and am giving up - def test_matrix_parameter(self, execute_kwargs, tol): + def test_matrix_parameter(self, execute_kwargs, device, shots): """Test that the jax interface works correctly with a matrix parameter""" - a = jax.numpy.array(0.1) - U = jax.numpy.array([[0, 1], [1, 0]]) - - def cost(a, U, device): - with qml.queuing.AnnotatedQueue() as q: - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + U = jnp.array([[0, 1], [1, 0]]) + a = jnp.array(0.1) - tape = qml.tape.QuantumScript.from_queue(q) - tape.trainable_params = [0] + def cost(a, U): + ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) - res = cost(a, U, device=dev) - assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) + res = cost(a, U) + assert np.allclose(res, -jnp.cos(a), atol=atol_for_shots(shots), rtol=0) - jac_fn = jax.grad(cost, argnums=0) - res = jac_fn(a, U, device=dev) - assert np.allclose(res, np.sin(a), atol=tol, rtol=0) + jac_fn = jax.jacobian(cost) + jac = jac_fn(a, U) + if not shots.has_partitioned_shots: + assert isinstance(jac, jnp.ndarray) + assert np.allclose(jac, jnp.sin(a), atol=atol_for_shots(shots), rtol=0) - def test_differentiable_expand(self, execute_kwargs, tol): + def test_differentiable_expand(self, execute_kwargs, device, shots): """Test that operation and nested tapes expansion is differentiable""" class U3(qml.U3): + """Dummy operator.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -619,258 +469,349 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - def cost_fn(a, p, device): - with qml.queuing.AnnotatedQueue() as q_tape: - qml.RX(a, wires=0) - U3(*p, wires=0) - qml.expval(qml.PauliX(0)) - - tape = qml.tape.QuantumScript.from_queue(q_tape) - tape = tape.expand(stop_at=lambda obj: device.supports_operation(obj.name)) - return execute([tape], device, **execute_kwargs)[0] + def cost_fn(a, p): + tape = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], + [qml.expval(qml.PauliX(0))], + shots=shots, + ) + gradient_fn = execute_kwargs["gradient_fn"] + if gradient_fn is None: + _gradient_method = None + elif isinstance(gradient_fn, str): + _gradient_method = gradient_fn + else: + _gradient_method = "gradient-transform" + conf = qml.devices.ExecutionConfig( + interface="autograd", + gradient_method=_gradient_method, + grad_on_execution=execute_kwargs.get("grad_on_execution", None), + ) + program, _ = device.preprocess(execution_config=conf) + return execute([tape], device, **execute_kwargs, transform_program=program)[0] - a = jax.numpy.array(0.1) - p = jax.numpy.array([0.1, 0.2, 0.3]) + a = jnp.array(0.1) + p = jnp.array([0.1, 0.2, 0.3]) - dev = qml.device("default.qubit.legacy", wires=1) - res = cost_fn(a, p, device=dev) - expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( - np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) + res = cost_fn(a, p) + expected = jnp.cos(a) * jnp.cos(p[1]) * jnp.sin(p[0]) + jnp.sin(a) * ( + jnp.cos(p[2]) * jnp.sin(p[1]) + jnp.cos(p[0]) * jnp.cos(p[1]) * jnp.sin(p[2]) ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - jac_fn = jax.grad(cost_fn, argnums=1) - res = jac_fn(a, p, device=dev) - expected = jax.numpy.array( + jac_fn = jax.jacobian(cost_fn, argnums=[1]) + res = jac_fn(a, p) + expected = jnp.array( [ - np.cos(p[1]) * (np.cos(a) * np.cos(p[0]) - np.sin(a) * np.sin(p[0]) * np.sin(p[2])), - np.cos(p[1]) * np.cos(p[2]) * np.sin(a) - - np.sin(p[1]) - * (np.cos(a) * np.sin(p[0]) + np.cos(p[0]) * np.sin(a) * np.sin(p[2])), - np.sin(a) - * (np.cos(p[0]) * np.cos(p[1]) * np.cos(p[2]) - np.sin(p[1]) * np.sin(p[2])), + jnp.cos(p[1]) + * (jnp.cos(a) * jnp.cos(p[0]) - jnp.sin(a) * jnp.sin(p[0]) * jnp.sin(p[2])), + jnp.cos(p[1]) * jnp.cos(p[2]) * jnp.sin(a) + - jnp.sin(p[1]) + * (jnp.cos(a) * jnp.sin(p[0]) + jnp.cos(p[0]) * jnp.sin(a) * jnp.sin(p[2])), + jnp.sin(a) + * (jnp.cos(p[0]) * jnp.cos(p[1]) * jnp.cos(p[2]) - jnp.sin(p[1]) * jnp.sin(p[2])), ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_independent_expval(self, execute_kwargs): - """Tests computing an expectation value that is independent of trainable - parameters.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1, 0.2, 0.3]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, cache=cache, **execute_kwargs) - return res[0] + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = jax.grad(cost)(params, cache=None) - assert res.shape == (3,) + def test_probability_differentiation(self, execute_kwargs, device, shots): + """Tests correct output shape and evaluation for a tape + with prob outputs""" + def cost(x, y): + ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.probs(wires=0), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return jnp.hstack(execute([tape], device, **execute_kwargs)[0]) -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) -class TestVectorValued: - """Test vector-valued jacobian returns for the JAX Python interface.""" - - def test_multiple_expvals(self, execute_kwargs): - """Tests computing multiple expectation values in a tape.""" - - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1, 0.2, 0.3]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = qml.execute([tape], dev, cache=cache, **execute_kwargs) - return res[0] - - res = jax.jacobian(cost)(params, cache=None) + x = jnp.array(0.543) + y = jnp.array(-0.654) - assert isinstance(res, tuple) - assert len(res) == 2 - - assert res[0].shape == (3,) - assert isinstance(res[0], jax.numpy.ndarray) - - assert res[1].shape == (3,) - assert isinstance(res[1], jax.numpy.ndarray) - - def test_multiple_expvals_single_par(self, execute_kwargs): - """Tests computing multiple expectation values in a tape with a single - trainable parameter.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.1]) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = qml.execute([tape], dev, cache=cache, **execute_kwargs) - return res[0] - - res = jax.jacobian(cost)(params, cache=None) - - assert isinstance(res, tuple) - - assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == (1,) - - assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (1,) + res = cost(x, y) + expected = jnp.array( + [ + [ + jnp.cos(x / 2) ** 2, + jnp.sin(x / 2) ** 2, + (1 + jnp.cos(x) * jnp.cos(y)) / 2, + (1 - jnp.cos(x) * jnp.cos(y)) / 2, + ], + ] + ) + if shots.has_partitioned_shots: + assert np.allclose(res[:, 0:2].flatten(), expected, atol=atol_for_shots(shots)) + assert np.allclose(res[:, 2:].flatten(), expected, atol=atol_for_shots(shots)) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + jac_fn = jax.jacobian(cost, argnums=[0, 1]) + res = jac_fn(x, y) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (2, 4) if shots.has_partitioned_shots else (4,) + assert res[1].shape == (2, 4) if shots.has_partitioned_shots else (4,) + + expected = ( + jnp.array( + [ + [ + -jnp.sin(x) / 2, + jnp.sin(x) / 2, + -jnp.sin(x) * jnp.cos(y) / 2, + jnp.sin(x) * jnp.cos(y) / 2, + ], + ] + ), + jnp.array( + [ + [0, 0, -jnp.cos(x) * jnp.sin(y) / 2, jnp.cos(x) * jnp.sin(y) / 2], + ] + ), + ) - def test_multi_tape_fwd(self, execute_kwargs): - """Test the forward evaluation of a cost function that uses the output - of multiple tapes that be vector-valued.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = jax.numpy.array([0.3, 0.2]) + if shots.has_partitioned_shots: + assert np.allclose(res[0][:, 0:2].flatten(), expected[0], atol=atol_for_shots(shots)) + assert np.allclose(res[0][:, 2:].flatten(), expected[0], atol=atol_for_shots(shots)) + assert np.allclose(res[1][:, :2].flatten(), expected[1], atol=atol_for_shots(shots)) + assert np.allclose(res[1][:, 2:].flatten(), expected[1], atol=atol_for_shots(shots)) + else: + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) + + def test_ragged_differentiation(self, execute_kwargs, device, shots): + """Tests correct output shape and evaluation for a tape + with prob and expval outputs""" + + def cost(x, y): + ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + res = qml.execute([tape], device, **execute_kwargs)[0] + return jnp.hstack(jax.tree_util.tree_leaves(res)) + + x = jnp.array(0.543) + y = jnp.array(-0.654) + + res = cost(x, y) + expected = jnp.array( + [jnp.cos(x), (1 + jnp.cos(x) * jnp.cos(y)) / 2, (1 - jnp.cos(x) * jnp.cos(y)) / 2] + ) + if shots.has_partitioned_shots: + assert np.allclose(res[:3], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[3:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + jac_fn = jax.jacobian(cost, argnums=[0, 1]) + res = jac_fn(x, y) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,) + assert res[1].shape == (3 * shots.num_copies,) if shots.has_partitioned_shots else (3,) + + expected = ( + jnp.array([-jnp.sin(x), -jnp.sin(x) * jnp.cos(y) / 2, jnp.sin(x) * jnp.cos(y) / 2]), + jnp.array([0, -jnp.cos(x) * jnp.sin(y) / 2, jnp.cos(x) * jnp.sin(y) / 2]), + ) + if shots.has_partitioned_shots: + assert np.allclose(res[0][:3], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[0][3:], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1][:3], expected[1], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1][3:], expected[1], atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) + + +class TestHigherOrderDerivatives: + """Test that the jax execute function can be differentiated""" + + @pytest.mark.parametrize( + "params", + [ + jnp.array([0.543, -0.654]), + jnp.array([0, -0.654]), + jnp.array([-2.0, 0]), + ], + ) + def test_parameter_shift_hessian(self, params, tol): + """Tests that the output of the parameter-shift transform + can be differentiated using jax, yielding second derivatives.""" + dev = DefaultQubit() def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.expval(qml.PauliY(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[1], wires=[0]) - qml.RX(x[1], wires=[0]) - qml.RX(-x[1], wires=[0]) - qml.expval(qml.PauliY(0)) - qml.expval(qml.PauliY(1)) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = qml.execute(tapes=[tape1, tape2], device=dev, **execute_kwargs) + ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) + + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) return result[0] + result[1][0] - expected = -jax.numpy.sin(params[0]) + -jax.numpy.sin(params[1]) res = cost_fn(params) - assert jax.numpy.allclose(expected, res) - - def test_multi_tape_jacobian(self, execute_kwargs): - """Test the jacobian computation with multiple tapes.""" + x, y = params + expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) - def cost(x, y, device, interface, ek): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + res = jax.grad(cost_fn)(params) + expected = jnp.array( + [-jnp.cos(x) * jnp.cos(2 * y) * jnp.sin(x), -jnp.cos(x) ** 2 * jnp.sin(2 * y)] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + res = jax.jacobian(jax.grad(cost_fn))(params) + expected = jnp.array( + [ + [-jnp.cos(2 * x) * jnp.cos(2 * y), jnp.sin(2 * x) * jnp.sin(2 * y)], + [jnp.sin(2 * x) * jnp.sin(2 * y), -2 * jnp.cos(x) ** 2 * jnp.cos(2 * y)], + ] + ) + assert np.allclose(res, expected, atol=tol, rtol=0) - tape2 = qml.tape.QuantumScript.from_queue(q2) - return qml.execute([tape1, tape2], device, **ek, interface=interface) + def test_max_diff(self, tol): + """Test that setting the max_diff parameter blocks higher-order + derivatives""" + dev = DefaultQubit() + params = jnp.array([0.543, -0.654]) - dev = qml.device("default.qubit.legacy", wires=2) - x = jax.numpy.array(0.543) - y = jax.numpy.array(-0.654) + def cost_fn(x): + ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - x_ = np.array(0.543) - y_ = np.array(-0.654) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - exec_jax = cost(x, y, dev, interface="jax-python", ek=execute_kwargs) - exec_autograd = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) + return result[0] + result[1][0] - assert np.allclose(exec_jax, exec_autograd) + res = cost_fn(params) + x, y = params + expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) - res = jax.jacobian(cost, argnums=(0, 1))( - x, y, dev, interface="jax-python", ek=execute_kwargs + res = jax.grad(cost_fn)(params) + expected = jnp.array( + [-jnp.cos(x) * jnp.cos(2 * y) * jnp.sin(x), -jnp.cos(x) ** 2 * jnp.sin(2 * y)] ) + assert np.allclose(res, expected, atol=tol, rtol=0) - import autograd.numpy as anp - - def cost_stack(x, y, device, interface, ek): - return anp.hstack(cost(x, y, device, interface, ek)) + res = jax.jacobian(jax.grad(cost_fn))(params) + expected = jnp.zeros([2, 2]) + assert np.allclose(res, expected, atol=tol, rtol=0) - exp = qml.jacobian(cost_stack, argnum=(0, 1))( - x_, y_, dev, interface="autograd", ek=execute_kwargs - ) - res_0 = jax.numpy.array([res[0][0][0], res[0][1][0], res[1][0][0], res[1][1][0]]) - res_1 = jax.numpy.array([res[0][0][1], res[0][1][1], res[1][0][1], res[1][1][1]]) - assert np.allclose(res_0, exp[0]) - assert np.allclose(res_1, exp[1]) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +@pytest.mark.usefixtures("use_legacy_and_new_opmath") +class TestHamiltonianWorkflows: + """Test that tapes ending with expectations + of Hamiltonians provide correct results and gradients""" - def test_multi_tape_jacobian_probs_expvals(self, execute_kwargs): - """Test the jacobian computation with multiple tapes with probability - and expectation value computations.""" + @pytest.fixture + def cost_fn(self, execute_kwargs, shots, device): + """Cost function for gradient tests""" - adjoint = execute_kwargs.get("gradient_kwargs", {}).get("method", "") == "adjoint_jacobian" - if adjoint: - pytest.skip("The adjoint diff method doesn't support probabilities.") + def _cost_fn(weights, coeffs1, coeffs2): + obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] + H1 = qml.Hamiltonian(coeffs1, obs1) + if qml.operation.active_new_opmath(): + H1 = qml.pauli.pauli_sentence(H1).operation() - def cost(x, y, device, interface, ek): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + obs2 = [qml.PauliZ(0)] + H2 = qml.Hamiltonian(coeffs2, obs2) + if qml.operation.active_new_opmath(): + H2 = qml.pauli.pauli_sentence(H2).operation() - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) + with qml.queuing.AnnotatedQueue() as q: + qml.RX(weights[0], wires=0) + qml.RY(weights[1], wires=1) qml.CNOT(wires=[0, 1]) - qml.probs(wires=[0]) - qml.probs(wires=[1]) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - return qml.execute([tape1, tape2], device, **ek, interface=interface) - - dev = qml.device("default.qubit.legacy", wires=2) - x = jax.numpy.array(0.543) - y = jax.numpy.array(-0.654) - - x_ = np.array(0.543) - y_ = np.array(-0.654) - - exec_jax = cost(x, y, dev, interface="jax-python", ek=execute_kwargs) - exec_autograd = cost(x_, y_, dev, interface="autograd", ek=execute_kwargs) - - assert all( - np.allclose(exec_jax[i][j], exec_autograd[i][j]) for i in range(2) for j in range(2) - ) - - res = jax.jacobian(cost, argnums=(0, 1))( - x, y, dev, interface="jax-python", ek=execute_kwargs + qml.expval(H1) + qml.expval(H2) + + tape = qml.tape.QuantumScript.from_queue(q, shots=shots) + res = execute([tape], device, **execute_kwargs)[0] + if shots.has_partitioned_shots: + return jnp.hstack(res[0] + res[1]) + return jnp.hstack(res) + + return _cost_fn + + @staticmethod + def cost_fn_expected(weights, coeffs1, coeffs2): + """Analytic value of cost_fn above""" + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return [-c * jnp.sin(x) * jnp.sin(y) + jnp.cos(x) * (a + b * jnp.sin(y)), d * jnp.cos(x)] + + @staticmethod + def cost_fn_jacobian(weights, coeffs1, coeffs2): + """Analytic jacobian of cost_fn above""" + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return jnp.array( + [ + [ + -c * jnp.cos(x) * jnp.sin(y) - jnp.sin(x) * (a + b * jnp.sin(y)), + b * jnp.cos(x) * jnp.cos(y) - c * jnp.cos(y) * jnp.sin(x), + jnp.cos(x), + jnp.cos(x) * jnp.sin(y), + -(jnp.sin(x) * jnp.sin(y)), + 0, + ], + [-d * jnp.sin(x), 0, 0, 0, 0, jnp.cos(x)], + ] ) - assert isinstance(res, TensorLike) - assert len(res) == 2 - - for r, exp_shape in zip(res, [(), (2,)]): - assert isinstance(r, tuple) - assert len(r) == 2 - assert len(r[0]) == 2 - assert isinstance(r[0][0], jax.numpy.ndarray) - assert r[0][0].shape == exp_shape - assert isinstance(r[0][1], jax.numpy.ndarray) - assert r[0][1].shape == exp_shape - assert len(r[1]) == 2 - assert isinstance(r[1][0], jax.numpy.ndarray) - assert r[1][0].shape == exp_shape - assert isinstance(r[1][1], jax.numpy.ndarray) - assert r[1][1].shape == exp_shape + def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with no trainable parameters.""" + + if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): + pytest.skip("adjoint differentiation does not suppport hamiltonians.") + + coeffs1 = jnp.array([0.1, 0.2, 0.3]) + coeffs2 = jnp.array([0.7]) + weights = jnp.array([0.4, 0.5]) + + res = cost_fn(weights, coeffs1, coeffs2) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + res = jax.jacobian(cost_fn)(weights, coeffs1, coeffs2) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] + if shots.has_partitioned_shots: + assert np.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with trainable parameters.""" + if execute_kwargs["gradient_fn"] == "adjoint": + pytest.skip("trainable hamiltonians not supported with adjoint") + if qml.operation.active_new_opmath(): + pytest.skip("parameter shift derivatives do not yet support sums.") + + coeffs1 = jnp.array([0.1, 0.2, 0.3]) + coeffs2 = jnp.array([0.7]) + weights = jnp.array([0.4, 0.5]) + + res = cost_fn(weights, coeffs1, coeffs2) + expected = self.cost_fn_expected(weights, coeffs1, coeffs2) + if shots.has_partitioned_shots: + assert np.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + res = jnp.hstack(jax.jacobian(cost_fn, argnums=[0, 1, 2])(weights, coeffs1, coeffs2)) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + if shots.has_partitioned_shots: + pytest.xfail( + "multiple hamiltonians with shot vectors does not seem to be differentiable." + ) + else: + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py index c3fd7df5d1a..4fe3d4946d2 100644 --- a/tests/interfaces/test_jax_jit.py +++ b/tests/interfaces/test_jax_jit.py @@ -14,7 +14,7 @@ """Unit tests for the JAX-JIT interface""" import numpy as np -# pylint: disable=protected-access,too-few-public-methods +# pylint: disable=protected-access,too-few-public-methods,unnecessary-lambda import pytest import pennylane as qml @@ -37,7 +37,7 @@ def test_jacobian_options(self, mocker): a = jax.numpy.array([0.1, 0.2]) - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, device): with qml.queuing.AnnotatedQueue() as q: @@ -64,7 +64,7 @@ def test_incorrect_gradients_on_execution(self): is used with grad_on_execution=True""" a = jax.numpy.array([0.1, 0.2]) - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, device): with qml.queuing.AnnotatedQueue() as q: @@ -90,7 +90,7 @@ def test_unknown_interface(self): """Test that an error is raised if the interface is unknown""" a = jax.numpy.array([0.1, 0.2]) - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, device): with qml.queuing.AnnotatedQueue() as q: @@ -112,8 +112,8 @@ def cost(a, device): def test_grad_on_execution(self, mocker): """Test that grad on execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(dev, "execute_and_gradients") + dev = qml.device("default.qubit", wires=1) + spy = mocker.spy(dev, "execute_and_compute_derivatives") def cost(a): with qml.queuing.AnnotatedQueue() as q: @@ -134,22 +134,22 @@ def cost(a): )[0] a = jax.numpy.array([0.1, 0.2]) - jax.jit(cost)(a) + + with qml.Tracker(dev) as tracker: + jax.jit(cost)(a) # adjoint method only performs a single device execution # gradients are not requested when we only want the results - assert dev.num_executions == 1 spy.assert_not_called() + assert tracker.totals["executions"] == 1 # when the jacobian is requested, we always calculate it at the same time as the results jax.grad(jax.jit(cost))(a) spy.assert_called() - def test_no_gradients_on_execution(self, mocker): - """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "execute_and_gradients") + def test_no_gradients_on_execution(self): + """Test that no grad on execution uses the `device.execute` and `device.compute_derivatives` pathway""" + dev = qml.device("default.qubit", wires=1) def cost(a): with qml.queuing.AnnotatedQueue() as q: @@ -168,14 +168,16 @@ def cost(a): )[0] a = jax.numpy.array([0.1, 0.2]) - jax.jit(cost)(a) - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() + with qml.Tracker(dev) as tracker: + jax.jit(cost)(a) - jax.grad(jax.jit(cost))(a) - spy_gradients.assert_called() + assert tracker.totals["executions"] == 1 + assert "derivatives" not in tracker.totals + + with qml.Tracker(dev) as tracker: + jax.grad(jax.jit(cost))(a) + assert tracker.totals["derivatives"] == 1 class TestCaching: @@ -183,7 +185,7 @@ class TestCaching: def test_cache_maxsize(self, mocker): """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cachesize): @@ -211,7 +213,7 @@ def cost(a, cachesize): def test_custom_cache(self, mocker): """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cache): @@ -238,7 +240,7 @@ def cost(a, cache): def test_custom_cache_multiple(self, mocker): """Test the use of a custom cache object with multiple tapes""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") a = jax.numpy.array(0.1) @@ -276,7 +278,7 @@ def cost(a, b, cache): def test_caching_param_shift(self, tol): """Test that, when using parameter-shift transform, caching produces the optimum number of evaluations.""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, cache): with qml.queuing.AnnotatedQueue() as q: @@ -296,33 +298,37 @@ def cost(a, cache): # Without caching, 5 evaluations are required to compute # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) params = jax.numpy.array([0.1, 0.2]) - jax.grad(cost)(params, cache=None) - assert dev.num_executions == 5 + with qml.Tracker(dev) as tracker: + jax.grad(cost)(params, cache=None) + + assert tracker.totals["executions"] == 5 # With caching, 5 evaluations are required to compute # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - jac_fn = jax.grad(cost) - grad1 = jac_fn(params, cache=True) - assert dev.num_executions == 5 + with qml.Tracker(dev) as tracker: + jac_fn = jax.grad(cost) + grad1 = jac_fn(params, cache=True) + assert tracker.totals["executions"] == 5 # Check that calling the cost function again # continues to evaluate the device (that is, the cache # is emptied between calls) - grad2 = jac_fn(params, cache=True) - assert dev.num_executions == 10 + with qml.Tracker(dev) as tracker: + grad2 = jac_fn(params, cache=True) + assert tracker.totals["executions"] == 5 assert np.allclose(grad1, grad2, atol=tol, rtol=0) # Check that calling the cost function again # with different parameters produces a different Jacobian - grad2 = jac_fn(2 * params, cache=True) - assert dev.num_executions == 15 + with qml.Tracker(dev) as tracker: + grad2 = jac_fn(2 * params, cache=True) + assert tracker.totals["executions"] == 5 assert not np.allclose(grad1, grad2, atol=tol, rtol=0) def test_caching_adjoint_backward(self): """Test that caching produces the optimum number of adjoint evaluations when mode=backward""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -346,15 +352,18 @@ def cost(a, cache): # Without caching, 2 evaluations are required. # 1 for the forward pass, and one per output dimension # on the backward pass. - jax.grad(cost)(params, cache=None) - assert dev.num_executions == 2 + with qml.Tracker(dev) as tracker: + jax.grad(cost)(params, cache=None) + assert tracker.totals["executions"] == 1 + assert tracker.totals["derivatives"] == 1 # With caching, also 2 evaluations are required. One # for the forward pass, and one for the backward pass. - dev._num_executions = 0 - jac_fn = jax.grad(cost) - jac_fn(params, cache=True) - assert dev.num_executions == 2 + with qml.Tracker(dev) as tracker: + jac_fn = jax.grad(cost) + jac_fn(params, cache=True) + assert tracker.totals["executions"] == 1 + assert tracker.totals["derivatives"] == 1 execute_kwargs_integration = [ @@ -379,7 +388,7 @@ class TestJaxExecuteIntegration: def test_execution(self, execute_kwargs): """Test execution""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) def cost(a, b): with qml.queuing.AnnotatedQueue() as q1: @@ -409,7 +418,7 @@ def cost(a, b): def test_scalar_jacobian(self, execute_kwargs, tol): """Test scalar jacobian calculation""" a = jax.numpy.array(0.1) - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) def cost(a): with qml.queuing.AnnotatedQueue() as q: @@ -432,7 +441,7 @@ def cost(a): tape.trainable_params = [0] tapes, fn = param_shift(tape) - expected = fn(dev.batch_execute(tapes)) + expected = fn(dev.execute(tapes)) assert expected.shape == () assert np.allclose(res, expected, atol=tol, rtol=0) @@ -442,7 +451,7 @@ def test_reusing_quantum_tape(self, execute_kwargs, tol): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) with qml.queuing.AnnotatedQueue() as q: qml.RY(a, wires=0) @@ -484,7 +493,7 @@ def cost(a, b): def test_grad_with_backward_mode(self, execute_kwargs): """Test jax grad for adjoint diff method in backward mode""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) expected_results = jax.numpy.array([-0.3875172, -0.18884787, -0.38355705]) @@ -523,14 +532,14 @@ def cost(a, b, c, device): return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) res = jax.jit(jax.grad(cost, argnums=(0, 1, 2)), static_argnums=3)(a, b, c, device=dev) assert len(res) == 3 def test_classical_processing_multiple_tapes(self, execute_kwargs): """Test classical processing within the quantum tape for multiple tapes""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.3, 0.2]) def cost_fn(x): @@ -558,7 +567,7 @@ def cost_fn(x): def test_multiple_tapes_output(self, execute_kwargs): """Test the output types for the execution of multiple quantum tapes""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.3, 0.2]) def cost_fn(x): @@ -602,7 +611,7 @@ def cost(a, U, device): tape.trainable_params = [0] return execute([tape], device, **execute_kwargs)[0] - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) res = jax.jit(cost, static_argnums=2)(a, U, device=dev) assert np.allclose(res, -np.cos(a), atol=tol, rtol=0) @@ -627,13 +636,15 @@ def cost_fn(a, p, device): qscript = qml.tape.QuantumScript( [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] ) - qscript = qscript.expand(stop_at=lambda obj: device.supports_operation(obj.name)) + qscript = qscript.expand( + stop_at=lambda obj: qml.devices.default_qubit.stopping_condition(obj) + ) return execute([qscript], device, **execute_kwargs)[0] a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.qubit", wires=1) res = jax.jit(cost_fn, static_argnums=2)(a, p, device=dev) expected = np.cos(a) * np.cos(p[1]) * np.sin(p[0]) + np.sin(a) * ( np.cos(p[2]) * np.sin(p[1]) + np.cos(p[0]) * np.cos(p[1]) * np.sin(p[2]) @@ -657,7 +668,7 @@ def cost_fn(a, p, device): def test_independent_expval(self, execute_kwargs): """Tests computing an expectation value that is independent of trainable parameters.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -694,7 +705,7 @@ def test_shapes(self, execute_kwargs, ret_type, shape, expected_type): if adjoint: pytest.skip("The adjoint diff method doesn't support probabilities.") - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -722,7 +733,7 @@ def cost(a, cache): def test_independent_expval(self, execute_kwargs): """Tests computing an expectation value that is independent of trainable parameters.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -754,7 +765,7 @@ def cost(a, cache): def test_vector_valued_qnode(self, execute_kwargs, ret, out_dim, expected_type): """Tests the shape of vector-valued QNode results.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) grad_meth = ( execute_kwargs["gradient_kwargs"]["method"] @@ -793,7 +804,7 @@ def cost(a, cache): def test_qnode_sample(self, execute_kwargs): """Tests computing multiple expectation values in a tape.""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) + dev = qml.device("default.qubit", wires=2, shots=10) params = jax.numpy.array([0.1, 0.2, 0.3]) grad_meth = ( @@ -817,11 +828,12 @@ def cost(a, cache): return res res = jax.jit(cost, static_argnums=1)(params, cache=None) - assert res.shape == (dev.shots,) + + assert res.shape == (dev.shots.total_shots,) def test_multiple_expvals_grad(self, execute_kwargs): """Tests computing multiple expectation values in a tape.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) params = jax.numpy.array([0.1, 0.2, 0.3]) def cost(a, cache): @@ -868,7 +880,7 @@ def cost(x, y, device, interface, ek): return qml.execute([tape1, tape2], device, **ek, interface=interface)[0] - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) diff --git a/tests/interfaces/test_jax_jit_qnode.py b/tests/interfaces/test_jax_jit_qnode.py index f109ddc0a4c..1c0b8ad1f59 100644 --- a/tests/interfaces/test_jax_jit_qnode.py +++ b/tests/interfaces/test_jax_jit_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 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. @@ -12,21 +12,35 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the JAX-JIT interface with a QNode""" +import copy + # pylint: disable=too-many-arguments,too-few-public-methods +from functools import partial + import pytest +from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml from pennylane import numpy as np from pennylane import qnode +from pennylane.devices import DefaultQubit +# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], + [DefaultQubit(), "backprop", True, False], + [DefaultQubit(), "finite-diff", False, False], + [DefaultQubit(), "parameter-shift", False, False], + [DefaultQubit(), "adjoint", True, False], + [DefaultQubit(), "adjoint", True, True], + [ParamShiftDerivativesDevice(), "device", False, True], + [DefaultQubit(), "adjoint", False, False], + [DefaultQubit(), "spsa", False, False], + [DefaultQubit(), "hadamard", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", False, True], + [qml.device("lightning.qubit", wires=5), "adjoint", True, False], + [qml.device("lightning.qubit", wires=5), "adjoint", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", True, True], + [qml.device("lightning.qubit", wires=5), "parameter-shift", False, False], ] interface_and_qubit_device_and_diff_method = [ ["auto"] + inner_list for inner_list in qubit_device_and_diff_method @@ -43,33 +57,33 @@ @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQNode: """Test that using the QNode with JAX integrates with the PennyLane stack""" - def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): + def test_execution_with_interface( + self, dev, diff_method, grad_on_execution, interface, device_vjp + ): """Test execution works with the interface""" - if diff_method == "backprop": - pytest.skip("Test does not support backprop") - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)) - a = np.array(0.1, requires_grad=True) + a = jax.numpy.array(0.1) jax.jit(circuit)(a) assert circuit.interface == interface @@ -82,7 +96,9 @@ def circuit(a): assert isinstance(grad, jax.Array) assert grad.shape == () - def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_changing_trainability( + self, dev, diff_method, grad_on_execution, interface, device_vjp, tol + ): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -91,13 +107,12 @@ def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, i a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - dev = qml.device(dev_name, wires=2) - @qnode( dev, interface=interface, diff_method="parameter-shift", grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -130,21 +145,21 @@ def circuit(a, b): circuit(a, b) assert circuit.qtape.trainable_params == [1] - def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): + def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test classical processing within the quantum tape""" a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) c = jax.numpy.array(0.3) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -159,21 +174,24 @@ def circuit(a, b, c): assert len(res) == 2 - def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_matrix_parameter( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test that the jax interface works correctly with a matrix parameter""" - U = jax.numpy.array([[0, 1], [1, 0]]) - a = jax.numpy.array(0.1) - num_wires = 2 + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + U = jax.numpy.array([[0, 1], [1, 0]]) + a = jax.numpy.array(0.1) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -186,16 +204,19 @@ def circuit(U, a): if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [1] - def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_differentiable_expand( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test that operation and nested tape expansion is differentiable""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs = { - "sampler_rng": SEED_FOR_SPSA, - "num_directions": 10, - } + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA class U3(qml.U3): @@ -207,13 +228,6 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) @@ -222,6 +236,7 @@ def decomposition(self): diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, p): @@ -248,15 +263,12 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): + def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test setting jacobian options""" if diff_method != "finite-diff": pytest.skip("Test only applies to finite diff.") - a = np.array([0.1, 0.2], requires_grad=True) - dev = qml.device(dev_name, wires=1) - @qnode( dev, interface=interface, @@ -264,6 +276,7 @@ def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interf h=1e-8, approx_order=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -278,36 +291,37 @@ def circuit(a): @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestVectorValuedQNode: """Test that using vector-valued QNodes with JAX integrate with the PennyLane stack""" - def test_diff_expval_expval(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_expval_expval( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test jacobian calculation""" - + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} if diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -349,29 +363,30 @@ def circuit(a, b): assert res[1][1].shape == () assert np.allclose(res[1][1], expected[1][1], atol=tol, rtol=0) - def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_jacobian_no_evaluate( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test jacobian calculation when no prior circuit evaluation has been performed""" - + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} + if diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -412,22 +427,21 @@ def circuit(a, b): assert r.shape == () assert np.allclose(r, e, atol=tol, rtol=0) - def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_single_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with a single prob output""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + if diff_method == "spsa": + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -436,6 +450,7 @@ def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, inter diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -465,22 +480,21 @@ def circuit(x, y): assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0) assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0) - def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_multi_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + if diff_method == "spsa": + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 3 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -489,6 +503,7 @@ def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interf diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -551,22 +566,21 @@ def circuit(x, y): assert jac[1][1].shape == (4,) assert np.allclose(jac[1][1], expected_1[1], atol=tol, rtol=0) - def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_expval_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + if diff_method == "spsa": + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -575,6 +589,7 @@ def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, inter diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -628,23 +643,19 @@ def circuit(x, y): assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) def test_diff_expval_probs_sub_argnums( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs with less trainable parameters (argnums) than parameters.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - tol = TOL_FOR_SPSA + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -653,6 +664,7 @@ def test_diff_expval_probs_sub_argnums( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -685,19 +697,21 @@ def circuit(x, y): assert jac[1][0].shape == (2,) assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning does not support device vjps with jax jacobians.") gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "hadamard": + if diff_method == "hadamard": pytest.skip("Hadamard does not support var") elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA} + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -706,6 +720,7 @@ def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interfac diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -766,8 +781,8 @@ class TestShotsIntegration: remains differentiable.""" def test_diff_method_None(self, interface): - """Test jax device works with diff_method=None.""" - dev = qml.device("default.qubit.jax", wires=1, shots=10) + """Test device works with diff_method=None.""" + dev = DefaultQubit() @jax.jit @qml.qnode(dev, diff_method=None, interface=interface) @@ -777,52 +792,40 @@ def circuit(x): assert jax.numpy.allclose(circuit(jax.numpy.array(0.0)), 1) - def test_changing_shots(self, interface, mocker, tol): + @pytest.mark.skip("jax.jit does not work with sample") + def test_changing_shots(self, interface): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) + @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + circuit(a, b) # execute with shots=100 res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # no additional calls + assert res.shape == (100, 2) # pylint:disable=comparison-with-callable def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" - dev = qml.device("default.qubit.legacy", wires=2, shots=1) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) + @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - # TODO: jit when https://github.com/PennyLaneAI/pennylane/issues/3474 is resolved - res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000) - assert dev.shots == 1 + jit_cost_fn = jax.jit(cost_fn, static_argnames=["shots"]) + res = jax.grad(jit_cost_fn, argnums=[0, 1])(a, b, shots=30000) expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] assert np.allclose(res, expected, atol=0.1, rtol=0) @@ -830,44 +833,79 @@ def cost_fn(a, b): def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" # pylint: disable=unused-argument - dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = jax.numpy.array([0.543, -0.654]) spy = mocker.spy(qml, "execute") - # We're choosing interface="jax" such that backprop can be used in the - # test later - @qnode(dev, interface="jax") + @qnode(DefaultQubit(), interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) + cost_fn(a, b, shots=100) # pylint:disable=unexpected-keyword-arg # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn is qml.gradients.param_shift - - cost_fn(a, b) + assert cost_fn.gradient_fn == qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + cost_fn(a, b) + assert cost_fn.gradient_fn == "backprop" # if we set the shots to None, backprop can now be used - cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg assert spy.call_args[1]["gradient_fn"] == "backprop" - assert cost_fn.gradient_fn == "backprop" - cost_fn(a, b) - assert cost_fn.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)]) + def test_shot_vectors_single_measurements(self, interface, shots): + """Test jax-jit can work with shot vectors.""" + + dev = qml.device("default.qubit", shots=shots, seed=4747) + + @jax.jit + @qml.qnode(dev, interface=interface, diff_method="parameter-shift") + def circuit(x): + qml.RX(x, wires=0) + return qml.var(qml.PauliZ(0)) + + res = circuit(0.5) + expected = 1 - np.cos(0.5) ** 2 + assert qml.math.allclose(res[0], expected, atol=1e-2) + assert qml.math.allclose(res[1], expected, atol=3e-2) + + g = jax.jacobian(circuit)(0.5) + + expected_g = 2 * np.cos(0.5) * np.sin(0.5) + assert qml.math.allclose(g[0], expected_g, atol=2e-2) + assert qml.math.allclose(g[1], expected_g, atol=2e-2) + + @pytest.mark.parametrize("shots", [(10000, 10000), (10000, 10005)]) + def test_shot_vectors_multiple_measurements(self, interface, shots): + """Test jax-jit can work with shot vectors.""" + + dev = qml.device("default.qubit", shots=shots, seed=987548) + + @jax.jit + @qml.qnode(dev, interface=interface, diff_method="parameter-shift") + def circuit(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)), qml.probs(wires=0) + + res = circuit(0.5) + assert qml.math.allclose(res[0][0], np.cos(0.5), atol=5e-3) + assert qml.math.allclose(res[1][0], np.cos(0.5), atol=5e-3) + expected_probs = np.array([np.cos(0.25) ** 2, np.sin(0.25) ** 2]) + assert qml.math.allclose(res[0][1], expected_probs, atol=5e-3) + assert qml.math.allclose(res[1][1], expected_probs, atol=5e-3) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - def test_sampling(self, dev_name, diff_method, grad_on_execution, interface): + def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test sampling works as expected""" if grad_on_execution: pytest.skip("Sampling not possible with forward grad_on_execution differentiation.") @@ -875,17 +913,19 @@ def test_sampling(self, dev_name, diff_method, grad_on_execution, interface): if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - dev = qml.device(dev_name, wires=2, shots=10) - - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) - return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) + return qml.sample(qml.Z(0)), qml.sample(qml.s_prod(2, qml.X(0) @ qml.Y(1))) - res = jax.jit(circuit)() + res = jax.jit(circuit, static_argnames="shots")(shots=10) assert isinstance(res, tuple) @@ -894,7 +934,7 @@ def circuit(): assert isinstance(res[1], jax.Array) assert res[1].shape == (10,) - def test_counts(self, dev_name, diff_method, grad_on_execution, interface): + def test_counts(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test counts works as expected""" if grad_on_execution: pytest.skip("Sampling not possible with forward grad_on_execution differentiation.") @@ -902,10 +942,12 @@ def test_counts(self, dev_name, diff_method, grad_on_execution, interface): if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - dev = qml.device(dev_name, wires=2, shots=10) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(): qml.Hadamard(wires=[0]) @@ -916,9 +958,9 @@ def circuit(): with pytest.raises( NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts." ): - jax.jit(circuit)() + jax.jit(circuit, static_argnames="shots")(shots=10) else: - res = jax.jit(circuit)() + res = jax.jit(circuit, static_argnames="shots")(shots=10) assert isinstance(res, tuple) @@ -927,28 +969,32 @@ def circuit(): assert isinstance(res[1], dict) assert len(res[1]) == 2 - def test_chained_qnodes(self, dev_name, diff_method, grad_on_execution, interface): + def test_chained_qnodes(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test that the gradient of chained QNodes works without error""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit2(data, weights): qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) @@ -974,35 +1020,76 @@ def cost(weights): assert len(res) == 2 + def test_postselection_differentiation( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): + """Test that when postselecting with default.qubit, differentiation works correctly.""" + + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention") + elif dev.name == "lightning.qubit": + pytest.xfail("lightning qubit does not support postselection.") + + @qml.qnode( + dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + ) + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + phi = jax.numpy.array(1.23) + theta = jax.numpy.array(4.56) + + assert np.allclose(jax.jit(circuit)(phi, theta), jax.jit(expected_circuit)(theta)) + + gradient = jax.jit(jax.grad(circuit, argnums=[0, 1]))(phi, theta) + exp_theta_grad = jax.jit(jax.grad(expected_circuit))(theta) + assert np.allclose(gradient, [0.0, exp_theta_grad]) + @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQubitIntegrationHigherOrder: """Tests that ensure various qubit circuits integrate correctly when computing higher-order derivatives""" - def test_second_derivative(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_second_derivative( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test second derivative calculation of a scalar-valued QNode""" - gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not second derivative.") + if diff_method in {"adjoint", "device"}: + pytest.skip("Adjoint does not support second derivatives.") elif diff_method == "spsa": - gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA, "num_directions": 10} + gradient_kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + gradient_kwargs["num_directions"] = 20 + gradient_kwargs["h"] = H_FOR_SPSA tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1033,31 +1120,25 @@ def circuit(x): else: assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Test hessian calculation of a scalar-valued QNode""" gradient_kwargs = {} - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { "h": H_FOR_SPSA, - "num_directions": 20, + "num_directions": 40, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1091,10 +1172,12 @@ def circuit(x): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_hessian_vector_valued( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test hessian calculation of a vector-valued QNode""" gradient_kwargs = {} - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1104,18 +1187,12 @@ def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, i } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1159,11 +1236,11 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, dev_name, diff_method, interface, grad_on_execution, tol + self, dev, diff_method, interface, device_vjp, grad_on_execution, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" gradient_kwargs = {} - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1173,18 +1250,12 @@ def test_hessian_vector_valued_postprocessing( } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1231,11 +1302,11 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" gradient_kwargs = {} - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": gradient_kwargs = { @@ -1245,19 +1316,13 @@ def test_hessian_vector_valued_separate_args( } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, max_diff=2, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(a, b): @@ -1282,6 +1347,7 @@ def circuit(a, b): ] ) assert np.allclose(g, expected_g.T, atol=tol, rtol=0) + hess = jax.jit(jax.jacobian(jac_fn, argnums=[0, 1]))(a, b) expected_hess = np.array( @@ -1301,23 +1367,24 @@ def circuit(a, b): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "lightning.qubit" and diff_method == "adjoint": + pytest.xfail("lightning.qubit does not support adjoint with the state.") x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) + if not dev.wires: + dev = copy.copy(dev) + dev._wires = qml.wires.Wires([0, 1]) # pylint:disable=protected-access @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1327,7 +1394,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is np.dtype("complex128") # pylint:disable=no-member + assert res.dtype is np.dtype("complex128") probs = jax.numpy.abs(res) ** 2 return probs[0] + probs[2] @@ -1341,9 +1408,13 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector(self, state, dev_name, diff_method, grad_on_execution, interface, tol): + def test_projector( + self, state, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test that the variance of a projector is correctly returned""" gradient_kwargs = {} + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") if diff_method == "adjoint": pytest.skip("Adjoint does not support projectors") elif diff_method == "hadamard": @@ -1352,7 +1423,6 @@ def test_projector(self, state, dev_name, diff_method, grad_on_execution, interf gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=2) P = jax.numpy.array(state) x, y = 0.765, -0.654 @@ -1361,6 +1431,7 @@ def test_projector(self, state, dev_name, diff_method, grad_on_execution, interf diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1383,97 +1454,9 @@ def circuit(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) -# TODO: Add CV test when return types and custom diff are compatible -@pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": H_FOR_SPSA}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -@pytest.mark.parametrize("interface", ["auto", "jax-jit", "jax"]) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, interface, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - r = 0.543 - phi = -0.654 - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - res = circuit(r, phi) - expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = jax.grad(circuit, argnums=[0, 1])(r, phi) - expected = np.array( - [ - 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, - 2 * np.sinh(2 * r) * np.sin(2 * phi), - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_second_order_observable(self, diff_method, kwargs, interface, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - n = 0.12 - a = 0.765 - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) - - res = circuit(n, a) - expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = jax.grad(circuit, argnums=[0, 1])(n, a) - expected = np.array([2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -# TODO: add support for fwd grad_on_execution to JAX-JIT -@pytest.mark.parametrize("interface", ["auto", "jax-jit"]) -def test_adjoint_reuse_device_state(mocker, interface): - """Tests that the jax interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=1) - - @qnode(dev, interface=interface, diff_method="adjoint") - def circ(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - jax.grad(circ)(1.0) - assert circ.device.num_executions == 1 - - spy.assert_called_with(mocker.ANY, use_device_state=True) - - @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly @@ -1481,20 +1464,13 @@ class TestTapeExpansion: @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff, interface + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - class PhaseShift(qml.PhaseShift): grad_method = None @@ -1507,6 +1483,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, + device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1517,16 +1494,19 @@ def circuit(x, y): x = jax.numpy.array(0.5) y = jax.numpy.array(0.7) circuit(x, y) + jax.grad(circuit, argnums=[0])(x, y) @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, interface, mocker, tol + self, dev, diff_method, grad_on_execution, max_diff, device_vjp, interface, mocker, tol ): """Test that the Hamiltonian is not expanded if there are non-commuting groups and the number of shots is None and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} + if dev.name == "param_shift.qubit": + pytest.xfail("gradients transforms have a different vjp shape convention.") if diff_method == "adjoint": pytest.skip("The adjoint method does not yet support Hamiltonians") elif diff_method == "hadamard": @@ -1539,16 +1519,17 @@ def test_hamiltonian_expansion_analytic( } tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3, shots=None) spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] + @jax.jit @qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1592,14 +1573,16 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, interface, max_diff, mocker + def test_hamiltonian_finite_shots( + self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff ): - """Test that the Hamiltonian is expanded if there + """Test that the Hamiltonian is correctly measured if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} tol = 0.3 + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") if diff_method in ("adjoint", "backprop", "finite-diff"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "hadamard": @@ -1608,8 +1591,6 @@ def test_hamiltonian_expansion_finite_shots( gradient_kwargs = {"sampler_rng": SEED_FOR_SPSA, "h": H_FOR_SPSA, "num_directions": 20} tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3, shots=50000) - spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1618,6 +1599,7 @@ def test_hamiltonian_expansion_finite_shots( diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=max_diff, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1633,13 +1615,12 @@ def circuit(data, weights, coeffs): c = jax.numpy.array([-0.6543, 0.24, 0.54]) # test output - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) # pylint:disable=unexpected-keyword-arg expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_called() # test gradients - grad = jax.grad(circuit, argnums=[1, 2])(d, w, c) + grad = jax.grad(circuit, argnums=[1, 2])(d, w, c, shots=50000) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1664,27 +1645,21 @@ def circuit(data, weights, coeffs): # assert np.allclose(grad2_w_c, expected, atol=tol) def test_vmap_compared_param_broadcasting( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when vectorized=True is specified for the callback when caching is disabled.""" + if ( + dev.name == "default.qubit" + and diff_method == "adjoint" + and grad_on_execution + and not device_vjp + ): + pytest.xfail("adjoint is incompatible with parameter broadcasting.") interface = "jax-jit" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not yet support Hamiltonians") - elif diff_method == "hadamard": - pytest.skip("The Hadamard method does not yet support Hamiltonians") - - if diff_method == "backprop": - pytest.skip( - "The backprop method does not yet support parameter-broadcasting with Hamiltonians" - ) - phys_qubits = 2 - if diff_method == "hadamard": - phys_qubits = 3 n_configs = 5 pars_q = np.random.rand(n_configs, 2) - dev = qml.device(dev_name, wires=tuple(range(phys_qubits)), shots=None) def minimal_circ(params): @qml.qnode( @@ -1692,42 +1667,80 @@ def minimal_circ(params): interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, cache=None, ) def _measure_operator(): qml.RY(params[..., 0], wires=0) qml.RY(params[..., 1], wires=1) - op = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)]) - return qml.expval(op) + return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) res = _measure_operator() return res - assert np.allclose( - jax.jit(minimal_circ)(pars_q), jax.jit(jax.vmap(minimal_circ))(pars_q), tol - ) + res1 = jax.jit(minimal_circ)(pars_q) + res2 = jax.jit(jax.vmap(minimal_circ))(pars_q) + assert np.allclose(res1, res2, tol) def test_vmap_compared_param_broadcasting_multi_output( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when vectorized=True is specified for the callback when caching is disabled and when multiple output values are returned.""" + if ( + dev.name == "default.qubit" + and diff_method == "adjoint" + and grad_on_execution + and not device_vjp + ): + pytest.xfail("adjoint is incompatible with parameter broadcasting.") interface = "jax-jit" - if diff_method == "adjoint": - pytest.skip("The adjoint method does not yet support Hamiltonians") - elif diff_method == "hadamard": - pytest.skip("The Hadamard method does not yet support Hamiltonians") - if diff_method == "backprop": - pytest.skip( - "The backprop method does not yet support parameter-broadcasting with Hamiltonians" + n_configs = 5 + pars_q = np.random.rand(n_configs, 2) + + def minimal_circ(params): + @qml.qnode( + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + cache=None, ) + def _measure_operator(): + qml.RY(params[..., 0], wires=0) + qml.RY(params[..., 1], wires=1) + return qml.expval(qml.Z(0) @ qml.Z(1)), qml.expval(qml.X(0) @ qml.X(1)) + + res = _measure_operator() + return res + + res1, res2 = jax.jit(minimal_circ)(pars_q) + vres1, vres2 = jax.jit(jax.vmap(minimal_circ))(pars_q) + assert np.allclose(res1, vres1, tol) + assert np.allclose(res2, vres2, tol) + + def test_vmap_compared_param_broadcasting_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): + """Test that jax.vmap works just as well as parameter-broadcasting with JAX JIT on the forward pass when + vectorized=True is specified for the callback when caching is disabled and when multiple output values + are returned.""" + if ( + dev.name == "default.qubit" + and diff_method == "adjoint" + and grad_on_execution + and not device_vjp + ): + pytest.xfail("adjoint is incompatible with parameter broadcasting.") + elif dev.name == "lightning.qubit" and diff_method == "adjoint": + pytest.xfail("lightning adjoign cannot differentiate probabilities.") + interface = "jax-jit" - phys_qubits = 2 n_configs = 5 pars_q = np.random.rand(n_configs, 2) - dev = qml.device(dev_name, wires=tuple(range(phys_qubits)), shots=None) def minimal_circ(params): @qml.qnode( @@ -1735,14 +1748,13 @@ def minimal_circ(params): interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, cache=None, ) def _measure_operator(): qml.RY(params[..., 0], wires=0) qml.RY(params[..., 1], wires=1) - op1 = qml.Hamiltonian([1.0], [qml.PauliZ(0) @ qml.PauliZ(1)]) - op2 = qml.Hamiltonian([1.0], [qml.PauliX(0) @ qml.PauliX(1)]) - return qml.expval(op1), qml.expval(op2) + return qml.probs(wires=0), qml.probs(wires=1) res = _measure_operator() return res @@ -1758,23 +1770,25 @@ def _measure_operator(): @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestJIT: """Test JAX JIT integration with the QNode and automatic resolution of the correct JAX interface variant.""" - def test_gradient(self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface): + def test_gradient( + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface + ): """Test derivative calculation of a scalar valued QNode""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - gradient_kwargs = {} - if diff_method == "spsa": + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if device_vjp and jacobian == jax.jacfwd: + pytest.skip("device vjps not compatible with forward diff.") + elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA @@ -1783,6 +1797,7 @@ def test_gradient(self, dev_name, diff_method, grad_on_execution, jacobian, tol, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x): @@ -1806,7 +1821,9 @@ def circuit(x): "ignore:Requested adjoint differentiation to be computed with finite shots." ) @pytest.mark.parametrize("shots", [10, 1000]) - def test_hermitian(self, dev_name, diff_method, grad_on_execution, shots, jacobian, interface): + def test_hermitian( + self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface + ): """Test that the jax device works with qml.Hermitian and jitting even when shots>0. @@ -1814,13 +1831,6 @@ def test_hermitian(self, dev_name, diff_method, grad_on_execution, shots, jacobi to different reasons, hence the parametrization in the test. """ # pylint: disable=unused-argument - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - if diff_method == "backprop": pytest.skip("Backpropagation is unsupported if shots > 0.") @@ -1830,7 +1840,11 @@ def test_hermitian(self, dev_name, diff_method, grad_on_execution, shots, jacobi projector = np.array(qml.matrix(qml.PauliZ(0) @ qml.PauliZ(1))) @qml.qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circ(projector): return qml.expval(qml.Hermitian(projector, wires=range(2))) @@ -1843,23 +1857,20 @@ def circ(projector): ) @pytest.mark.parametrize("shots", [10, 1000]) def test_probs_obs_none( - self, dev_name, diff_method, grad_on_execution, shots, jacobian, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, jacobian, interface ): """Test that the jax device works with qml.probs, a MeasurementProcess that has obs=None even when shots>0.""" # pylint: disable=unused-argument - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - if diff_method in ["backprop", "adjoint"]: pytest.skip("Backpropagation is unsupported if shots > 0.") @qml.qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(): return qml.probs(wires=0) @@ -1870,23 +1881,30 @@ def circuit(): # reason="Non-trainable parameters are not being correctly unwrapped by the interface" # ) def test_gradient_subset( - self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface ): """Test derivative calculation of a scalar valued QNode with respect to a subset of arguments""" - if diff_method == "spsa" and not grad_on_execution: + if diff_method == "spsa" and not grad_on_execution and not device_vjp: pytest.xfail(reason="incorrect jacobian results") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") - if diff_method == "hadamard" and not grad_on_execution: - pytest.xfail(reason="XLA raised wire error") + if diff_method == "device" and not grad_on_execution and device_vjp: + pytest.xfail(reason="various runtime-related errors") + + if diff_method == "adjoint" and device_vjp and jacobian is jax.jacfwd: + pytest.xfail(reason="TypeError applying forward-mode autodiff.") a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - dev = qml.device(dev_name, wires=1) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a, wires=0) @@ -1903,20 +1921,17 @@ def circuit(a, b, c): assert np.allclose(g, expected_g, atol=tol, rtol=0) def test_gradient_scalar_cost_vector_valued_qnode( - self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface ): """Test derivative calculation of a scalar valued cost function that uses the output of a vector-valued QNode""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - gradient_kwargs = {} - if diff_method == "adjoint": - pytest.xfail(reason="Adjoint does not support probs.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + elif jacobian == jax.jacfwd and device_vjp: + pytest.skip("device vjps are not compatible with forward differentiation.") elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA @@ -1926,6 +1941,7 @@ def test_gradient_scalar_cost_vector_valued_qnode( diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1958,21 +1974,26 @@ def cost(x, y, idx): assert np.allclose(g0, expected_g[0][idx], atol=tol, rtol=0) assert np.allclose(g1, expected_g[1][idx], atol=tol, rtol=0) + # pylint: disable=unused-argument def test_matrix_parameter( - self, dev_name, diff_method, grad_on_execution, jacobian, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, tol, interface ): """Test that the JAX-JIT interface works correctly with a matrix parameter""" - # pylint: disable=unused-argument - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("device vjps are not compatible with forward differentiation.") + # pylint: disable=unused-argument @qml.qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circ(p, U): qml.QubitUnitary(U, wires=0) @@ -1984,7 +2005,7 @@ def circ(p, U): res = jax.jit(circ)(p, U) assert np.allclose(res, -np.cos(p), atol=tol, rtol=0) - jac_fn = jax.jit(jax.grad(circ, argnums=(0))) + jac_fn = jax.jit(jacobian(circ, argnums=0)) res = jac_fn(p, U) assert np.allclose(res, np.sin(p), atol=tol, rtol=0) @@ -1992,27 +2013,31 @@ def circ(p, U): @pytest.mark.parametrize("shots", [None, 10000]) @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestReturn: """Class to test the shape of the Grad/Jacobian with different return types.""" def test_grad_single_measurement_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For one measurement and one param, the gradient is a float.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -2021,27 +2046,30 @@ def circuit(a): a = jax.numpy.array(0.1) - grad = jax.jit(jacobian(circuit))(a) + grad = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -2051,7 +2079,9 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - grad = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) + grad = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( + a, b, shots=shots + ) assert isinstance(grad, tuple) assert len(grad) == 2 @@ -2059,21 +2089,24 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2082,31 +2115,31 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - grad = jax.jit(jacobian(circuit))(a) + grad = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -2115,30 +2148,31 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -2148,7 +2182,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")(a, b, shots=shots) assert isinstance(jac, tuple) @@ -2159,24 +2193,25 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2184,29 +2219,33 @@ def circuit(a): return qml.probs(wires=[0, 1]) a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2214,7 +2253,9 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( + par_0, par_1, shots=shots + ) assert isinstance(jac, tuple) @@ -2233,20 +2274,24 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2255,7 +2300,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2267,23 +2312,29 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_var_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of var.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - dev = qml.device(dev_name, wires=2, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if diff_method == "adjoint": + pytest.skip("adjoint supports either all expvals or only diagonal measurements") par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2291,7 +2342,9 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")( + par_0, par_1, shots=shots + ) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2311,20 +2364,26 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_var_var_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of var.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - dev = qml.device(dev_name, wires=2, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if diff_method == "adjoint": + pytest.skip("adjoint supports either all expvals or only diagonal measurements") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2333,7 +2392,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2345,23 +2404,24 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a single params return an array.""" + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if device_vjp and jacobian == jax.jacfwd: + pytest.skip("device vjp not compatible with forward differentiation.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -2370,7 +2430,7 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2382,23 +2442,24 @@ def circuit(a): assert jac[1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -2408,7 +2469,7 @@ def circuit(a, b): a = np.array(0.1, requires_grad=True) b = np.array(0.2, requires_grad=True) - jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) + jac = jax.jit(jacobian(circuit, argnums=[0, 1]), static_argnames="shots")(a, b, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2428,23 +2489,24 @@ def circuit(a, b): assert jac[1][1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transforms have a different vjp shape convention.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2453,7 +2515,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jax.jit(jacobian(circuit))(a) + jac = jax.jit(jacobian(circuit), static_argnames="shots")(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2474,23 +2536,17 @@ def circuit(a): @pytest.mark.parametrize("hessian", hessian_fn) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestReturnHessian: """Class to test the shape of the Hessian with different return types.""" def test_hessian_expval_multiple_params( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") par_0 = jax.numpy.array(0.1) @@ -2502,6 +2558,7 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2529,20 +2586,13 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2551,6 +2601,7 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2564,12 +2615,10 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_var_multiple_params( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, device_vjp, grad_on_execution, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) - - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") @@ -2583,6 +2632,7 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2610,16 +2660,14 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") - dev = qml.device(dev_name, wires=2) - params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2628,6 +2676,7 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2641,17 +2690,10 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of non commuting obs.") @@ -2665,6 +2707,7 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2691,22 +2734,15 @@ def circuit(x, y): assert h_comp.shape == (2,) def test_hessian_probs_expval_multiple_param_array( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of non commuting obs.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2715,6 +2751,7 @@ def test_hessian_probs_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2733,12 +2770,10 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) - - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") @@ -2752,6 +2787,7 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2778,16 +2814,14 @@ def circuit(x, y): assert h_comp.shape == (2,) def test_hessian_probs_var_multiple_param_array( - self, dev_name, diff_method, hessian, grad_on_execution, interface + self, dev, diff_method, hessian, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - if diff_method == "adjoint": + if diff_method in {"adjoint", "device"}: pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": pytest.skip("Test does not supports hadamard because of var.") - dev = qml.device(dev_name, wires=2) - params = jax.numpy.array([0.1, 0.2], dtype=jax.numpy.float64) @qnode( @@ -2796,6 +2830,7 @@ def test_hessian_probs_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2818,15 +2853,9 @@ def circuit(x): @pytest.mark.parametrize("diff_method", ["parameter-shift", "hadamard"]) def test_jax_device_hessian_shots(hessian, diff_method): """The hessian of multiple measurements with a multiple param array return a single array.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - dev = qml.device("default.qubit.jax", wires=num_wires, shots=10000) - - @jax.jit - @qml.qnode(dev, diff_method=diff_method, max_diff=2) + @partial(jax.jit, static_argnames="shots") + @qml.qnode(DefaultQubit(), diff_method=diff_method, max_diff=2) def circuit(x): qml.RY(x[0], wires=0) qml.RX(x[1], wires=0) @@ -2835,7 +2864,7 @@ def circuit(x): x = jax.numpy.array([1.0, 2.0]) a, b = x - hess = jax.jit(hessian(circuit))(x) + hess = jax.jit(hessian(circuit), static_argnames="shots")(x, shots=10000) expected_hess = [ [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)], @@ -2849,28 +2878,33 @@ def circuit(x): @pytest.mark.parametrize("argnums", [0, 1, [0, 1]]) @pytest.mark.parametrize("jacobian", jacobian_fn) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestSubsetArgnums: def test_single_measurement( self, interface, - dev_name, + dev, diff_method, grad_on_execution, + device_vjp, jacobian, argnums, jit_inside, tol, ): """Test single measurement with different diff methods with argnums.""" - - dev = qml.device(dev_name, wires=3) - kwargs = {} + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transform have a different vjp shape convention.") if diff_method == "spsa": - tol = TOL_FOR_SPSA kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA @qml.qnode( dev, @@ -2878,6 +2912,7 @@ def test_single_measurement( diff_method=diff_method, grad_on_execution=grad_on_execution, cache=False, + device_vjp=device_vjp, **kwargs, ) def circuit(a, b): @@ -2907,27 +2942,34 @@ def circuit(a, b): def test_multi_measurements( self, interface, - dev_name, + dev, diff_method, grad_on_execution, + device_vjp, jacobian, argnums, jit_inside, tol, ): """Test multiple measurements with different diff methods with argnums.""" - dev = qml.device(dev_name, wires=3) + if jacobian == jax.jacfwd and device_vjp: + pytest.skip("jacfwd is not compatible with device_vjp=True.") + if "lightning" in dev.name: + pytest.xfail("lightning device vjps are not compatible with jax jaocbians") + if dev.name == "param_shift.qubit": + pytest.xfail("gradient transform have a different vjp shape convention.") kwargs = {} if diff_method == "spsa": - tol = TOL_FOR_SPSA kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + tol = TOL_FOR_SPSA @qml.qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **kwargs, ) def circuit(a, b): @@ -2953,3 +2995,72 @@ def circuit(a, b): else: assert np.allclose(jac[0], expected[0], atol=tol) assert np.allclose(jac[1], expected[1], atol=tol) + + +class TestSinglePrecision: + """Tests for compatibility with single precision mode.""" + + # pylint: disable=import-outside-toplevel + def test_type_conversion_fallback(self): + """Test that if the type isn't int, float, or complex, we still have a fallback.""" + from pennylane.workflow.interfaces.jax_jit import _jax_dtype + + assert _jax_dtype(bool) == jax.numpy.dtype(bool) + + @pytest.mark.parametrize("diff_method", ("adjoint", "parameter-shift")) + def test_float32_return(self, diff_method): + """Test that jax jit works when float64 mode is disabled.""" + jax.config.update("jax_enable_x64", False) + + try: + + @jax.jit + @qml.qnode(qml.device("default.qubit"), diff_method=diff_method) + def circuit(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)) + + grad = jax.grad(circuit)(jax.numpy.array(0.1)) + assert qml.math.allclose(grad, -np.sin(0.1)) + finally: + jax.config.update("jax_enable_x64", True) + jax.config.update("jax_enable_x64", True) + + @pytest.mark.parametrize("diff_method", ("adjoint", "finite-diff")) + def test_complex64_return(self, diff_method): + """Test that jax jit works with differentiating the state.""" + jax.config.update("jax_enable_x64", False) + + try: + tol = 2e-2 if diff_method == "finite-diff" else 1e-6 + + @jax.jit + @qml.qnode(qml.device("default.qubit", wires=1), diff_method=diff_method) + def circuit(x): + qml.RX(x, wires=0) + return qml.state() + + j = jax.jacobian(circuit, holomorphic=True)(jax.numpy.array(0.1 + 0j)) + assert qml.math.allclose(j, [-np.sin(0.05) / 2, -np.cos(0.05) / 2 * 1j], atol=tol) + + finally: + jax.config.update("jax_enable_x64", True) + jax.config.update("jax_enable_x64", True) + + def test_int32_return(self): + """Test that jax jit forward execution works with samples and int32""" + + jax.config.update("jax_enable_x64", False) + + try: + + @jax.jit + @qml.qnode(qml.device("default.qubit", shots=10), diff_method=qml.gradients.param_shift) + def circuit(x): + qml.RX(x, wires=0) + return qml.sample(wires=0) + + _ = circuit(jax.numpy.array(0.1)) + finally: + jax.config.update("jax_enable_x64", True) + jax.config.update("jax_enable_x64", True) diff --git a/tests/interfaces/test_jax_qnode.py b/tests/interfaces/test_jax_qnode.py index fa88e315155..3ea650f96e8 100644 --- a/tests/interfaces/test_jax_qnode.py +++ b/tests/interfaces/test_jax_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 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. @@ -12,27 +12,39 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the JAX-Python interface with a QNode""" -# pylint: disable=too-many-arguments,too-few-public-methods,too-many-public-methods +# pylint: disable=no-member, too-many-arguments, unexpected-keyword-arg, use-implicit-booleaness-not-comparison + +from itertools import product + +import numpy as np import pytest import pennylane as qml -from pennylane import numpy as np from pennylane import qnode -from pennylane.tape import QuantumScript - -qubit_device_and_diff_method = [ - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], +from pennylane.devices import DefaultQubit + +device_seed = 42 + +# device, diff_method, grad_on_execution, device_vjp +device_and_diff_method = [ + [DefaultQubit(seed=device_seed), "backprop", True, False], + [DefaultQubit(seed=device_seed), "finite-diff", False, False], + [DefaultQubit(seed=device_seed), "parameter-shift", False, False], + [DefaultQubit(seed=device_seed), "adjoint", True, False], + [DefaultQubit(seed=device_seed), "adjoint", False, False], + [DefaultQubit(seed=device_seed), "adjoint", True, True], + [DefaultQubit(seed=device_seed), "adjoint", False, True], + [DefaultQubit(seed=device_seed), "spsa", False, False], + [DefaultQubit(seed=device_seed), "hadamard", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", False, True], + [qml.device("lightning.qubit", wires=5), "adjoint", True, True], + [qml.device("lightning.qubit", wires=5), "adjoint", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", True, False], ] -interface_and_qubit_device_and_diff_method = [ - ["auto"] + inner_list for inner_list in qubit_device_and_diff_method -] + [["jax"] + inner_list for inner_list in qubit_device_and_diff_method] +interface_and_device_and_diff_method = [ + ["auto"] + inner_list for inner_list in device_and_diff_method +] + [["jax"] + inner_list for inner_list in device_and_diff_method] pytestmark = pytest.mark.jax @@ -46,46 +58,49 @@ @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", interface_and_device_and_diff_method ) class TestQNode: """Test that using the QNode with JAX integrates with the PennyLane stack""" - def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): + def test_execution_with_interface( + self, dev, diff_method, grad_on_execution, interface, device_vjp + ): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)) - a = np.array(0.1, requires_grad=True) + a = jax.numpy.array(0.1) circuit(a) assert circuit.interface == interface - # the tape is able to deduce trainable parameters - assert circuit.qtape.trainable_params == [0] + # jax doesn't set trainable parameters on regular execution + assert circuit.qtape.trainable_params == [] # gradients should work grad = jax.grad(circuit)(a) assert isinstance(grad, jax.Array) + # the tape is able to deduce trainable parameters + assert circuit.qtape.trainable_params == [0] assert grad.shape == () - def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_changing_trainability( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): # pylint:disable=unused-argument """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -94,14 +109,7 @@ def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, i a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - dev = qml.device(dev_name, wires=2) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, interface=interface, diff_method="parameter-shift") def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -127,27 +135,18 @@ def circuit(a, b): expected = [-np.sin(a) + np.sin(a) * np.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - # trainability also updates on evaluation - a = np.array(0.54, requires_grad=False) - b = np.array(0.8, requires_grad=True) - circuit(a, b) - assert circuit.qtape.trainable_params == [1] - - def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): + def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test classical processing within the quantum tape""" a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) c = jax.numpy.array(0.3) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -162,21 +161,20 @@ def circuit(a, b, c): assert len(res) == 2 - def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_matrix_parameter( + self, dev, diff_method, grad_on_execution, interface, device_vjp, tol + ): """Test that the jax interface works correctly with a matrix parameter""" U = jax.numpy.array([[0, 1], [1, 0]]) a = jax.numpy.array(0.1) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -189,39 +187,32 @@ def circuit(U, a): if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [1] - def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_differentiable_expand( + self, dev, diff_method, grad_on_execution, interface, device_vjp, tol + ): """Test that operation and nested tape expansion is differentiable""" + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) if diff_method == "spsa": - spsa_kwargs = dict( - sampler_rng=np.random.default_rng(SEED_FOR_SPSA), - num_directions=10, - ) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 10 tol = TOL_FOR_SPSA - class U3(qml.U3): + class U3(qml.U3): # pylint:disable=too-few-public-methods def decomposition(self): theta, phi, lam = self.data wires = self.wires + return [ + qml.Rot(lam, theta, -lam, wires=wires), + qml.PhaseShift(phi + lam, wires=wires), + ] - with qml.queuing.AnnotatedQueue() as q_tape: - qml.Rot(lam, theta, -lam, wires=wires) - qml.PhaseShift(phi + lam, wires=wires) - - tape = QuantumScript.from_queue(q_tape) - return tape.operations - - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) a = jax.numpy.array(0.1) p = jax.numpy.array([0.1, 0.2, 0.3]) @@ -250,23 +241,16 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): + def test_jacobian_options( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): # pylint:disable=unused-argument """Test setting jacobian options""" if diff_method != "finite-diff": pytest.skip("Test only applies to finite diff.") - a = np.array([0.1, 0.2], requires_grad=True) - - dev = qml.device(dev_name, wires=1) + a = jax.numpy.array([0.1, 0.2]) - @qnode( - dev, - interface=interface, - diff_method="finite-diff", - h=1e-8, - approx_order=2, - grad_on_execution=grad_on_execution, - ) + @qnode(dev, interface=interface, diff_method="finite-diff", h=1e-8, approx_order=2) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -276,30 +260,31 @@ def circuit(a): @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method ) class TestVectorValuedQNode: """Test that using vector-valued QNodes with JAX integrate with the PennyLane stack""" - def test_diff_expval_expval(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_expval_expval( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test jacobian calculation""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) @qnode(dev, **kwargs) def circuit(a, b): @@ -310,7 +295,7 @@ def circuit(a, b): res = circuit(a, b) - assert circuit.qtape.trainable_params == [0, 1] + assert circuit.qtape.trainable_params == [] assert isinstance(res, tuple) assert len(res) == 2 @@ -320,6 +305,7 @@ def circuit(a, b): res = jax.jacobian(circuit, argnums=[0, 1])(a, b) expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) + assert circuit.qtape.trainable_params == [0, 1] assert isinstance(res, tuple) assert len(res) == 2 @@ -339,11 +325,20 @@ def circuit(a, b): assert res[1][1].shape == () assert np.allclose(res[1][1], expected[1][1], atol=tol, rtol=0) - def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_jacobian_no_evaluate( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test jacobian calculation when no prior circuit evaluation has been performed""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA @@ -351,13 +346,6 @@ def test_jacobian_no_evaluate(self, dev_name, diff_method, grad_on_execution, in a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -373,11 +361,10 @@ def circuit(a, b): expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - for _res, _exp in zip(res, expected): - for r, e in zip(_res, _exp): - assert isinstance(r, jax.numpy.ndarray) - assert r.shape == () - assert np.allclose(r, e, atol=tol, rtol=0) + assert isinstance(res[0][0], jax.numpy.ndarray) + for i, j in product((0, 1), (0, 1)): + assert res[i][j].shape == () + assert np.allclose(res[i][j], expected[i][j], atol=tol, rtol=0) # call the Jacobian with new parameters a = jax.numpy.array(0.6) @@ -390,30 +377,27 @@ def circuit(a, b): expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - for _res, _exp in zip(res, expected): - for r, e in zip(_res, _exp): - assert isinstance(r, jax.numpy.ndarray) - assert r.shape == () - assert np.allclose(r, e, atol=tol, rtol=0) + for i, j in product((0, 1), (0, 1)): + assert res[i][j].shape == () + assert np.allclose(res[i][j], expected[i][j], atol=tol, rtol=0) - def test_diff_single_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_single_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with a single prob output""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -445,24 +429,24 @@ def circuit(x, y): assert np.allclose(res[0], expected.T[0], atol=tol, rtol=0) assert np.allclose(res[1], expected.T[1], atol=tol, rtol=0) - def test_diff_multi_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_multi_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with multiple prob outputs""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - num_wires = 3 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -484,11 +468,11 @@ def circuit(x, y): ] assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == (2,) + assert res[0].shape == (2,) # pylint:disable=comparison-with-callable assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (4,) + assert res[1].shape == (4,) # pylint:disable=comparison-with-callable assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -527,24 +511,23 @@ def circuit(x, y): assert jac[1][1].shape == (4,) assert np.allclose(jac[1][1], expected_1[1], atol=tol, rtol=0) - def test_diff_expval_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_expval_probs( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -561,11 +544,11 @@ def circuit(x, y): assert len(res) == 2 assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == () + assert res[0].shape == () # pylint:disable=comparison-with-callable assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) + assert res[1].shape == (2,) # pylint:disable=comparison-with-callable assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -599,32 +582,28 @@ def circuit(x, y): assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) def test_diff_expval_probs_sub_argnums( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs with less trainable parameters (argnums) than parameters.""" kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "spsa": - tol = TOL_FOR_SPSA + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + tol = TOL_FOR_SPSA x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) + if diff_method == "adjoint": + x = x + 0j + y = y + 0j + @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -633,6 +612,8 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1]) + if "lightning" in dev.name: + pytest.xfail("lightning does not support measuring probabilities with adjoint.") jac = jax.jacobian(circuit, argnums=[0])(x, y) expected = [ @@ -657,21 +638,24 @@ def circuit(x, y): assert jac[1][0].shape == (2,) assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - def test_diff_var_probs(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_diff_var_probs(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Tests correct output shape and evaluation for a tape with prob and variance outputs""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - if diff_method == "adjoint": - pytest.skip("Adjoint does not support probs") - elif diff_method == "hadamard": + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + } + + if diff_method == "hadamard": pytest.skip("Hadamard does not support var") + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3) x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @@ -690,11 +674,11 @@ def circuit(x, y): ] assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == () + assert res[0].shape == () # pylint:disable=comparison-with-callable assert np.allclose(res[0], expected[0], atol=tol, rtol=0) assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) + assert res[1].shape == (2,) # pylint:disable=comparison-with-callable assert np.allclose(res[1], expected[1], atol=tol, rtol=0) jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) @@ -734,53 +718,40 @@ class TestShotsIntegration: remains differentiable.""" def test_diff_method_None(self, interface): - """Test jax device works with diff_method=None.""" - dev = qml.device("default.qubit.jax", wires=1, shots=10) + """Test device works with diff_method=None.""" - @qml.qnode(dev, diff_method=None, interface=interface) + @qml.qnode(DefaultQubit(), diff_method=None, interface=interface) def circuit(x): qml.RX(x, wires=0) return qml.expval(qml.PauliZ(0)) - assert jax.numpy.allclose(circuit(jax.numpy.array(0.0)), 1) + assert jax.numpy.allclose(circuit(jax.numpy.array(0.0), shots=10), 1) - def test_changing_shots(self, interface, mocker, tol): + def test_changing_shots(self, interface): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) + @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + circuit(a, b) # execute with shots=100 - res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(a, b) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # no additional calls + res = circuit(a, b, shots=100) + assert res.shape == (100, 2) # pylint: disable=comparison-with-callable def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" - dev = qml.device("default.qubit.legacy", wires=2, shots=1) a, b = jax.numpy.array([0.543, -0.654]) - @qnode(dev, diff_method=qml.gradients.param_shift, interface=interface) + @qnode(DefaultQubit(), diff_method=qml.gradients.param_shift, interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -788,51 +759,40 @@ def cost_fn(a, b): return qml.expval(qml.PauliY(1)) res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000) - assert dev.shots == 1 expected = [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] assert np.allclose(res, expected, atol=0.1, rtol=0) - def test_update_diff_method(self, mocker, interface): + def test_update_diff_method(self, interface, mocker): """Test that temporarily setting the shots updates the diff method""" - # pylint: disable=unused-argument - dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = jax.numpy.array([0.543, -0.654]) spy = mocker.spy(qml, "execute") - # We're choosing interface="jax" such that backprop can be used in the - # test later - @qnode(dev, interface="jax") + @qnode(DefaultQubit(), interface=interface) def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) + cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn is qml.gradients.param_shift - - cost_fn(a, b) + assert cost_fn.gradient_fn == qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - assert cost_fn.gradient_fn is qml.gradients.param_shift - - # if we set the shots to None, backprop can now be used - cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg - assert spy.call_args[1]["gradient_fn"] == "backprop" - assert cost_fn.gradient_fn == "backprop" + # if we use the default shots value of None, backprop can now be used cost_fn(a, b) - assert cost_fn.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift + assert cost_fn.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] == "backprop" -@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev,diff_method,grad_on_execution, device_vjp", device_and_diff_method) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" - def test_sampling(self, dev_name, diff_method, grad_on_execution): + def test_sampling(self, dev, diff_method, grad_on_execution, device_vjp): """Test sampling works as expected""" if grad_on_execution is True: pytest.skip("Sampling not possible with grad_on_execution differentiation.") @@ -840,43 +800,48 @@ def test_sampling(self, dev_name, diff_method, grad_on_execution): if diff_method == "adjoint": pytest.skip("Adjoint warns with finite shots") - dev = qml.device(dev_name, wires=2, shots=10) - - @qnode(dev, diff_method=diff_method, interface="jax", grad_on_execution=grad_on_execution) + @qnode( + dev, + diff_method=diff_method, + interface="jax", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert isinstance(res, tuple) assert isinstance(res[0], jax.Array) - assert res[0].shape == (10,) + assert res[0].shape == (10,) # pylint:disable=comparison-with-callable assert isinstance(res[1], jax.Array) - assert res[1].shape == (10,) + assert res[1].shape == (10,) # pylint:disable=comparison-with-callable - def test_counts(self, dev_name, diff_method, grad_on_execution): + def test_counts(self, dev, diff_method, grad_on_execution, device_vjp): """Test counts works as expected""" if grad_on_execution is True: pytest.skip("Sampling not possible with grad_on_execution differentiation.") if diff_method == "adjoint": - pytest.skip("Adjoint warns with finite shots") + pytest.skip("Adjoint errors with finite shots") - dev = qml.device(dev_name, wires=2, shots=10) - - @qnode(dev, diff_method=diff_method, interface="jax", grad_on_execution=grad_on_execution) + @qnode( + dev, + diff_method=diff_method, + interface="jax", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) - return ( - qml.counts(qml.PauliZ(0), all_outcomes=True), - qml.counts(qml.PauliX(1), all_outcomes=True), - ) + return qml.counts(qml.PauliZ(0)), qml.counts(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert isinstance(res, tuple) @@ -885,25 +850,32 @@ def circuit(): assert isinstance(res[1], dict) assert len(res[1]) == 2 - def test_chained_qnodes(self, dev_name, diff_method, grad_on_execution): + def test_chained_qnodes(self, dev, diff_method, grad_on_execution, device_vjp): """Test that the gradient of chained QNodes works without error""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + # pylint:disable=too-few-public-methods class Template(qml.templates.StronglyEntanglingLayers): def decomposition(self): return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] - @qnode(dev, interface="jax", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="jax", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit1(weights): Template(weights, wires=[0, 1]) return qml.expval(qml.PauliZ(0)) - @qnode(dev, interface="jax", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="jax", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit2(data, weights): qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) Template(weights, wires=[0, 1]) @@ -928,37 +900,72 @@ def cost(weights): assert len(res) == 2 + def test_postselection_differentiation(self, dev, diff_method, grad_on_execution, device_vjp): + """Test that when postselecting with default.qubit, differentiation works correctly.""" + + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") + + @qml.qnode( + dev, + diff_method=diff_method, + interface="jax", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface="jax", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + phi = jax.numpy.array(1.23) + theta = jax.numpy.array(4.56) + + assert np.allclose(circuit(phi, theta), expected_circuit(theta)) + + gradient = jax.grad(circuit, argnums=[0, 1])(phi, theta) + exp_theta_grad = jax.grad(expected_circuit)(theta) + assert np.allclose(gradient, [0.0, exp_theta_grad]) + @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method ) class TestQubitIntegrationHigherOrder: """Tests that ensure various qubit circuits integrate correctly when computing higher-order derivatives""" - def test_second_derivative(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_second_derivative( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test second derivative calculation of a scalar-valued QNode""" - kwargs = dict( - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - max_diff=2, - ) + kwargs = { + "diff_method": diff_method, + "interface": interface, + "grad_on_execution": grad_on_execution, + "device_vjp": device_vjp, + "max_diff": 2, + } + if diff_method == "adjoint": pytest.skip("Adjoint does not second derivative.") elif diff_method == "spsa": - spsa_kwargs = dict( - sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=100, h=0.001 - ) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(x): qml.RY(x[0], wires=0) @@ -987,7 +994,7 @@ def circuit(x): else: assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Test hessian calculation of a scalar-valued QNode""" gradient_kwargs = {} if diff_method == "adjoint": @@ -1000,18 +1007,12 @@ def test_hessian(self, dev_name, diff_method, grad_on_execution, interface, tol) } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1045,12 +1046,15 @@ def circuit(x): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_hessian_vector_valued( + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test hessian calculation of a vector-valued QNode""" gradient_kwargs = {} if diff_method == "adjoint": pytest.skip("Adjoint does not support second derivative.") elif diff_method == "spsa": + qml.math.random.seed(42) gradient_kwargs = { "h": H_FOR_SPSA, "num_directions": 20, @@ -1058,18 +1062,12 @@ def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, i } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1113,7 +1111,7 @@ def circuit(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, dev_name, diff_method, interface, grad_on_execution, tol + self, dev, diff_method, interface, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" gradient_kwargs = {} @@ -1127,18 +1125,12 @@ def test_hessian_vector_valued_postprocessing( } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1185,7 +1177,7 @@ def cost_fn(x): assert np.allclose(hess, expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_separate_args( - self, dev_name, diff_method, grad_on_execution, interface, tol + self, dev, diff_method, grad_on_execution, device_vjp, interface, tol ): """Test hessian calculation of a vector valued QNode that has separate input arguments""" gradient_kwargs = {} @@ -1199,18 +1191,12 @@ def test_hessian_vector_valued_separate_args( } tol = TOL_FOR_SPSA - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=2, **gradient_kwargs, ) @@ -1236,6 +1222,7 @@ def circuit(a, b): ] ) assert np.allclose(g, expected_g.T, atol=tol, rtol=0) + hess = jax.jacobian(jac_fn, argnums=[0, 1])(a, b) expected_hess = np.array( @@ -1255,23 +1242,21 @@ def circuit(a, b): else: assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_state(self, dev_name, diff_method, grad_on_execution, interface, tol): + def test_state(self, dev, diff_method, grad_on_execution, device_vjp, interface, tol): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("Lightning does not support state adjoint differentiation.") x = jax.numpy.array(0.543) y = jax.numpy.array(-0.654) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1281,7 +1266,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is np.dtype("complex128") + assert res.dtype is np.dtype("complex128") # pylint:disable=no-member probs = jax.numpy.abs(res) ** 2 return probs[0] + probs[2] @@ -1295,18 +1280,19 @@ def cost_fn(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector(self, state, dev_name, diff_method, grad_on_execution, interface, tol): + def test_projector( + self, state, dev, diff_method, grad_on_execution, device_vjp, interface, tol + ): """Test that the variance of a projector is correctly returned""" gradient_kwargs = {} if diff_method == "adjoint": - pytest.skip("Adjoint does not support projectors") - elif diff_method == "hadamard": + pytest.skip("adjoint supports all expvals or only diagonal measurements.") + if diff_method == "hadamard": pytest.skip("Hadamard does not support var.") elif diff_method == "spsa": gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=2) P = jax.numpy.array(state) x, y = 0.765, -0.654 @@ -1315,6 +1301,7 @@ def test_projector(self, state, dev_name, diff_method, grad_on_execution, interf diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(x, y): @@ -1338,94 +1325,7 @@ def circuit(x, y): @pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": H_FOR_SPSA}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -@pytest.mark.parametrize("interface", ["jax", "jax-python"]) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, interface, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - r = 0.543 - phi = -0.654 - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - res = circuit(r, phi) - expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = jax.grad(circuit, argnums=[0, 1])(r, phi) - expected = np.array( - [ - 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, - 2 * np.sinh(2 * r) * np.sin(2 * phi), - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_second_order_observable(self, diff_method, kwargs, interface, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - n = 0.12 - a = 0.765 - - @qnode(dev, interface=interface, diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) - - res = circuit(n, a) - expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - res = jax.grad(circuit, argnums=[0, 1])(n, a) - expected = np.array([2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize("interface", ["auto", "jax", "jax-python"]) -def test_adjoint_reuse_device_state(mocker, interface): - """Tests that the jax interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=1) - - @qnode(dev, interface=interface, diff_method="adjoint") - def circ(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - jax.grad(circ)(1.0) - assert circ.device.num_executions == 1 - - spy.assert_called_with(mocker.ANY, use_device_state=True) - - -@pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly @@ -1433,21 +1333,14 @@ class TestTapeExpansion: @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff, interface + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - class PhaseShift(qml.PhaseShift): + class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods grad_method = None def decomposition(self): @@ -1459,6 +1352,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, + device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1474,7 +1368,7 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, interface, mocker, tol + self, dev, diff_method, grad_on_execution, max_diff, interface, device_vjp, mocker, tol ): """Test that the Hamiltonian is not expanded if there are non-commuting groups and the number of shots is None @@ -1492,7 +1386,6 @@ def test_hamiltonian_expansion_analytic( } tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3, shots=None) spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @@ -1501,6 +1394,7 @@ def test_hamiltonian_expansion_analytic( interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=max_diff, **gradient_kwargs, ) @@ -1545,11 +1439,11 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, interface, max_diff, mocker + def test_hamiltonian_finite_shots( + self, dev, diff_method, grad_on_execution, device_vjp, interface, max_diff, mocker ): - """Test that the Hamiltonian is expanded if there - are non-commuting groups and the number of shots is finite + """Test that the Hamiltonian is correctly measured (and not expanded) + if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} tol = 0.3 @@ -1565,7 +1459,6 @@ def test_hamiltonian_expansion_finite_shots( } tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=3, shots=50000) spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @@ -1574,6 +1467,7 @@ def test_hamiltonian_expansion_finite_shots( interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, max_diff=max_diff, **gradient_kwargs, ) @@ -1590,13 +1484,13 @@ def circuit(data, weights, coeffs): c = jax.numpy.array([-0.6543, 0.24, 0.54]) # test output - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_called() + spy.assert_not_called() # test gradients - grad = jax.grad(circuit, argnums=[1, 2])(d, w, c) + grad = jax.grad(circuit, argnums=[1, 2])(d, w, c, shots=50000) expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1626,27 +1520,24 @@ def circuit(data, weights, coeffs): @pytest.mark.parametrize("shots", [None, 10000]) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", interface_and_device_and_diff_method ) -class TestReturn: +class TestReturn: # pylint:disable=too-many-public-methods """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" def test_grad_single_measurement_param( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """For one measurement and one param, the gradient is a float.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -1655,27 +1546,24 @@ def circuit(a): a = jax.numpy.array(0.1) - grad = jax.grad(circuit)(a) + grad = jax.grad(circuit)(a, shots=shots) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, shots, device_vjp, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -1685,7 +1573,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - grad = jax.grad(circuit, argnums=[0, 1])(a, b) + grad = jax.grad(circuit, argnums=[0, 1])(a, b, shots=shots) assert isinstance(grad, tuple) assert len(grad) == 2 @@ -1693,21 +1581,18 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, shots, device_vjp, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1716,14 +1601,14 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - grad = jax.grad(circuit)(a) + grad = jax.grad(circuit)(a, shots=shots) assert isinstance(grad, jax.numpy.ndarray) assert grad.shape == (2,) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" @@ -1733,15 +1618,12 @@ def test_jacobian_single_measurement_param_probs( if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -1750,14 +1632,14 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4,) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1766,15 +1648,12 @@ def test_jacobian_single_measurement_probs_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -1784,7 +1663,7 @@ def circuit(a, b): a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) - jac = jacobian(circuit, argnums=[0, 1])(a, b) + jac = jacobian(circuit, argnums=[0, 1])(a, b, shots=shots) assert isinstance(jac, tuple) @@ -1796,7 +1675,7 @@ def circuit(a, b): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1805,15 +1684,12 @@ def test_jacobian_single_measurement_probs_multiple_param_single_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1821,26 +1697,29 @@ def circuit(a): return qml.probs(wires=[0, 1]) a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, jax.numpy.ndarray) assert jac.shape == (4, 2) @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, + dev, + diff_method, + grad_on_execution, + jacobian, + shots, + interface, + device_vjp, ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - + if device_vjp and jacobian is jax.jacfwd: + pytest.skip("forward pass can't be done with registered vjp.") + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -1850,6 +1729,7 @@ def test_jacobian_expval_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1857,7 +1737,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(jac, tuple) @@ -1877,21 +1757,22 @@ def circuit(x, y): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if device_vjp and jacobian is jax.jacfwd: + pytest.skip("forward pass can't be done with registered vjp.") + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1900,7 +1781,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1913,18 +1794,16 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_var_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of var.") - elif diff_method == "hadamard": + pytest.skip("adjoint supports either all measurements or only diagonal measurements.") + if diff_method == "hadamard": pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -1934,6 +1813,7 @@ def test_jacobian_var_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1941,7 +1821,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1) + jac = jacobian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -1962,20 +1842,22 @@ def circuit(x, y): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_var_var_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, jacobian, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of var.") - elif diff_method == "hadamard": + pytest.skip("adjoint supports either all expvals or all diagonal measurements.") + if diff_method == "hadamard": pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1984,7 +1866,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1997,24 +1879,22 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface ): """The jacobian of multiple measurements with a single params return an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + if diff_method == "adjoint" and jacobian == jax.jacfwd: + pytest.skip("jacfwd doesn't like complex numbers") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -2023,7 +1903,7 @@ def circuit(a): a = jax.numpy.array(0.1) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2036,33 +1916,32 @@ def circuit(a): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + if diff_method == "adjoint" and jacobian == jax.jacfwd: + pytest.skip("jacfwd doesn't like complex numbers") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) + a = jax.numpy.array(0.1) + b = jax.numpy.array(0.2) - jac = jacobian(circuit, argnums=[0, 1])(a, b) + jac = jacobian(circuit, argnums=[0, 1])(a, b, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2083,23 +1962,22 @@ def circuit(a, b): @pytest.mark.parametrize("jacobian", jacobian_fn) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, jacobian, shots, interface + self, dev, diff_method, grad_on_execution, jacobian, device_vjp, shots, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) + if "lightning" in dev.name: + pytest.xfail("lightning device_vjp not compatible with jax.jacobian.") + if diff_method == "adjoint" and jacobian == jax.jacfwd: + pytest.skip("jacfwd doesn't like complex numbers") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -2108,7 +1986,7 @@ def circuit(a): a = jax.numpy.array([0.1, 0.2]) - jac = jacobian(circuit)(a) + jac = jacobian(circuit)(a, shots=shots) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2120,19 +1998,12 @@ def circuit(a): assert jac[1].shape == (4, 2) def test_hessian_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2145,6 +2016,7 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2152,7 +2024,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2170,7 +2042,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2178,13 +2050,6 @@ def test_hessian_expval_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2193,6 +2058,7 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2200,13 +2066,13 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit)(params) + hess = jax.hessian(circuit)(params, shots=shots) assert isinstance(hess, jax.numpy.ndarray) assert hess.shape == (2, 2) def test_hessian_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2215,7 +2081,6 @@ def test_hessian_var_multiple_params( pytest.skip("Test does not support Hadamard because of var.") if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) par_0 = jax.numpy.array(0.1) par_1 = jax.numpy.array(0.2) @@ -2226,6 +2091,7 @@ def test_hessian_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2233,7 +2099,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2251,7 +2117,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2261,8 +2127,6 @@ def test_hessian_var_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2271,6 +2135,7 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2278,22 +2143,15 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = jax.hessian(circuit)(params) + hess = jax.hessian(circuit)(params, shots=shots) assert isinstance(hess, jax.numpy.ndarray) assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2311,6 +2169,7 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2318,7 +2177,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2354,7 +2213,7 @@ def circuit(x, y): assert hess[1][1][1].shape == (2,) def test_hessian_expval_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2364,19 +2223,13 @@ def test_hessian_expval_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - params = jax.numpy.array([0.1, 0.2]) @qnode( dev, interface=interface, diff_method=diff_method, + device_vjp=device_vjp, max_diff=2, grad_on_execution=grad_on_execution, ) @@ -2386,7 +2239,7 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit)(params) + hess = jax.hessian(circuit)(params, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2398,7 +2251,7 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2408,8 +2261,6 @@ def test_hessian_probs_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = qml.numpy.array(0.1) par_1 = qml.numpy.array(0.2) @@ -2419,6 +2270,7 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2426,7 +2278,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1) + hess = jax.hessian(circuit, argnums=[0, 1])(par_0, par_1, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2462,7 +2314,7 @@ def circuit(x, y): assert hess[1][1][1].shape == (2,) def test_hessian_var_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots, interface + self, dev, diff_method, grad_on_execution, device_vjp, shots, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2472,8 +2324,6 @@ def test_hessian_var_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - params = jax.numpy.array([0.1, 0.2]) @qnode( @@ -2482,6 +2332,7 @@ def test_hessian_var_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2489,7 +2340,7 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) - hess = jax.hessian(circuit)(params) + hess = jax.hessian(circuit)(params, shots=shots) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2501,14 +2352,11 @@ def circuit(x): assert hess[1].shape == (2, 2, 2) -@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) -def test_no_ops(dev_name): +def test_no_ops(): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = qml.device(dev_name, wires=1) - - @qml.qnode(dev, interface="jax") + @qml.qnode(DefaultQubit(), interface="jax") def circuit(): qml.Hadamard(wires=0) return qml.state() diff --git a/tests/interfaces/test_jax_qnode_shot_vector.py b/tests/interfaces/test_jax_qnode_shot_vector.py index 0ef5c8291de..697a1e90223 100644 --- a/tests/interfaces/test_jax_qnode_shot_vector.py +++ b/tests/interfaces/test_jax_qnode_shot_vector.py @@ -28,9 +28,9 @@ all_shots = [(1, 20, 100), (1, (20, 1), 100), (1, (5, 4), 100)] qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], - ["default.qubit.legacy", "parameter-shift", {}], - ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}], + ["default.qubit", "finite-diff", {"h": 10e-2}], + ["default.qubit", "parameter-shift", {}], + ["default.qubit", "spsa", {"h": 10e-2, "num_directions": 20}], ] interface_and_qubit_device_and_diff_method = [ @@ -766,59 +766,9 @@ def circuit(x): assert h[1].shape == (2, 2, 2) -@pytest.mark.parametrize("shots", all_shots) -class TestReturnShotVectorsDevice: - """Test for shot vectors with device method adjoint_jacobian.""" - - def test_jac_adjoint_fwd_error(self, shots): - """Test that an error is raised for adjoint forward.""" - dev = qml.device("default.qubit.legacy", wires=1, shots=shots) - - with pytest.warns( - UserWarning, match="Requested adjoint differentiation to be computed with finite shots." - ): - - @qnode(dev, interface="jax", diff_method="adjoint", grad_on_execution=True) - def circuit(a): - qml.RY(a, wires=0) - qml.RX(0.2, wires=0) - return qml.expval(qml.PauliZ(0)) - - a = jax.numpy.array(0.1) - - if isinstance(shots, tuple): - with pytest.raises( - qml.QuantumFunctionError, - match="Adjoint does not support shot vectors.", - ): - jax.jacobian(circuit)(a) - - def test_jac_adjoint_bwd_error(self, shots): - """Test that an error is raised for adjoint backward.""" - dev = qml.device("default.qubit.legacy", wires=1, shots=shots) - - with pytest.warns( - UserWarning, match="Requested adjoint differentiation to be computed with finite shots." - ): - - @qnode(dev, interface="jax", diff_method="adjoint", grad_on_execution=False) - def circuit(a): - qml.RY(a, wires=0) - qml.RX(0.2, wires=0) - return qml.expval(qml.PauliZ(0)) - - a = jax.numpy.array(0.1) - - with pytest.raises( - qml.QuantumFunctionError, - match="Adjoint does not support shot vectors.", - ): - jax.jacobian(circuit)(a) - - qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], - ["default.qubit.legacy", "parameter-shift", {}], + ["default.qubit", "finite-diff", {"h": 10e-2}], + ["default.qubit", "parameter-shift", {}], ] shots_large = [(1000000, 900000, 800000), (1000000, (900000, 2))] diff --git a/tests/interfaces/test_set_shots.py b/tests/interfaces/test_set_shots.py index 5e6a53b83d3..efe8c9cf7ca 100644 --- a/tests/interfaces/test_set_shots.py +++ b/tests/interfaces/test_set_shots.py @@ -18,7 +18,6 @@ import pytest import pennylane as qml -from pennylane.measurements import Shots from pennylane.workflow import set_shots @@ -30,29 +29,3 @@ def test_shots_new_device_interface(): with pytest.raises(ValueError): with set_shots(dev, 10): pass - - -def test_set_with_shots_class(): - """Test that shots can be set on the old device interface with a Shots class.""" - - dev = qml.devices.DefaultQubitLegacy(wires=1) - with set_shots(dev, Shots(10)): - assert dev.shots == 10 - - assert dev.shots is None - - shot_tuples = Shots((10, 10)) - with set_shots(dev, shot_tuples): - assert dev.shots == 20 - assert dev.shot_vector == list(shot_tuples.shot_vector) - - assert dev.shots is None - - -def test_shots_not_altered_if_False(): - """Test a value of False can be passed to shots, indicating to not override - shots on the device.""" - - dev = qml.devices.DefaultQubitLegacy(wires=1) - with set_shots(dev, False): - assert dev.shots is None diff --git a/tests/interfaces/test_tensorflow.py b/tests/interfaces/test_tensorflow.py index 057f93d0038..11a54575cfa 100644 --- a/tests/interfaces/test_tensorflow.py +++ b/tests/interfaces/test_tensorflow.py @@ -1,4 +1,4 @@ -# Copyright 2018-2021 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 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. @@ -11,227 +11,32 @@ # 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. -"""Unit tests for the TensorFlow interface""" -# pylint: disable=protected-access,too-few-public-methods +"""Tensorflow specific tests for execute and default qubit 2.""" import numpy as np import pytest import pennylane as qml from pennylane import execute -from pennylane.gradients import finite_diff, param_shift +from pennylane.devices import DefaultQubit +from pennylane.gradients import param_shift pytestmark = pytest.mark.tf - -tf = pytest.importorskip("tensorflow", minversion="2.1") - - -class TestTensorFlowExecuteUnitTests: - """Unit tests for TensorFlow execution""" - - def test_jacobian_options(self, mocker): - """Test setting jacobian options""" - spy = mocker.spy(qml.gradients, "param_shift") - - a = tf.Variable([0.1, 0.2], dtype=tf.float64) - - dev = qml.device("default.qubit.legacy", wires=1) - - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute( - [tape], - dev, - gradient_fn=param_shift, - gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, - interface="tf", - )[0] - - res = t.jacobian(res, a) - - for args in spy.call_args_list: - assert args[1]["shifts"] == [(np.pi / 4,)] * 2 - - def test_incorrect_grad_on_execution(self): - """Test that an error is raised if a gradient transform - is used with grad_on_execution""" - a = tf.Variable([0.1, 0.2]) - - dev = qml.device("default.qubit.legacy", wires=1) - - with tf.GradientTape(): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - with pytest.raises( - ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" - ): - execute([tape], dev, gradient_fn=param_shift, grad_on_execution=True, interface="tf") - - def test_grad_on_execution(self, mocker): - """Test that grad on execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - a = tf.Variable([0.1, 0.2]) - spy = mocker.spy(dev, "execute_and_gradients") - - with tf.GradientTape(): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - interface="tf", - ) - - # adjoint method only performs a single device execution, but gets both result and gradient - assert dev.num_executions == 1 - spy.assert_called() - - def test_no_grad_execution(self, mocker): - """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") - a = tf.Variable([0.1, 0.2]) - - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute( - [tape], - dev, - gradient_fn="device", - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - interface="tf", - )[0] - - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() - - t.jacobian(res, a) - spy_gradients.assert_called() +tf = pytest.importorskip("tensorflow") +# pylint: disable=too-few-public-methods class TestCaching: - """Test for caching behaviour""" - - def test_cache_maxsize(self, mocker): - """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - a = tf.Variable([0.1, 0.2]) - - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, gradient_fn=param_shift, cachesize=2, interface="tf")[0] - - t.jacobian(res, a) - cache = spy.call_args.kwargs["cache"] - - assert cache.maxsize == 2 - assert cache.currsize == 2 - assert len(cache) == 2 - - def test_custom_cache(self, mocker): - """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - a = tf.Variable([0.1, 0.2]) - custom_cache = {} - - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, gradient_fn=param_shift, cache=custom_cache, interface="tf")[ - 0 - ] - - t.jacobian(res, a) - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache - - unwrapped_tape = qml.transforms.convert_to_numpy_parameters(tape)[0][0] - h = unwrapped_tape.hash - - assert h in cache - assert np.allclose(cache[h], res) - - def test_caching_param_shift(self): - """Test that, when using parameter-shift transform, - caching reduces the number of evaluations to their optimum.""" - dev = qml.device("default.qubit.legacy", wires=1) - a = tf.Variable([0.1, 0.2], dtype=tf.float64) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="tf")[0] - - # Without caching, and non-vectorized, 9 evaluations are required to compute - # the Jacobian: 1 (forward pass) + 2 (backward pass) * (2 shifts * 2 params) - with tf.GradientTape(persistent=True) as t: - res = cost(a, cache=None) - t.jacobian(res, a, experimental_use_pfor=False) - assert dev.num_executions == 9 - - # With caching, and non-vectorized, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - with tf.GradientTape(persistent=True) as t: - res = cost(a, cache=True) - t.jacobian(res, a) - assert dev.num_executions == 5 - - # In vectorized mode, 5 evaluations are required to compute - # the Jacobian regardless of caching: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - with tf.GradientTape() as t: - res = cost(a, cache=None) - t.jacobian(res, a) - assert dev.num_executions == 5 + """Tests for caching behaviour""" @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params, tol): + def test_caching_param_shift_hessian(self, num_params): """Test that, when using parameter-shift transform, caching reduces the number of evaluations to their optimum when computing Hessians.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = tf.Variable(np.arange(1, num_params + 1) / 10, dtype=tf.float64) + dev = DefaultQubit() + params = tf.Variable(tf.range(1, num_params + 1) / 10) - N = params.shape[0] + N = num_params def cost(x, cache): with qml.queuing.AnnotatedQueue() as q: @@ -242,362 +47,425 @@ def cost(x, cache): qml.RZ(x[i], wires=[i % 2]) qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], dev, gradient_fn=param_shift, cache=cache, interface="tf", max_diff=2 + return qml.execute( + [tape], dev, gradient_fn=qml.gradients.param_shift, cache=cache, max_diff=2 )[0] # No caching: number of executions is not ideal - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - res = cost(params, cache=False) - grad = t1.gradient(res, params) - hess1 = t2.jacobian(grad, params) + with qml.Tracker(dev) as tracker: + with tf.GradientTape() as jac_tape: + with tf.GradientTape() as grad_tape: + res = cost(params, cache=False) + grad = grad_tape.gradient(res, params) + hess1 = jac_tape.jacobian(grad, params) if num_params == 2: # compare to theoretical result - x, y, *_ = params * 1.0 - expected = np.array( + x, y, *_ = params + expected = tf.convert_to_tensor( [ - [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + [2 * tf.cos(2 * x) * tf.sin(y) ** 2, tf.sin(2 * x) * tf.sin(2 * y)], + [tf.sin(2 * x) * tf.sin(2 * y), -2 * tf.cos(x) ** 2 * tf.cos(2 * y)], ] ) - assert np.allclose(expected, hess1, atol=tol, rtol=0) + assert np.allclose(expected, hess1) + + expected_runs = 1 # forward pass + + # Jacobian of an involutory observable: + # ------------------------------------ + # + # 2 * N execs: evaluate the analytic derivative of + # 1 execs: Get , the expectation value of the tape with unshifted parameters. + num_shifted_evals = 2 * N + runs_for_jacobian = num_shifted_evals + 1 + expected_runs += runs_for_jacobian - nonideal_runs = dev.num_executions + # Each tape used to compute the Jacobian is then shifted again + expected_runs += runs_for_jacobian * num_shifted_evals + assert tracker.totals["executions"] == expected_runs # Use caching: number of executions is ideal - dev._num_executions = 0 - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - res = cost(params, cache=True) - grad = t1.gradient(res, params) - hess2 = t2.jacobian(grad, params) - assert np.allclose(hess1, hess2, atol=tol, rtol=0) + with qml.Tracker(dev) as tracker2: + with tf.GradientTape() as jac_tape: + with tf.GradientTape() as grad_tape: + res = cost(params, cache=True) + grad = grad_tape.gradient(res, params) + hess2 = jac_tape.jacobian(grad, params) + assert np.allclose(hess1, hess2) expected_runs_ideal = 1 # forward pass expected_runs_ideal += 2 * N # Jacobian expected_runs_ideal += N + 1 # Hessian diagonal expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert dev.num_executions == expected_runs_ideal - assert expected_runs_ideal < nonideal_runs - - -execute_kwargs_integration = [ - {"gradient_fn": param_shift, "interface": "tf"}, - {"gradient_fn": param_shift, "interface": "auto"}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - "interface": "tf", - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - "interface": "tf", - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - "interface": "auto", - }, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - "interface": "auto", - }, + assert tracker2.totals["executions"] == expected_runs_ideal + assert expected_runs_ideal < expected_runs + + +# add tests for lightning 2 when possible +# set rng for device when possible +test_matrix = [ + ({"gradient_fn": param_shift, "interface": "tensorflow"}, 100000, DefaultQubit(seed=42)), # 0 + ({"gradient_fn": param_shift, "interface": "tensorflow"}, None, DefaultQubit()), # 1 + ({"gradient_fn": "backprop", "interface": "tensorflow"}, None, DefaultQubit()), # 2 + ({"gradient_fn": "adjoint", "interface": "tensorflow"}, None, DefaultQubit()), # 3 + ({"gradient_fn": param_shift, "interface": "tf-autograph"}, 100000, DefaultQubit(seed=42)), # 4 + ({"gradient_fn": param_shift, "interface": "tf-autograph"}, None, DefaultQubit()), # 5 + ({"gradient_fn": "backprop", "interface": "tf-autograph"}, None, DefaultQubit()), # 6 + ({"gradient_fn": "adjoint", "interface": "tf-autograph"}, None, DefaultQubit()), # 7 + ({"gradient_fn": "adjoint", "interface": "tf", "device_vjp": True}, None, DefaultQubit()), # 8 ] -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) -class TestTensorFlowExecuteIntegration: - """Test the TensorFlow interface execute function +def atol_for_shots(shots): + """Return higher tolerance if finite shots.""" + return 1e-2 if shots else 1e-6 + + +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +class TestTensorflowExecuteIntegration: + """Test the tensorflow interface execute function integrates well for both forward and backward execution""" - def test_execution(self, execute_kwargs): + def test_execution(self, execute_kwargs, shots, device): """Test execution""" - dev = qml.device("default.qubit.legacy", wires=1) - a = tf.Variable(0.1) - b = tf.Variable(0.2) - with tf.GradientTape(): - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, b): + ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] + tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) + + ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] + tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + return execute([tape1, tape2], device, **execute_kwargs) - tape2 = qml.tape.QuantumScript.from_queue(q2) - res = execute([tape1, tape2], dev, **execute_kwargs) + a = tf.Variable(0.1, dtype="float64") + b = tf.constant(0.2, dtype="float64") + with device.tracker: + res = cost(a, b) + + if execute_kwargs.get("gradient_fn", None) == "adjoint" and not execute_kwargs.get( + "device_vjp", False + ): + assert device.tracker.totals["execute_and_derivative_batches"] == 1 + else: + assert device.tracker.totals["batches"] == 1 + assert device.tracker.totals["executions"] == 2 # different wires so different hashes assert len(res) == 2 assert res[0].shape == () assert res[1].shape == () - assert isinstance(res[0], tf.Tensor) - assert isinstance(res[1], tf.Tensor) - def test_scalar_jacobian(self, execute_kwargs, tol): + assert qml.math.allclose(res[0], tf.cos(a) * tf.cos(b), atol=atol_for_shots(shots)) + assert qml.math.allclose(res[1], tf.cos(a) * tf.cos(b), atol=atol_for_shots(shots)) + + def test_scalar_jacobian(self, execute_kwargs, shots, device): """Test scalar jacobian calculation""" a = tf.Variable(0.1, dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] + device_vjp = execute_kwargs.get("device_vjp", False) - res = t.jacobian(res, a) - assert res.shape == () + def cost(a): + tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - # compare to standard tape jacobian - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost(a) + res = tape.jacobian(cost_res, a, experimental_use_pfor=not device_vjp) + assert res.shape == () # pylint: disable=no-member - tape = qml.tape.QuantumScript.from_queue(q) + # compare to standard tape jacobian + tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) tape.trainable_params = [0] tapes, fn = param_shift(tape) - expected = fn(dev.batch_execute(tapes)) + expected = fn(device.execute(tapes)) assert expected.shape == () - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res, -tf.sin(a), atol=atol_for_shots(shots)) - def test_jacobian(self, execute_kwargs, tol): + def test_jacobian(self, execute_kwargs, shots, device): """Test jacobian calculation""" - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) + a = tf.Variable(0.1) + b = tf.Variable(0.2) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, max_diff=2, **execute_kwargs)[0] - res = tf.stack(res) + device_vjp = execute_kwargs.get("device_vjp", False) - expected = [np.cos(a), -np.cos(a) * np.sin(b)] - assert np.allclose(res, expected, atol=tol, rtol=0) + def cost(a, b): + ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") + + with tf.GradientTape(persistent=device_vjp) as tape: + res = cost(a, b) + expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - (agrad, bgrad) = t.jacobian(res, [a, b]) - assert agrad.shape == (2,) - assert bgrad.shape == (2,) + jac = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) + assert isinstance(jac, list) and len(jac) == 2 + assert jac[0].shape == (2,) + assert jac[1].shape == (2,) - expected = [[-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]] - assert np.allclose(expected, [agrad, bgrad], atol=tol, rtol=0) + expected = ([-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)]) + for _r, _e in zip(jac, expected): + assert np.allclose(_r, _e, atol=atol_for_shots(shots)) - def test_tape_no_parameters(self, execute_kwargs, tol): + def test_tape_no_parameters(self, execute_kwargs, shots, device): """Test that a tape with no parameters is correctly ignored during the gradient computation""" - dev = qml.device("default.qubit.legacy", wires=1) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) - x, y = 1.0 * params - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.expval(qml.PauliX(0)) + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots + ) - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(0.5, wires=0) - qml.expval(qml.PauliZ(0)) + tape2 = qml.tape.QuantumScript( + [qml.RY(tf.constant(0.5), wires=0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - tape2 = qml.tape.QuantumScript.from_queue(q2) - with qml.queuing.AnnotatedQueue() as q3: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - tape3 = qml.tape.QuantumScript.from_queue(q3) - res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs)) - res = tf.stack(res) + tape4 = qml.tape.QuantumScript( + [qml.RY(tf.constant(0.5), 0)], + [qml.probs(wires=[0, 1])], + shots=shots, + ) + return tf.reduce_sum( + qml.math.hstack( + execute([tape1, tape2, tape3, tape4], device, **execute_kwargs), + like="tensorflow", + ) + ) - expected = 1 + np.cos(0.5) + np.cos(x) * np.cos(y) - assert np.allclose(res, expected, atol=tol, rtol=0) + params = tf.Variable([0.1, 0.2]) + x, y = params - grad = t.gradient(res, params) - expected = [-np.cos(y) * np.sin(x), -np.cos(x) * np.sin(y)] - assert np.allclose(grad, expected, atol=tol, rtol=0) + with tf.GradientTape() as tape: + res = cost(params) + expected = 2 + tf.cos(0.5) + tf.cos(x) * tf.cos(y) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_reusing_quantum_tape(self, execute_kwargs, tol): - """Test re-using a quantum tape by passing new parameters""" - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) + if ( + execute_kwargs.get("interface", "") == "tf-autograph" + and execute_kwargs.get("gradient_fn", "") == "adjoint" + ): + with pytest.raises(NotImplementedError): + tape.gradient(res, params) + return - dev = qml.device("default.qubit.legacy", wires=2) + grad = tape.gradient(res, params) + expected = [-tf.cos(y) * tf.sin(x), -tf.cos(x) * tf.sin(y)] + assert np.allclose(grad, expected, atol=atol_for_shots(shots), rtol=0) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) + def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): + """Test that tapes wit different can be executed and differentiated.""" - tape = qml.tape.QuantumScript.from_queue(q) - assert tape.trainable_params == [0, 1] - res = execute([tape], dev, **execute_kwargs)[0] - res = tf.stack(res) + if ( + execute_kwargs["gradient_fn"] == "adjoint" + and execute_kwargs["interface"] == "tf-autograph" + ): + pytest.skip("Cannot compute the jacobian with adjoint-differentation and tf-autograph") - t.jacobian(res, [a, b]) + device_vjp = execute_kwargs.get("device_vjp", False) - a = tf.Variable(0.54, dtype=tf.float64) - b = tf.Variable(0.8, dtype=tf.float64) + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + shots=shots, + ) - # check that the cost function continues to depend on the - # values of the parameters for subsequent calls - with tf.GradientTape() as t: - tape = tape.bind_new_parameters([2 * a, b], [0, 1]) - res2 = execute([tape], dev, **execute_kwargs)[0] - res2 = tf.stack(res2) + tape2 = qml.tape.QuantumScript( + [qml.RY(tf.constant(0.5, dtype=tf.float64), 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] - assert np.allclose(res2, expected, atol=tol, rtol=0) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + return qml.math.hstack( + execute([tape1, tape2, tape3], device, **execute_kwargs), like="tensorflow" + ) - jac2 = t.jacobian(res2, [a, b]) - expected = [ - [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], - [0, -tf.cos(2 * a) * tf.cos(b)], - ] - assert np.allclose(jac2, expected, atol=tol, rtol=0) + params = tf.Variable([0.1, 0.2], dtype=tf.float64) + x, y = params - def test_reusing_pre_constructed_quantum_tape(self, execute_kwargs, tol): - """Test re-using a quantum tape that was previously constructed - *outside of* a gradient tape, by passing new parameters""" - a = tf.Variable(0.1, dtype=tf.float64) - b = tf.Variable(0.2, dtype=tf.float64) + with tf.GradientTape(persistent=device_vjp) as tape: + res = cost(params) + + assert isinstance(res, tf.Tensor) + assert res.shape == (4,) + + assert np.allclose(res[0], tf.cos(x) * tf.cos(y), atol=atol_for_shots(shots)) + assert np.allclose(res[1], 1, atol=atol_for_shots(shots)) + assert np.allclose(res[2], tf.cos(0.5), atol=atol_for_shots(shots)) + assert np.allclose(res[3], tf.cos(x) * tf.cos(y), atol=atol_for_shots(shots)) + + jac = tape.jacobian(res, params, experimental_use_pfor=not device_vjp) + assert isinstance(jac, tf.Tensor) + assert jac.shape == (4, 2) # pylint: disable=no-member + + assert np.allclose(jac[1:3], 0, atol=atol_for_shots(shots)) - dev = qml.device("default.qubit.legacy", wires=2) + d1 = -tf.sin(x) * tf.cos(y) + assert np.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) + + d2 = -tf.cos(x) * tf.sin(y) + assert np.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) + assert np.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) + + def test_reusing_quantum_tape(self, execute_kwargs, shots, device): + """Test re-using a quantum tape by passing new parameters""" + a = tf.Variable(0.1) + b = tf.Variable(0.2) + + tape = qml.tape.QuantumScript( + [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], + ) + assert tape.trainable_params == [0, 1] - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) + device_vjp = execute_kwargs.get("device_vjp", False) - tape = qml.tape.QuantumScript.from_queue(q) - with tf.GradientTape() as t: - tape = tape.bind_new_parameters([a, b], [0, 1]) - assert tape.trainable_params == [0, 1] - res = execute([tape], dev, **execute_kwargs)[0] - res = qml.math.stack(res) + def cost(a, b): + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return qml.math.hstack( + execute([new_tape], device, **execute_kwargs)[0], like="tensorflow" + ) - t.jacobian(res, [a, b]) + with tf.GradientTape(persistent=device_vjp) as t: + res = cost(a, b) + jac = t.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) a = tf.Variable(0.54, dtype=tf.float64) b = tf.Variable(0.8, dtype=tf.float64) - with tf.GradientTape() as t: - tape = tape.bind_new_parameters([2 * a, b], [0, 1]) - res2 = execute([tape], dev, **execute_kwargs)[0] - res2 = qml.math.stack(res2) + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + + with tf.GradientTape(persistent=device_vjp): + res2 = cost(2 * a, b) expected = [tf.cos(2 * a), -tf.cos(2 * a) * tf.sin(b)] - assert np.allclose(res2, expected, atol=tol, rtol=0) + assert np.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) - jac2 = t.jacobian(res2, [a, b]) - expected = [ + with tf.GradientTape(persistent=device_vjp) as t: + res = cost(2 * a, b) + + jac = t.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) + expected = ( [-2 * tf.sin(2 * a), 2 * tf.sin(2 * a) * tf.sin(b)], [0, -tf.cos(2 * a) * tf.cos(b)], - ] - assert np.allclose(jac2, expected, atol=tol, rtol=0) + ) + assert isinstance(jac, list) and len(jac) == 2 + for _j, _e in zip(jac, expected): + assert np.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) - def test_classical_processing(self, execute_kwargs): + def test_classical_processing(self, execute_kwargs, device, shots): """Test classical processing within the quantum tape""" a = tf.Variable(0.1, dtype=tf.float64) b = tf.constant(0.2, dtype=tf.float64) c = tf.Variable(0.3, dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=1) + device_vjp = execute_kwargs.get("device_vjp", False) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a * c, wires=0) - qml.RZ(b, wires=0) - qml.RX(c + c**2 + tf.sin(a), wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, b, c): + ops = [ + qml.RY(a * c, wires=0), + qml.RZ(b, wires=0), + qml.RX(c + c**2 + tf.sin(a), wires=0), + ] - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [0, 2] - assert tape.get_parameters() == [a * c, c + c**2 + tf.sin(a)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] + + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost(a, b, c) + + res = tape.jacobian(cost_res, [a, c], experimental_use_pfor=not device_vjp) + + # Only two arguments are trainable + assert isinstance(res, list) and len(res) == 2 + assert res[0].shape == () + assert res[1].shape == () - res = t.jacobian(res, [a, b, c]) - assert isinstance(res[0], tf.Tensor) - assert res[1] is None - assert isinstance(res[2], tf.Tensor) + # I tried getting analytic results for this circuit but I kept being wrong and am giving up - def test_no_trainable_parameters(self, execute_kwargs): + def test_no_trainable_parameters(self, execute_kwargs, shots, device): """Test evaluation and Jacobian if there are no trainable parameters""" - b = tf.constant(0.2, dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) + a = tf.constant(0.1) + b = tf.constant(0.2) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(0.2, wires=0) - qml.RX(b, wires=0) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + def cost(a, b): + ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - res = qml.math.stack(res) + with tf.GradientTape() as tape: + cost_res = cost(a, b) - assert res.shape == (2,) - assert isinstance(res, tf.Tensor) + assert cost_res.shape == (2,) + + res = tape.jacobian(cost_res, [a, b]) + assert len(res) == 2 + assert all(r is None for r in res) + + def loss(a, b): + return tf.reduce_sum(cost(a, b)) + + with tf.GradientTape() as tape: + loss_res = loss(a, b) - res = t.jacobian(res, b) - assert res is None + res = tape.gradient(loss_res, [a, b]) + assert all(r is None for r in res) - @pytest.mark.parametrize("U", [tf.constant([[0, 1], [1, 0]]), np.array([[0, 1], [1, 0]])]) - def test_matrix_parameter(self, execute_kwargs, U, tol): - """Test that the TF interface works correctly + def test_matrix_parameter(self, execute_kwargs, device, shots): + """Test that the tensorflow interface works correctly with a matrix parameter""" - a = tf.Variable(0.1, dtype=tf.float64) + U = tf.constant([[0, 1], [1, 0]], dtype=tf.complex128) + a = tf.Variable(0.1) - dev = qml.device("default.qubit.legacy", wires=2) + device_vjp = execute_kwargs.get("device_vjp", False) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, U): + ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) + return execute([tape], device, **execute_kwargs)[0] - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [1] + with tf.GradientTape(persistent=device_vjp) as tape: + res = cost(a, U) - assert np.allclose(res, -tf.cos(a), atol=tol, rtol=0) + assert np.allclose(res, -tf.cos(a), atol=atol_for_shots(shots)) - res = t.jacobian(res, a) - assert np.allclose(res, tf.sin(a), atol=tol, rtol=0) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) + assert isinstance(jac, tf.Tensor) + assert np.allclose(jac, tf.sin(a), atol=atol_for_shots(shots), rtol=0) - def test_differentiable_expand(self, execute_kwargs, tol): - """Test that operation and nested tape expansion + def test_differentiable_expand(self, execute_kwargs, device, shots): + """Test that operation and nested tapes expansion is differentiable""" + device_vjp = execute_kwargs.get("device_vjp", False) + class U3(qml.U3): + """Dummy operator.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -606,26 +474,38 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - dev = qml.device("default.qubit.legacy", wires=1) - a = np.array(0.1) - p = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) + def cost_fn(a, p): + tape = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] + ) + gradient_fn = execute_kwargs["gradient_fn"] + if gradient_fn is None: + _gradient_method = None + elif isinstance(gradient_fn, str): + _gradient_method = gradient_fn + else: + _gradient_method = "gradient-transform" + config = qml.devices.ExecutionConfig( + interface="autograd", + gradient_method=_gradient_method, + grad_on_execution=execute_kwargs.get("grad_on_execution", None), + ) + program, _ = device.preprocess(execution_config=config) + return execute([tape], device, **execute_kwargs, transform_program=program)[0] - with tf.GradientTape() as tape: - with qml.queuing.AnnotatedQueue() as q_qtape: - qml.RX(a, wires=0) - U3(p[0], p[1], p[2], wires=0) - qml.expval(qml.PauliX(0)) + a = tf.constant(0.1) + p = tf.Variable([0.1, 0.2, 0.3]) - qtape = qml.tape.QuantumScript.from_queue(q_qtape) - res = execute([qtape], dev, **execute_kwargs)[0] + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost_fn(a, p) expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( tf.cos(p[2]) * tf.sin(p[1]) + tf.cos(p[0]) * tf.cos(p[1]) * tf.sin(p[2]) ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - res = tape.jacobian(res, p) - expected = np.array( + res = tape.jacobian(cost_res, p, experimental_use_pfor=not device_vjp) + expected = tf.convert_to_tensor( [ tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), tf.cos(p[1]) * tf.cos(p[2]) * tf.sin(a) @@ -635,121 +515,119 @@ def decomposition(self): * (tf.cos(p[0]) * tf.cos(p[1]) * tf.cos(p[2]) - tf.sin(p[1]) * tf.sin(p[2])), ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_probability_differentiation(self, execute_kwargs, tol): + def test_probability_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") + def cost(x, y): + ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.probs(wires=0), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") + + device_vjp = execute_kwargs.get("device_vjp", False) - dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=[0]) - qml.probs(wires=[1]) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - res = qml.math.stack(res) + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost(x, y) - expected = np.array( + expected = tf.convert_to_tensor( [ - [tf.cos(x / 2) ** 2, tf.sin(x / 2) ** 2], - [(1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2], + [ + tf.cos(x / 2) ** 2, + tf.sin(x / 2) ** 2, + (1 + tf.cos(x) * tf.cos(y)) / 2, + (1 - tf.cos(x) * tf.cos(y)) / 2, + ], ] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - res = t.jacobian(res, [x, y]) - expected = np.array( - [ + if ( + execute_kwargs.get("interface", "") == "tf-autograph" + and execute_kwargs.get("gradient_fn", "") == "adjoint" + ): + with pytest.raises(tf.errors.UnimplementedError): + tape.jacobian(cost_res, [x, y]) + return + res = tape.jacobian(cost_res, [x, y], experimental_use_pfor=not device_vjp) + assert isinstance(res, list) and len(res) == 2 + assert res[0].shape == (4,) + assert res[1].shape == (4,) + + expected = ( + tf.convert_to_tensor( [ - [-tf.sin(x) / 2, tf.sin(x) / 2], - [-tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], - ], + [ + -tf.sin(x) / 2, + tf.sin(x) / 2, + -tf.sin(x) * tf.cos(y) / 2, + tf.sin(x) * tf.cos(y) / 2, + ], + ] + ), + tf.convert_to_tensor( [ - [0, 0], - [-tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], - ], - ] + [0, 0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], + ] + ), ) - assert np.allclose(res, expected, atol=tol, rtol=0) - def test_ragged_differentiation(self, execute_kwargs, tol): + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) + + def test_ragged_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") - dev = qml.device("default.qubit.legacy", wires=2) + device_vjp = execute_kwargs.get("device_vjp", False) + + def cost(x, y): + ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") + x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) - with tf.GradientTape() as t: - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.probs(wires=[1]) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - res = tf.experimental.numpy.hstack(res) + with tf.GradientTape(persistent=device_vjp) as tape: + cost_res = cost(x, y) - expected = np.array( + expected = tf.convert_to_tensor( [tf.cos(x), (1 + tf.cos(x) * tf.cos(y)) / 2, (1 - tf.cos(x) * tf.cos(y)) / 2] ) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(cost_res, expected, atol=atol_for_shots(shots), rtol=0) - res = t.jacobian(res, [x, y]) - expected = np.array( - [ - [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], - [0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2], - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_sampling(self, execute_kwargs): - """Test sampling works as expected""" if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True + execute_kwargs.get("interface", "") == "tf-autograph" + and execute_kwargs.get("gradient_fn", "") == "adjoint" ): - pytest.skip("Adjoint differentiation does not support samples") - - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - - with tf.GradientTape(): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(tf.Variable(0.1), wires=0) - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.sample(qml.PauliZ(0)) - qml.sample(qml.PauliX(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute([tape], dev, **execute_kwargs)[0] - res = qml.math.stack(res) - - assert res.shape == (2, 10) - assert isinstance(res, tf.Tensor) + with pytest.raises(tf.errors.UnimplementedError): + tape.jacobian(cost_res, [x, y]) + return + res = tape.jacobian(cost_res, [x, y], experimental_use_pfor=not device_vjp) + assert isinstance(res, list) and len(res) == 2 + assert res[0].shape == (3,) + assert res[1].shape == (3,) + + expected = ( + tf.convert_to_tensor( + [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.sin(x) * tf.cos(y) / 2] + ), + tf.convert_to_tensor([0, -tf.cos(x) * tf.sin(y) / 2, tf.cos(x) * tf.sin(y) / 2]), + ) + assert np.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert np.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) -@pytest.mark.parametrize("interface", ["auto", "tf"]) class TestHigherOrderDerivatives: - """Test that the TensorFlow execute function can be differentiated""" + """Test that the tensorflow execute function can be differentiated""" - @pytest.mark.slow @pytest.mark.parametrize( "params", [ @@ -758,196 +636,96 @@ class TestHigherOrderDerivatives: tf.Variable([-2.0, 0], dtype=tf.float64), ], ) - def test_parameter_shift_hessian(self, params, tol, interface): + def test_parameter_shift_hessian(self, params, tol): """Tests that the output of the parameter-shift transform can be differentiated using tensorflow, yielding second derivatives.""" - dev = qml.device("default.qubit.tf", wires=2) - x, y = params * 1.0 - - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(params[0], wires=[0]) - qml.RY(params[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(params[0], wires=0) - qml.RY(params[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = execute( - [tape1, tape2], dev, gradient_fn=param_shift, interface=interface, max_diff=2 - ) - res = result[0] + result[1][0] + dev = DefaultQubit() - expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) + def cost_fn(x): + ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - grad = t1.gradient(res, params) - expected = np.array( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert np.allclose(grad, expected, atol=tol, rtol=0) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) + return result[0] + result[1][0] - hess = t2.jacobian(grad, params) - expected = np.array( - [ - [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], - ] - ) - assert np.allclose(hess, expected, atol=tol, rtol=0) - - def test_hessian_vector_valued(self, tol, interface): - """Test hessian calculation of a vector valued QNode""" - dev = qml.device("default.qubit.tf", wires=1) - params = tf.Variable([0.543, -0.654], dtype=tf.float64) - - with tf.GradientTape() as t2: - with tf.GradientTape(persistent=True) as t1: - with qml.queuing.AnnotatedQueue() as q: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute( - [tape], dev, gradient_fn=param_shift, interface=interface, max_diff=2 - )[0] - res = tf.stack(res) - - g = t1.jacobian(res, params, experimental_use_pfor=False) - - hess = t2.jacobian(g, params) - - a, b = params * 1.0 + with tf.GradientTape() as jac_tape: + with tf.GradientTape() as grad_tape: + res = cost_fn(params) + grad = grad_tape.gradient(res, params) + hess = jac_tape.jacobian(grad, params) - expected_res = [ - 0.5 + 0.5 * tf.cos(a) * tf.cos(b), - 0.5 - 0.5 * tf.cos(a) * tf.cos(b), - ] - assert np.allclose(res, expected_res, atol=tol, rtol=0) + x, y = params + expected = 0.5 * (3 + tf.cos(x) ** 2 * tf.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) - expected_g = [ - [-0.5 * tf.sin(a) * tf.cos(b), -0.5 * tf.cos(a) * tf.sin(b)], - [0.5 * tf.sin(a) * tf.cos(b), 0.5 * tf.cos(a) * tf.sin(b)], - ] - assert np.allclose(g, expected_g, atol=tol, rtol=0) + expected = tf.convert_to_tensor( + [-tf.cos(x) * tf.cos(2 * y) * tf.sin(x), -tf.cos(x) ** 2 * tf.sin(2 * y)] + ) + assert np.allclose(grad, expected, atol=tol, rtol=0) - expected_hess = [ - [ - [-0.5 * tf.cos(a) * tf.cos(b), 0.5 * tf.sin(a) * tf.sin(b)], - [0.5 * tf.sin(a) * tf.sin(b), -0.5 * tf.cos(a) * tf.cos(b)], - ], + expected = tf.convert_to_tensor( [ - [0.5 * tf.cos(a) * tf.cos(b), -0.5 * tf.sin(a) * tf.sin(b)], - [-0.5 * tf.sin(a) * tf.sin(b), 0.5 * tf.cos(a) * tf.cos(b)], - ], - ] - - np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) - - def test_adjoint_hessian(self, interface): - """Since the adjoint hessian is not a differentiable transform, - higher-order derivatives are not supported.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = tf.Variable([0.543, -0.654], dtype=tf.float64) - - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - with qml.queuing.AnnotatedQueue() as q: - qml.RX(params[0], wires=[0]) - qml.RY(params[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - res = execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - interface=interface, - )[0] - - grad = t1.gradient(res, params) - assert grad is not None - assert grad.dtype == tf.float64 - assert grad.shape == params.shape - - hess = t2.jacobian(grad, params) - assert hess is None + [-tf.cos(2 * x) * tf.cos(2 * y), tf.sin(2 * x) * tf.sin(2 * y)], + [tf.sin(2 * x) * tf.sin(2 * y), -2 * tf.cos(x) ** 2 * tf.cos(2 * y)], + ] + ) + assert np.allclose(hess, expected, atol=tol, rtol=0) - def test_max_diff(self, tol, interface): + def test_max_diff(self, tol): """Test that setting the max_diff parameter blocks higher-order derivatives""" - dev = qml.device("default.qubit.tf", wires=2) - params = tf.Variable([0.543, -0.654], dtype=tf.float64) - x, y = params * 1.0 - - with tf.GradientTape() as t2: - with tf.GradientTape() as t1: - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(params[0], wires=[0]) - qml.RY(params[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(params[0], wires=0) - qml.RY(params[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - result = execute( - [tape1, tape2], dev, gradient_fn=param_shift, max_diff=1, interface=interface - ) - res = result[0] + result[1][0] + dev = DefaultQubit() + params = tf.Variable([0.543, -0.654]) - expected = 0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y)) - assert np.allclose(res, expected, atol=tol, rtol=0) + def cost_fn(x): + ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - grad = t1.gradient(res, params) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - expected = np.array( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert np.allclose(grad, expected, atol=tol, rtol=0) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) + return result[0] + result[1][0] - hess = t2.jacobian(grad, params) - assert hess is None + with tf.GradientTape() as jac_tape: + with tf.GradientTape() as grad_tape: + res = cost_fn(params) + grad = grad_tape.gradient(res, params) + hess = jac_tape.gradient(grad, params) + x, y = params + expected = 0.5 * (3 + tf.cos(x) ** 2 * tf.cos(2 * y)) + assert np.allclose(res, expected, atol=tol, rtol=0) -execute_kwargs_hamiltonian = [ - {"gradient_fn": param_shift, "interface": "tensorflow"}, - {"gradient_fn": finite_diff, "interface": "tensorflow"}, - {"gradient_fn": param_shift, "interface": "auto"}, - {"gradient_fn": finite_diff, "interface": "auto"}, -] + expected = tf.convert_to_tensor( + [-tf.cos(x) * tf.cos(2 * y) * tf.sin(x), -tf.cos(x) ** 2 * tf.sin(2 * y)] + ) + assert np.allclose(grad, expected, atol=tol, rtol=0) + assert hess is None -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +@pytest.mark.usefixtures("use_legacy_and_new_opmath") class TestHamiltonianWorkflows: """Test that tapes ending with expectations of Hamiltonians provide correct results and gradients""" @pytest.fixture - def cost_fn(self, execute_kwargs): + def cost_fn(self, execute_kwargs, shots, device): """Cost function for gradient tests""" - def _cost_fn(weights, coeffs1, coeffs2, dev=None): + def _cost_fn(weights, coeffs1, coeffs2): obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] H1 = qml.Hamiltonian(coeffs1, obs1) + if qml.operation.active_new_opmath(): + H1 = qml.pauli.pauli_sentence(H1).operation() obs2 = [qml.PauliZ(0)] H2 = qml.Hamiltonian(coeffs2, obs2) + if qml.operation.active_new_opmath(): + H2 = qml.pauli.pauli_sentence(H2).operation() with qml.queuing.AnnotatedQueue() as q: qml.RX(weights[0], wires=0) @@ -956,73 +734,111 @@ def _cost_fn(weights, coeffs1, coeffs2, dev=None): qml.expval(H1) qml.expval(H2) - tape = qml.tape.QuantumScript.from_queue(q) - return tf.stack(execute([tape], dev, **execute_kwargs)[0]) + tape = qml.tape.QuantumScript.from_queue(q, shots=shots) + return qml.math.hstack(execute([tape], device, **execute_kwargs)[0], like="tensorflow") return _cost_fn @staticmethod def cost_fn_expected(weights, coeffs1, coeffs2): """Analytic value of cost_fn above""" - a, b, c = coeffs1.numpy() - d = coeffs2.numpy()[0] - x, y = weights.numpy() - return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return [-c * tf.sin(x) * tf.sin(y) + tf.cos(x) * (a + b * tf.sin(y)), d * tf.cos(x)] @staticmethod - def cost_fn_jacobian_expected(weights, coeffs1, coeffs2): + def cost_fn_jacobian(weights, coeffs1, coeffs2): """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1.numpy() - d = coeffs2.numpy()[0] - x, y = weights.numpy() - return np.array( + a, b, c = coeffs1 + d = coeffs2[0] + x, y = weights + return tf.convert_to_tensor( [ [ - -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), - b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), - np.cos(x), - np.cos(x) * np.sin(y), - -(np.sin(x) * np.sin(y)), + -c * tf.cos(x) * tf.sin(y) - tf.sin(x) * (a + b * tf.sin(y)), + b * tf.cos(x) * tf.cos(y) - c * tf.cos(y) * tf.sin(x), + tf.cos(x), + tf.cos(x) * tf.sin(y), + -(tf.sin(x) * tf.sin(y)), 0, ], - [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + [-d * tf.sin(x), 0, 0, 0, 0, tf.cos(x)], ] ) - def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): - # pylint: disable=unused-argument - weights = tf.Variable([0.4, 0.5], dtype=tf.float64) + def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with no trainable parameters.""" + + if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): + pytest.skip("adjoint differentiation does not suppport hamiltonians.") + + device_vjp = execute_kwargs.get("device_vjp", False) + coeffs1 = tf.constant([0.1, 0.2, 0.3], dtype=tf.float64) coeffs2 = tf.constant([0.7], dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) + weights = tf.Variable([0.4, 0.5], dtype=tf.float64) - with tf.GradientTape() as tape: - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + with tf.GradientTape(persistent=device_vjp) as tape: + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = tape.jacobian(res, [weights, coeffs1, coeffs2]) - expected = self.cost_fn_jacobian_expected(weights, coeffs1, coeffs2) - assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0) - assert res[1] is None - assert res[2] is None + jac = tape.jacobian(res, [weights], experimental_use_pfor=not device_vjp) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] + assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0) + + def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, shots): + """Test hamiltonian with trainable parameters.""" + if execute_kwargs["gradient_fn"] == "adjoint": + pytest.skip("trainable hamiltonians not supported with adjoint") + if qml.operation.active_new_opmath(): + pytest.skip("parameter shift derivatives do not yet support sums.") - def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): - # pylint: disable=unused-argument - weights = tf.Variable([0.4, 0.5], dtype=tf.float64) coeffs1 = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) coeffs2 = tf.Variable([0.7], dtype=tf.float64) - dev = qml.device("default.qubit.legacy", wires=2) + weights = tf.Variable([0.4, 0.5], dtype=tf.float64) with tf.GradientTape() as tape: - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res, expected, atol=tol, rtol=0) + assert np.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + jac = qml.math.hstack(tape.jacobian(res, [weights, coeffs1, coeffs2]), like="tensorflow") + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) + assert np.allclose(jac, expected, atol=atol_for_shots(shots), rtol=0) + + +@pytest.mark.parametrize("diff_method", ("adjoint", "parameter-shift")) +def test_device_returns_float32(diff_method): + """Test that if the device returns float32, the derivative succeeds.""" + + def _to_float32(results): + if isinstance(results, (list, tuple)): + return tuple(_to_float32(r) for r in results) + return np.array(results, dtype=np.float32) + + class Float32Dev(qml.devices.DefaultQubit): + def execute(self, circuits, execution_config=qml.devices.DefaultExecutionConfig): + results = super().execute(circuits, execution_config) + return _to_float32(results) + + dev = Float32Dev() + + @qml.qnode(dev, diff_method=diff_method) + def circuit(x): + qml.RX(tf.cos(x), wires=0) + return qml.expval(qml.Z(0)) + + x = tf.Variable(0.1, dtype=tf.float64) + + with tf.GradientTape() as tape: + y = circuit(x) + + assert qml.math.allclose(y, np.cos(np.cos(0.1))) - res = tape.jacobian(res, [weights, coeffs1, coeffs2]) - expected = self.cost_fn_jacobian_expected(weights, coeffs1, coeffs2) - assert np.allclose(res[0], expected[:, :2], atol=tol, rtol=0) - assert np.allclose(res[1], expected[:, 2:5], atol=tol, rtol=0) - assert np.allclose(res[2], expected[:, 5:], atol=tol, rtol=0) + g = tape.gradient(y, x) + expected_g = np.sin(np.cos(0.1)) * np.sin(0.1) + assert qml.math.allclose(g, expected_g) diff --git a/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py b/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py index ab3147dfc95..695438310c1 100644 --- a/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py +++ b/tests/interfaces/test_tensorflow_autograph_qnode_shot_vector.py @@ -1,4 +1,4 @@ -# Copyright 2022 Xanadu Quantum Technologies Inc. +# Copyright 2023 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. @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the TF interface with shot vectors and with a QNode""" -# pylint: disable=too-many-arguments,too-few-public-methods +# pylint: disable=too-many-arguments,too-few-public-methods,unexpected-keyword-arg,redefined-outer-name import pytest import pennylane as qml from pennylane import numpy as np from pennylane import qnode +from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf @@ -26,22 +27,35 @@ shots_and_num_copies = [((1, (5, 2), 10), 4)] shots_and_num_copies_hess = [((10, (5, 1)), 2)] -spsa_kwargs = {"h": 10e-2, "num_directions": 30, "sampler_rng": np.random.default_rng(42)} +kwargs = { + "finite-diff": {"h": 10e-2}, + "parameter-shift": {}, + "spsa": {"h": 10e-2, "num_directions": 20}, +} + qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], - ["default.qubit.legacy", "parameter-shift", {}], - ["default.qubit.legacy", "spsa", spsa_kwargs], + [DefaultQubit(), "finite-diff"], + [DefaultQubit(), "parameter-shift"], + [DefaultQubit(), "spsa"], ] TOLS = { "finite-diff": 0.3, "parameter-shift": 1e-2, - "spsa": 0.3, + "spsa": 0.5, } +@pytest.fixture +def gradient_kwargs(request): + diff_method = request.node.funcargs["diff_method"] + return kwargs[diff_method] | ( + {"sampler_rng": np.random.default_rng(42)} if diff_method == "spsa" else {} + ) + + @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) -@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -50,14 +64,13 @@ class TestReturnWithShotVectors: """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" def test_jac_single_measurement_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and one param, the gradient is a float.""" - dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.expval(qml.PauliZ(0)) @@ -65,7 +78,7 @@ def circuit(a): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -74,14 +87,13 @@ def circuit(a): assert jac.shape == (num_copies,) def test_jac_single_measurement_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b): + def circuit(a, b, **_): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)) @@ -90,7 +102,7 @@ def circuit(a, b): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -102,14 +114,13 @@ def circuit(a, b): assert j.shape == (num_copies,) def test_jacobian_single_measurement_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - dev = qml.device(dev_name, wires=1, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) @@ -117,7 +128,7 @@ def circuit(a): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -126,15 +137,14 @@ def circuit(a): assert jac.shape == (num_copies, 2) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.probs(wires=[0, 1]) @@ -142,7 +152,7 @@ def circuit(a): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -151,15 +161,14 @@ def circuit(a): assert jac.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b): + def circuit(a, b, **_): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.probs(wires=[0, 1]) @@ -168,7 +177,7 @@ def circuit(a, b): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -180,15 +189,14 @@ def circuit(a, b): assert j.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) @@ -196,7 +204,7 @@ def circuit(a): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -205,24 +213,22 @@ def circuit(a): assert jac.shape == (num_copies, 4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The gradient of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = tf.Variable(1.5, dtype=tf.float64) par_1 = tf.Variable(0.7, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, max_diff=1, **gradient_kwargs) - def circuit(x, y): + def circuit(x, y, **_): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) with tf.GradientTape() as tape: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, (par_0, par_1)) @@ -234,14 +240,13 @@ def circuit(x, y): assert j.shape == (num_copies, 2) def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) qml.RY(a[2], wires=0) @@ -250,7 +255,7 @@ def circuit(a): a = tf.Variable([0.7, 0.9, 1.1], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, a) @@ -259,14 +264,13 @@ def circuit(a): assert jac.shape == (num_copies, 2, 3) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a single params return an array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a, wires=0) qml.RX(0.7, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -274,7 +278,7 @@ def circuit(a): a = tf.Variable(1.5, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -283,14 +287,13 @@ def circuit(a): assert jac.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a, b): + def circuit(a, b, **_): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -299,7 +302,7 @@ def circuit(a, b): b = tf.Variable(0.7, dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, (a, b)) @@ -311,14 +314,13 @@ def circuit(a, b): assert j.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(a): + def circuit(a, **_): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) @@ -326,7 +328,7 @@ def circuit(a): a = tf.Variable([1.5, 0.7], dtype=tf.float64) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -337,7 +339,7 @@ def circuit(a): @pytest.mark.slow @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) -@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -346,7 +348,7 @@ class TestReturnShotVectorHessian: """Class to test the shape of the Hessian with different return types and shot vectors.""" def test_hessian_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """The hessian of a single measurement with multiple params return a tuple of arrays.""" @@ -354,14 +356,12 @@ def test_hessian_expval_multiple_params( # TODO: Find out why. pytest.skip("SPSA gradient does not support this particular test case") - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = tf.Variable(1.5, dtype=tf.float64) par_1 = tf.Variable(0.7, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, max_diff=2, **gradient_kwargs) - def circuit(x, y): + def circuit(x, y, **_): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) @@ -369,7 +369,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack(res) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -384,11 +384,11 @@ def circuit(x, y): assert h.shape == (2, num_copies) -shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)] +shots_and_num_copies = [((30000, 28000, 26000), 3), ((30000, (28000, 2)), 3)] @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) -@pytest.mark.parametrize("dev_name,diff_method,gradient_kwargs", qubit_device_and_diff_method) +@pytest.mark.parametrize("dev,diff_method", qubit_device_and_diff_method) @pytest.mark.parametrize( "decorator,interface", [(tf.function, "tf"), (lambda x: x, "tf-autograph")], @@ -397,24 +397,23 @@ class TestReturnShotVectorIntegration: """Tests for the integration of shots with the TF interface.""" def test_single_expectation_value( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """Tests correct output shape and evaluation for a tape with a single expval output""" - dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(x, y): + def circuit(x, y, **_): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) with tf.GradientTape() as tape: - res = circuit(x, y) + res = circuit(x, y, shots=shots) res = qml.math.stack(res) all_res = tape.jacobian(res, (x, y)) @@ -431,24 +430,23 @@ def circuit(x, y): assert np.allclose(res, exp, atol=tol, rtol=0) def test_prob_expectation_values( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, decorator, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, decorator, interface ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - dev = qml.device(dev_name, wires=2, shots=shots) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) - def circuit(x, y): + def circuit(x, y, **_): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) with tf.GradientTape() as tape: - res = circuit(x, y) + res = circuit(x, y, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) all_res = tape.jacobian(res, (x, y)) diff --git a/tests/interfaces/test_tensorflow_qnode.py b/tests/interfaces/test_tensorflow_qnode.py index 5b03f9da10f..e3c4597ae06 100644 --- a/tests/interfaces/test_tensorflow_qnode.py +++ b/tests/interfaces/test_tensorflow_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2018-2020 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 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. @@ -13,26 +13,31 @@ # limitations under the License. """Integration tests for using the TensorFlow interface with a QNode""" import numpy as np + +# pylint: disable=too-many-arguments,too-few-public-methods,comparison-with-callable, use-implicit-booleaness-not-comparison import pytest import pennylane as qml from pennylane import qnode - -# pylint: disable=too-many-arguments,too-few-public-methods, use-dict-literal, use-implicit-booleaness-not-comparison - +from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf tf = pytest.importorskip("tensorflow") - +# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], + [DefaultQubit(), "finite-diff", False, False], + [DefaultQubit(), "parameter-shift", False, False], + [DefaultQubit(), "backprop", True, False], + [DefaultQubit(), "adjoint", True, False], + [DefaultQubit(), "adjoint", False, False], + [DefaultQubit(), "adjoint", False, True], + [DefaultQubit(), "spsa", False, False], + [DefaultQubit(), "hadamard", False, False], + [qml.device("lightning.qubit", wires=4), "adjoint", False, True], + [qml.device("lightning.qubit", wires=4), "adjoint", False, False], + [qml.device("lightning.qubit", wires=4), "adjoint", True, True], + [qml.device("lightning.qubit", wires=4), "adjoint", True, False], ] TOL_FOR_SPSA = 1.0 @@ -45,25 +50,25 @@ @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQNode: """Test that using the QNode with TensorFlow integrates with the PennyLane stack""" - def test_execution_with_interface(self, dev_name, diff_method, grad_on_execution, interface): + def test_execution_with_interface( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -84,7 +89,7 @@ def circuit(a): # with the interface, the tape returns tensorflow tensors assert isinstance(res, tf.Tensor) - assert res.shape == tuple() + assert res.shape == () # the tape is able to deduce trainable parameters assert circuit.qtape.trainable_params == [0] @@ -92,23 +97,20 @@ def circuit(a): # gradients should work grad = tape.gradient(res, a) assert isinstance(grad, tf.Tensor) - assert grad.shape == tuple() + assert grad.shape == () - def test_interface_swap(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_interface_swap(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test that the TF interface can be applied to a QNode with a pre-existing interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface="autograd", diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface="autograd", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -135,51 +137,44 @@ def circuit(a): assert np.allclose(res1, res2, atol=tol, rtol=0) assert np.allclose(grad1, grad2, atol=tol, rtol=0) - def test_drawing(self, dev_name, diff_method, grad_on_execution, interface): + def test_drawing(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test circuit drawing when using the TF interface""" x = tf.Variable(0.1, dtype=tf.float64) y = tf.Variable([0.2, 0.3], dtype=tf.float64) z = tf.Variable(0.4, dtype=tf.float64) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(p1, p2=y, **kwargs): qml.RX(p1, wires=0) qml.RY(p2[0] * p2[1], wires=1) qml.RX(kwargs["p3"], wires=0) qml.CNOT(wires=[0, 1]) - return qml.state() + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) result = qml.draw(circuit)(p1=x, p3=z) - expected = "0: ──RX(0.10)──RX(0.40)─╭●── State\n1: ──RY(0.06)───────────╰X── State" + expected = "0: ──RX(0.10)──RX(0.40)─╭●── \n1: ──RY(0.06)───────────╰X── " assert result == expected - def test_jacobian(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_jacobian(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test jacobian calculation""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) - + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - a = tf.Variable(0.1, dtype=tf.float64) b = tf.Variable(0.2, dtype=tf.float64) @@ -190,7 +185,7 @@ def circuit(a, b): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.stack(res) @@ -202,61 +197,17 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [a, b]) + res = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) expected = [[-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)]] assert np.allclose(res, expected, atol=tol, rtol=0) - def test_jacobian_dtype(self, dev_name, diff_method, grad_on_execution, interface): - """Test calculating the jacobian with a different datatype""" - if diff_method == "backprop": - pytest.skip("Test does not support backprop") - - a = tf.Variable(0.1, dtype=tf.float32) - b = tf.Variable(0.2, dtype=tf.float32) - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, r_dtype=np.float32) - - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) - def circuit(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] - - with tf.GradientTape() as tape: - res = circuit(a, b) - res = tf.stack(res) - - assert circuit.qtape.trainable_params == [0, 1] - - assert isinstance(res, tf.Tensor) - assert res.shape == (2,) - assert res.dtype is tf.float32 - - res = tape.jacobian(res, [a, b]) - assert [r.dtype is tf.float32 for r in res] - - def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interface): + def test_jacobian_options(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test setting finite-difference jacobian options""" if diff_method not in {"finite-diff", "spsa"}: pytest.skip("Test only works with finite diff and spsa.") a = tf.Variable([0.1, 0.2]) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, interface=interface, @@ -264,18 +215,21 @@ def test_jacobian_options(self, dev_name, diff_method, grad_on_execution, interf approx_order=2, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) - tape.jacobian(res, a) + tape.jacobian(res, a, experimental_use_pfor=not device_vjp) - def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_changing_trainability( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Test changing the trainability of parameters changes the number of differentiation requests made""" if diff_method in ["backprop", "adjoint", "spsa"]: @@ -284,21 +238,16 @@ def test_changing_trainability(self, dev_name, diff_method, grad_on_execution, t a = tf.Variable(0.1, dtype=tf.float64) b = tf.Variable(0.2, dtype=tf.float64) - num_wires = 2 - diff_kwargs = {} - if diff_method == "hadamard": - num_wires = 3 - elif diff_method == "finite-diff": + if diff_method == "finite-diff": diff_kwargs = {"approx_order": 2, "strategy": "center"} - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, **diff_kwargs, ) def circuit(a, b): @@ -307,7 +256,7 @@ def circuit(a, b): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.stack(res) @@ -317,7 +266,7 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - jac = tape.jacobian(res, [a, b]) + jac = tape.jacobian(res, [a, b], experimental_use_pfor=not device_vjp) expected = [ [-tf.sin(a), tf.sin(a) * tf.sin(b)], [0, -tf.cos(a) * tf.cos(b)], @@ -328,7 +277,7 @@ def circuit(a, b): a = tf.Variable(0.54, dtype=tf.float64) b = tf.constant(0.8, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.stack(res) @@ -338,25 +287,22 @@ def circuit(a, b): expected = [tf.cos(a), -tf.cos(a) * tf.sin(b)] assert np.allclose(res, expected, atol=tol, rtol=0) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) expected = [-tf.sin(a), tf.sin(a) * tf.sin(b)] assert np.allclose(jac, expected, atol=tol, rtol=0) - def test_classical_processing(self, dev_name, diff_method, grad_on_execution, interface): + def test_classical_processing(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test classical processing within the quantum tape""" a = tf.Variable(0.1, dtype=tf.float64) b = tf.constant(0.2, dtype=tf.float64) c = tf.Variable(0.3, dtype=tf.float64) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(x, y, z): qml.RY(x * z, wires=0) @@ -364,30 +310,30 @@ def circuit(x, y, z): qml.RX(z + z**2 + tf.sin(a), wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b, c) if diff_method == "finite-diff": assert circuit.qtape.trainable_params == [0, 2] assert circuit.qtape.get_parameters() == [a * c, c + c**2 + tf.sin(a)] - res = tape.jacobian(res, [a, b, c]) + res = tape.jacobian(res, [a, b, c], experimental_use_pfor=not device_vjp) assert isinstance(res[0], tf.Tensor) assert res[1] is None assert isinstance(res[2], tf.Tensor) - def test_no_trainable_parameters(self, dev_name, diff_method, grad_on_execution, interface): + def test_no_trainable_parameters( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): """Test evaluation if there are no trainable parameters""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -398,7 +344,7 @@ def circuit(a, b): a = 0.1 b = tf.constant(0.2, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.stack(res) @@ -409,31 +355,30 @@ def circuit(a, b): assert isinstance(res, tf.Tensor) # can't take the gradient with respect to "a" since it's a Python scalar - grad = tape.jacobian(res, b) + grad = tape.jacobian(res, b, experimental_use_pfor=not device_vjp) assert grad is None @pytest.mark.parametrize("U", [tf.constant([[0, 1], [1, 0]]), np.array([[0, 1], [1, 0]])]) - def test_matrix_parameter(self, dev_name, diff_method, grad_on_execution, U, tol, interface): + def test_matrix_parameter( + self, dev, diff_method, grad_on_execution, device_vjp, U, tol, interface + ): """Test that the TF interface works correctly with a matrix parameter""" a = tf.Variable(0.1, dtype=tf.float64) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) qml.RY(a, wires=0) return qml.expval(qml.PauliZ(0)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(U, a) if diff_method == "finite-diff": @@ -441,18 +386,23 @@ def circuit(U, a): assert np.allclose(res, -tf.cos(a), atol=tol, rtol=0) - res = tape.jacobian(res, a) + res = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert np.allclose(res, tf.sin(a), atol=tol, rtol=0) - def test_differentiable_expand(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_differentiable_expand( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Test that operation and nested tapes expansion is differentiable""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } if diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA class U3(qml.U3): @@ -464,13 +414,6 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - a = np.array(0.1) p = tf.Variable([0.1, 0.2, 0.3], dtype=tf.float64) @@ -480,7 +423,7 @@ def circuit(a, p): U3(p[0], p[1], p[2], wires=0) return qml.expval(qml.PauliX(0)) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, p) expected = tf.cos(a) * tf.cos(p[1]) * tf.sin(p[0]) + tf.sin(a) * ( @@ -488,7 +431,7 @@ def circuit(a, p): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, p) + res = tape.jacobian(res, p, experimental_use_pfor=not device_vjp) expected = np.array( [ tf.cos(p[1]) * (tf.cos(a) * tf.cos(p[0]) - tf.sin(a) * tf.sin(p[0]) * tf.sin(p[2])), @@ -507,9 +450,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and differentiates it.""" - def test_changing_shots(self, mocker, tol, interface): + def test_changing_shots(self, interface): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -518,31 +461,21 @@ def circuit(weights): qml.RY(weights[0], wires=0) qml.RX(weights[1], wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(weights) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + circuit(weights) # execute with shots=100 - circuit(weights, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(weights) - assert np.allclose(res, -np.cos(a) * np.sin(b), atol=tol, rtol=0) - spy.assert_called_once() + res = circuit(weights, shots=100) # pylint: disable=unexpected-keyword-arg + assert res.shape == (100, 2) def test_gradient_integration(self, interface): """Test that temporarily setting the shots works for gradient computations""" # pylint: disable=unexpected-keyword-arg - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -557,7 +490,6 @@ def circuit(weights): res = circuit(weights, shots=[10000, 10000, 10000]) res = tf.transpose(tf.stack(res)) - assert dev.shots is None assert len(res) == 3 jacobian = tape.jacobian(res, weights) @@ -568,7 +500,7 @@ def test_multiple_gradient_integration(self, tol, interface): """Test that temporarily setting the shots works for gradient computations, even if the QNode has been re-evaluated with a different number of shots in the meantime.""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = [0.543, -0.654] weights = tf.Variable([a, b], dtype=tf.float64) @@ -582,7 +514,7 @@ def circuit(weights): with tf.GradientTape() as tape: res1 = circuit(weights) - assert qml.math.shape(res1) == tuple() + assert qml.math.shape(res1) == () res2 = circuit(weights, shots=[(1, 1000)]) # pylint: disable=unexpected-keyword-arg assert qml.math.shape(res2) == (1000,) @@ -593,7 +525,7 @@ def circuit(weights): def test_update_diff_method(self, mocker, interface): """Test that temporarily setting the shots updates the diff method""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100) + dev = DefaultQubit() weights = tf.Variable([0.543, -0.654], dtype=tf.float64) spy = mocker.spy(qml, "execute") @@ -605,115 +537,44 @@ def circuit(weights): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) + circuit(weights, shots=100) # pylint:disable=unexpected-keyword-arg # since we are using finite shots, parameter-shift will # be chosen - circuit(weights) - assert circuit.gradient_fn is qml.gradients.param_shift - - circuit(weights) + assert circuit.gradient_fn == qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - assert circuit.gradient_fn is qml.gradients.param_shift - - # if we set the shots to None, backprop can now be used - circuit(weights, shots=None) # pylint: disable=unexpected-keyword-arg - assert spy.call_args[1]["gradient_fn"] == "backprop" - assert circuit.gradient_fn == "backprop" + # if we use the default shots value of None, backprop can now be used circuit(weights) - assert circuit.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - - -@pytest.mark.parametrize("interface", ["auto", "tf"]) -class TestAdjoint: - """Specific integration tests for the adjoint method""" - - def test_reuse_state(self, mocker, interface): - """Tests that the TF interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=2) - - @qnode(dev, diff_method="adjoint", interface=interface) - def circ(x): - qml.RX(x[0], wires=0) - qml.RY(x[1], wires=1) - qml.CNOT(wires=(0, 1)) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - weights = tf.Variable([0.1, 0.2], dtype=tf.float64) - x, y = 1.0 * weights - - with tf.GradientTape() as tape: - res = tf.reduce_sum(circ(weights)) - - grad = tape.gradient(res, weights) - expected_grad = [-tf.sin(x), tf.cos(y)] - - assert np.allclose(grad, expected_grad) - assert circ.device.num_executions == 1 - spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) - - def test_reuse_state_multiple_evals(self, mocker, tol, interface): - """Tests that the TF interface reuses the device state for adjoint differentiation, - even where there are intermediate evaluations.""" - dev = qml.device("default.qubit.legacy", wires=2) - - x_val = 0.543 - y_val = -0.654 - x = tf.Variable(x_val, dtype=tf.float64) - y = tf.Variable(y_val, dtype=tf.float64) - - @qnode(dev, diff_method="adjoint", interface=interface) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - with tf.GradientTape() as tape: - res1 = circuit(x, y) - - assert np.allclose(res1, np.cos(x_val), atol=tol, rtol=0) - - # intermediate evaluation with different values - circuit(tf.math.tan(x), tf.math.cosh(y)) - - # the adjoint method will continue to compute the correct derivative - grad = tape.gradient(res1, x) - assert np.allclose(grad, -np.sin(x_val), atol=tol, rtol=0) - assert dev.num_executions == 2 - spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) + assert circuit.gradient_fn == "backprop" + assert spy.call_args[1]["gradient_fn"] == "backprop" @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" def test_probability_differentiation( - self, dev_name, diff_method, grad_on_execution, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface ): """Tests correct output shape and evaluation for a tape with multiple probs outputs""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - num_wires = 2 + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 + tol = TOL_FOR_SPSA x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @@ -725,7 +586,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.probs(wires=[0]), qml.probs(wires=[1]) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(x, y) res = tf.stack(res) @@ -737,7 +598,7 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [x, y]) + res = tape.jacobian(res, [x, y], experimental_use_pfor=not device_vjp) expected = np.array( [ [ @@ -752,24 +613,25 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_ragged_differentiation(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_ragged_differentiation( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - num_wires = 2 + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } + if diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 + tol = TOL_FOR_SPSA x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @@ -781,7 +643,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1]) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(x, y) res = tf.experimental.numpy.hstack(res) @@ -794,7 +656,7 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - res = tape.jacobian(res, [x, y]) + res = tape.jacobian(res, [x, y], experimental_use_pfor=not device_vjp) expected = np.array( [ [-tf.sin(x), -tf.sin(x) * tf.cos(y) / 2, tf.cos(y) * tf.sin(x) / 2], @@ -803,24 +665,20 @@ def circuit(x, y): ) assert np.allclose(res, expected, atol=tol, rtol=0) - def test_second_derivative(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_second_derivative( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Test second derivative calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -850,24 +708,18 @@ def circuit(x): ] assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - def test_hessian(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_hessian(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test hessian calculation of a scalar valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -896,24 +748,20 @@ def circuit(x): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_hessian_vector_valued( + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface + ): """Test hessian calculation of a vector valued QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -956,25 +804,19 @@ def circuit(x): np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) def test_hessian_vector_valued_postprocessing( - self, dev_name, diff_method, grad_on_execution, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, tol, interface ): """Test hessian calculation of a vector valued QNode with post-processing""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=0) @@ -1013,24 +855,18 @@ def circuit(x): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_hessian_ragged(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_hessian_ragged(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test hessian calculation of a ragged QNode""" if diff_method not in {"parameter-shift", "backprop", "hadamard"}: pytest.skip("Test only supports parameter-shift or backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.RY(x[0], wires=0) @@ -1081,23 +917,21 @@ def circuit(x): ] np.testing.assert_allclose(hess, expected_hess, atol=tol, rtol=0, verbose=True) - def test_state(self, dev_name, diff_method, grad_on_execution, tol, interface): + def test_state(self, dev, diff_method, grad_on_execution, device_vjp, tol, interface): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1107,7 +941,6 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is tf.complex128 probs = tf.math.abs(res) ** 2 return probs[0] + probs[2] @@ -1122,21 +955,27 @@ def cost_fn(x, y): assert np.allclose(grad, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector(self, state, dev_name, diff_method, grad_on_execution, tol, interface): + @pytest.mark.parametrize("dtype", ("int32", "int64")) + def test_projector( + self, state, dev, diff_method, grad_on_execution, device_vjp, tol, interface, dtype + ): """Test that the variance of a projector is correctly returned""" - kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface - ) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "interface": interface, + "device_vjp": device_vjp, + } if diff_method == "adjoint": - pytest.skip("Adjoint does not support projectors") - elif diff_method == "hadamard": + pytest.skip("adjoint supports either all expvals or all diagonal measurements.") + if diff_method == "hadamard": pytest.skip("Variance not implemented yet.") elif diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - dev = qml.device(dev_name, wires=2) - P = tf.constant(state) + P = tf.constant(state, dtype=dtype) x, y = 0.765, -0.654 weights = tf.Variable([x, y], dtype=tf.float64) @@ -1161,95 +1000,69 @@ def circuit(weights): ] assert np.allclose(grad, expected, atol=tol, rtol=0) + def test_postselection_differentiation( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): + """Test that when postselecting with default.qubit, differentiation works correctly.""" -@pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": 0.05}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - r = tf.Variable(0.543, dtype=tf.float64) - phi = tf.Variable(-0.654, dtype=tf.float64) - - @qnode(dev, diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - with tf.GradientTape() as tape: - res = circuit(r, phi) - - expected = np.exp(2 * r) * np.sin(phi) ** 2 + np.exp(-2 * r) * np.cos(phi) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) - - # circuit jacobians - grad = tape.gradient(res, [r, phi]) - expected = [ - 2 * np.exp(2 * r) * np.sin(phi) ** 2 - 2 * np.exp(-2 * r) * np.cos(phi) ** 2, - 2 * np.sinh(2 * r) * np.sin(2 * phi), - ] - assert np.allclose(grad, expected, atol=tol, rtol=0) + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") - def test_second_order_observable(self, diff_method, kwargs, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) - n = tf.Variable(0.12, dtype=tf.float64) - a = tf.Variable(0.765, dtype=tf.float64) + phi = tf.Variable(1.23) + theta = tf.Variable(4.56) - @qnode(dev, diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) + assert np.allclose(circuit(phi, theta), expected_circuit(theta)) - with tf.GradientTape() as tape: - res = circuit(n, a) + with tf.GradientTape() as res_tape: + res = circuit(phi, theta) + gradient = res_tape.gradient(res, [phi, theta]) - expected = n**2 + n + np.abs(a) ** 2 * (1 + 2 * n) - assert np.allclose(res, expected, atol=tol, rtol=0) + with tf.GradientTape() as expected_tape: + expected = expected_circuit(theta) + exp_theta_grad = expected_tape.gradient(expected, theta) - # circuit jacobians - grad = tape.gradient(res, [n, a]) - expected = [2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)] - assert np.allclose(grad, expected, atol=tol, rtol=0) + assert np.allclose(gradient, [0.0, exp_theta_grad]) @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the TF interface""" - def test_gradient_expansion(self, dev_name, diff_method, grad_on_execution, interface): + def test_gradient_expansion(self, dev, diff_method, grad_on_execution, device_vjp, interface): """Test that a *supported* operation with no gradient recipe is expanded for both parameter-shift and finite-differences, but not for execution.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) class PhaseShift(qml.PhaseShift): grad_method = None @@ -1263,6 +1076,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, ) def circuit(x): qml.Hadamard(wires=0) @@ -1274,7 +1088,6 @@ def circuit(x): with tf.GradientTape() as t2: with tf.GradientTape() as t1: loss = circuit(x) - res = t1.gradient(loss, x) assert np.allclose(res, -3 * np.sin(3 * x)) @@ -1286,20 +1099,13 @@ def circuit(x): @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff, interface + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - class PhaseShift(qml.PhaseShift): grad_method = None @@ -1312,6 +1118,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, + device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1329,24 +1136,24 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, tol, interface + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, tol, interface ): """Test that if there are non-commuting groups and the number of shots is None the first and second order gradients are correctly evaluated""" - kwargs = dict( - diff_method=diff_method, - grad_on_execution=grad_on_execution, - max_diff=max_diff, - interface=interface, - ) + kwargs = { + "diff_method": diff_method, + "grad_on_execution": grad_on_execution, + "max_diff": max_diff, + "interface": interface, + "device_vjp": device_vjp, + } if diff_method in ["adjoint", "hadamard"]: pytest.skip("The adjoint/hadamard method does not yet support Hamiltonians") elif diff_method == "spsa": + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} - dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1393,14 +1200,14 @@ def circuit(data, weights, coeffs): assert np.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, max_diff, mocker, interface + def test_hamiltonian_finite_shots( + self, dev, diff_method, grad_on_execution, device_vjp, max_diff, interface ): - """Test that the Hamiltonian is expanded if there + """Test that the Hamiltonian is correctly measured if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.3 + tol = 0.1 if diff_method in ("adjoint", "backprop", "hadamard"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": @@ -1413,8 +1220,6 @@ def test_hamiltonian_expansion_finite_shots( elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} - dev = qml.device(dev_name, wires=3, shots=50000) - spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1423,6 +1228,7 @@ def test_hamiltonian_expansion_finite_shots( grad_on_execution=grad_on_execution, max_diff=max_diff, interface=interface, + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1438,17 +1244,19 @@ def circuit(data, weights, coeffs): c = tf.Variable([-0.6543, 0.24, 0.54], dtype=tf.float64) # test output - with tf.GradientTape(persistent=True) as t2: + with tf.GradientTape(persistent=True) as _t2: with tf.GradientTape() as t1: - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) # pylint:disable=unexpected-keyword-arg expected = c[2] * np.cos(d[1] + w[1]) - c[1] * np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]) assert np.allclose(res, expected, atol=tol) - spy.assert_called() # test gradients grad = t1.gradient(res, [d, w, c]) + if diff_method in ["finite-diff", "spsa"]: + pytest.skip(f"{diff_method} not compatible") + expected_w = [ -c[1] * np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), -c[1] * np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]) - c[2] * np.sin(d[1] + w[1]), @@ -1458,33 +1266,36 @@ def circuit(data, weights, coeffs): assert np.allclose(grad[2], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2: - grad2_c = t2.jacobian(grad[2], c) - assert grad2_c is None + # TODO: figure out why grad2_c is np.zeros((3,3)) instead of None + # if diff_method == "parameter-shift" and max_diff == 2: + # grad2_c = _t2.jacobian(grad[2], c) + # print(grad2_c, grad[2], c) + # assert grad2_c is None - grad2_w_c = t2.jacobian(grad[1], c) - expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ - 0, - -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), - -np.sin(d[1] + w[1]), - ] - assert np.allclose(grad2_w_c, expected, atol=tol) + # grad2_w_c = _t2.jacobian(grad[1], c) + # expected = [0, -np.cos(d[0] + w[0]) * np.sin(d[1] + w[1]), 0], [ + # 0, + # -np.cos(d[1] + w[1]) * np.sin(d[0] + w[0]), + # -np.sin(d[1] + w[1]), + # ] + # assert np.allclose(grad2_w_c, expected, atol=tol) class TestSample: """Tests for the sample integration""" + # pylint:disable=unexpected-keyword-arg + def test_sample_dimension(self): """Test sampling works as expected""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert isinstance(res, tuple) assert len(res) == 2 @@ -1496,15 +1307,14 @@ def circuit(): def test_sampling_expval(self): """Test sampling works as expected if combined with expectation values""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert len(res) == 2 assert isinstance(res, tuple) @@ -1515,76 +1325,67 @@ def circuit(): def test_sample_combination(self): """Test the output of combining expval, var and sample""" - n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit() + result = circuit(shots=10) assert len(result) == 3 - assert result[0].shape == (n_sample,) + assert result[0].shape == (10,) assert result[1].shape == () assert result[2].shape == () assert isinstance(result[0], tf.Tensor) assert isinstance(result[1], tf.Tensor) assert isinstance(result[2], tf.Tensor) - assert result[0].dtype is tf.float64 - assert result[1].dtype is tf.float64 - assert result[2].dtype is tf.float64 + assert result[0].dtype is tf.float64 # pylint:disable=no-member + assert result[1].dtype is tf.float64 # pylint:disable=no-member + assert result[2].dtype is tf.float64 # pylint:disable=no-member def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" - n_sample = 10 - - dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit() + result = circuit(shots=10) assert isinstance(result, tf.Tensor) - assert np.array_equal(result.shape, (n_sample,)) + assert np.array_equal(result.shape, (10,)) def test_multi_wire_sample_regular_shape(self): """Test the return type and shape of sampling multiple wires where a rectangular array is expected""" - n_sample = 10 - - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(dev, diff_method="parameter-shift", interface="tf") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="tf") def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit() + result = circuit(shots=10) result = tf.stack(result) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tf.Tensor) - assert np.array_equal(result.shape, (3, n_sample)) + assert np.array_equal(result.shape, (3, 10)) assert result.dtype == tf.float64 def test_counts(self): """Test counts works as expected for TF""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100) - @qnode(dev, interface="tf") + # pylint:disable=unsubscriptable-object,no-member + @qnode(DefaultQubit(), interface="tf") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.counts(qml.PauliZ(0)) - res = circuit() + res = circuit(shots=100) assert isinstance(res, dict) assert list(res.keys()) == [-1, 1] @@ -1614,12 +1415,11 @@ class TestAutograph: def test_autograph_gradients(self, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated""" - dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(dev, diff_method="parameter-shift", interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1628,7 +1428,7 @@ def circuit(x, y): with tf.GradientTape() as tape: p0, p1 = circuit(x, y) - loss = p0[0] + p1[1] + loss = p0[0] + p1[1] # pylint:disable=unsubscriptable-object expected = tf.cos(x / 2) ** 2 + (1 - tf.cos(x) * tf.cos(y)) / 2 assert np.allclose(loss, expected, atol=tol, rtol=0) @@ -1640,12 +1440,11 @@ def circuit(x, y): def test_autograph_jacobian(self, decorator, interface, tol): """Test that a parameter-shift vector-valued QNode can be compiled using @tf.function, and differentiated""" - dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(dev, diff_method="parameter-shift", max_diff=1, interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1683,10 +1482,14 @@ def circuit(x, y): def test_autograph_adjoint_single_param(self, grad_on_execution, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" - dev = qml.device("default.qubit.legacy", wires=1) @decorator - @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution) + @qnode( + DefaultQubit(), + diff_method="adjoint", + interface=interface, + grad_on_execution=grad_on_execution, + ) def circuit(x): qml.RY(x, wires=0) return qml.expval(qml.PauliZ(0)) @@ -1707,10 +1510,14 @@ def circuit(x): def test_autograph_adjoint_multi_params(self, grad_on_execution, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" - dev = qml.device("default.qubit.legacy", wires=1) @decorator - @qnode(dev, diff_method="adjoint", interface=interface, grad_on_execution=grad_on_execution) + @qnode( + DefaultQubit(), + diff_method="adjoint", + interface=interface, + grad_on_execution=grad_on_execution, + ) def circuit(x): qml.RY(x[0], wires=0) qml.RX(x[1], wires=0) @@ -1729,16 +1536,44 @@ def circuit(x): expected_g = [-tf.sin(a) * tf.cos(b), -tf.cos(a) * tf.sin(b)] assert np.allclose(g, expected_g, atol=tol, rtol=0) + @pytest.mark.xfail + @pytest.mark.parametrize("grad_on_execution", [True, False]) + def test_autograph_adjoint_multi_out(self, grad_on_execution, decorator, interface, tol): + """Test that a parameter-shift QNode can be compiled + using @tf.function, and differentiated to second order""" + + @decorator + @qnode( + DefaultQubit(), + diff_method="adjoint", + interface=interface, + grad_on_execution=grad_on_execution, + ) + def circuit(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(0)) + + x = tf.Variable(0.5, dtype=tf.float64) + + with tf.GradientTape() as tape: + res = qml.math.hstack(circuit(x)) + g = tape.jacobian(res, x) + a = x * 1.0 + + expected_res = [tf.cos(a), tf.sin(a)] + assert np.allclose(res, expected_res, atol=tol, rtol=0) + + expected_g = [-tf.sin(a), tf.cos(a)] + assert np.allclose(g, expected_g, atol=tol, rtol=0) + def test_autograph_ragged_differentiation(self, decorator, interface, tol): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - dev = qml.device("default.qubit.legacy", wires=2) - x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(dev, diff_method="parameter-shift", max_diff=1, interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=1, interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1770,12 +1605,11 @@ def circuit(x, y): def test_autograph_hessian(self, decorator, interface, tol): """Test that a parameter-shift QNode can be compiled using @tf.function, and differentiated to second order""" - dev = qml.device("default.qubit.legacy", wires=1) a = tf.Variable(0.543, dtype=tf.float64) b = tf.Variable(-0.654, dtype=tf.float64) @decorator - @qnode(dev, diff_method="parameter-shift", max_diff=2, interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", max_diff=2, interface=interface) def circuit(x, y): qml.RY(x, wires=0) qml.RX(y, wires=0) @@ -1801,35 +1635,15 @@ def circuit(x, y): ] assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - def test_autograph_sample(self, decorator, interface): - """Test that a QNode returning raw samples can be compiled - using @tf.function""" - dev = qml.device("default.qubit", wires=2, shots=100) - x = tf.Variable(0.543, dtype=tf.float64) - y = tf.Variable(-0.654, dtype=tf.float64) - - @decorator - @qnode(dev, diff_method="parameter-shift", interface=interface) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.sample() - - result = circuit(x, y) - result = np.array(result).flatten() - assert np.all([r in [1, 0] for r in result]) - def test_autograph_state(self, decorator, interface, tol): """Test that a parameter-shift QNode returning a state can be compiled using @tf.function""" - dev = qml.device("default.qubit.legacy", wires=2) x = tf.Variable(0.543, dtype=tf.float64) y = tf.Variable(-0.654, dtype=tf.float64) # TODO: fix this for diff_method=None @decorator - @qnode(dev, diff_method="parameter-shift", interface=interface) + @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) def circuit(x, y): qml.RX(x, wires=[0]) qml.RY(y, wires=[1]) @@ -1846,16 +1660,15 @@ def circuit(x, y): def test_autograph_dimension(self, decorator, interface): """Test sampling works as expected""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) @decorator - @qnode(dev, diff_method="parameter-shift", interface=interface) - def circuit(): + @qnode(DefaultQubit(), diff_method="parameter-shift", interface=interface) + def circuit(**_): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) # pylint:disable=unexpected-keyword-arg res = tf.stack(res) assert res.shape == (2, 10) @@ -1863,24 +1676,23 @@ def circuit(): @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution,device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" def test_grad_single_measurement_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For one measurement and one param, the gradient is a float.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -1898,19 +1710,16 @@ def circuit(a): assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -1931,19 +1740,16 @@ def circuit(a, b): assert grad[1].shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -1962,68 +1768,64 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable(0.1) + a = tf.Variable(0.1, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert isinstance(jac, tf.Tensor) assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable(0.1) - b = tf.Variable(0.2) + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) - jac = tape.jacobian(res, (a, b)) + jac = tape.jacobian(res, (a, b), experimental_use_pfor=not device_vjp) assert isinstance(jac, tuple) @@ -2034,101 +1836,94 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """For a multi dimensional measurement (probs), check that a single array is returned.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) - a = tf.Variable([0.1, 0.2]) + a = tf.Variable([0.1, 0.2], dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert isinstance(jac, tf.Tensor) assert jac.shape == (4, 2) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The jacobian of multiple measurements with a single params return an array.""" - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable(0.1) + a = tf.Variable(0.1, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert isinstance(jac, tf.Tensor) assert jac.shape == (5,) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable(0.1) - b = tf.Variable(0.2) + a = tf.Variable(0.1, dtype=tf.float64) + b = tf.Variable(0.2, dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a, b) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, (a, b)) + jac = tape.jacobian(res, (a, b), experimental_use_pfor=not device_vjp) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2140,50 +1935,40 @@ def circuit(a, b): assert jac[1].shape == (5,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because of probabilities.") - - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("state adjoint diff not supported with lightning") @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) - a = tf.Variable([0.1, 0.2]) + a = tf.Variable([0.1, 0.2], dtype=tf.float64) - with tf.GradientTape() as tape: + with tf.GradientTape(persistent=device_vjp) as tape: res = circuit(a) res = tf.experimental.numpy.hstack(res) - jac = tape.jacobian(res, a) + jac = tape.jacobian(res, a, experimental_use_pfor=not device_vjp) assert isinstance(jac, tf.Tensor) assert jac.shape == (5, 2) def test_hessian_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2196,6 +1981,7 @@ def test_hessian_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2222,20 +2008,13 @@ def circuit(x, y): assert hess[1].shape == (2,) def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2244,6 +2023,7 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2262,10 +2042,10 @@ def circuit(x): assert isinstance(hess, tf.Tensor) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution, interface): + def test_hessian_var_multiple_params( + self, dev, diff_method, grad_on_execution, device_vjp, interface + ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2281,6 +2061,7 @@ def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execut diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2307,7 +2088,7 @@ def circuit(x, y): assert hess[1].shape == (2,) def test_hessian_var_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2316,8 +2097,6 @@ def test_hessian_var_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because of variance.") - dev = qml.device(dev_name, wires=2) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2326,6 +2105,7 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2345,16 +2125,9 @@ def circuit(x): assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - num_wires = 2 - - if diff_method == "hadamard": - pytest.skip("Test does not support hadamard because multiple measurements.") - - dev = qml.device(dev_name, wires=num_wires) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2367,6 +2140,7 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2394,7 +2168,7 @@ def circuit(x, y): assert hess[1].shape == (6,) def test_hessian_probs_expval_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" @@ -2404,10 +2178,6 @@ def test_hessian_probs_expval_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because multiple measurements.") - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2416,6 +2186,7 @@ def test_hessian_probs_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2436,11 +2207,9 @@ def circuit(x): assert hess.shape == (3, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") if diff_method == "hadamard": @@ -2455,6 +2224,7 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2482,7 +2252,7 @@ def circuit(x, y): assert hess[1].shape == (6,) def test_hessian_probs_var_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, interface + self, dev, diff_method, grad_on_execution, device_vjp, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" if diff_method == "adjoint": @@ -2490,8 +2260,6 @@ def test_hessian_probs_var_multiple_param_array( if diff_method == "hadamard": pytest.skip("Test does not support hadamard because of variance.") - dev = qml.device(dev_name, wires=2) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode( @@ -2500,6 +2268,7 @@ def test_hessian_probs_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2520,17 +2289,49 @@ def circuit(x): assert hess.shape == (3, 2, 2) -@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) -def test_no_ops(dev_name): +def test_no_ops(): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = qml.device(dev_name, wires=1) - - @qml.qnode(dev, interface="tf") + @qml.qnode(DefaultQubit(), interface="tf") def circuit(): qml.Hadamard(wires=0) return qml.state() res = circuit() assert isinstance(res, tf.Tensor) + + +def test_error_device_vjp_jacobian(): + """Test a ValueError is raised if a jacobian is attempted to be computed with device_vjp=True.""" + + dev = qml.device("default.qubit") + + @qml.qnode(dev, diff_method="adjoint", device_vjp=True) + def circuit(x): + qml.RX(x, wires=0) + return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(0)) + + x = tf.Variable(0.1) + + with tf.GradientTape() as tape: + y = qml.math.hstack(circuit(x)) + + with pytest.raises(ValueError): + tape.jacobian(y, x) + + +def test_error_device_vjp_state_float32(): + """Test a ValueError is raised is state differentiation is attemped with float32 parameters.""" + + dev = qml.device("default.qubit") + + @qml.qnode(dev, diff_method="adjoint", device_vjp=True) + def circuit(x): + qml.RX(x, wires=0) + return qml.probs(wires=0) + + x = tf.Variable(0.1, dtype=tf.float32) + with pytest.raises(ValueError, match="tensorflow with adjoint differentiation of the state"): + with tf.GradientTape(): + circuit(x) diff --git a/tests/interfaces/test_tensorflow_qnode_shot_vector.py b/tests/interfaces/test_tensorflow_qnode_shot_vector.py index 614c71ee7d4..c0340aaab9a 100644 --- a/tests/interfaces/test_tensorflow_qnode_shot_vector.py +++ b/tests/interfaces/test_tensorflow_qnode_shot_vector.py @@ -1,4 +1,4 @@ -# Copyright 2022 Xanadu Quantum Technologies Inc. +# Copyright 2023 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. @@ -12,12 +12,13 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the TF interface with shot vectors and with a QNode""" -# pylint: disable=too-many-arguments +# pylint: disable=too-many-arguments,unexpected-keyword-arg import pytest import pennylane as qml from pennylane import numpy as np from pennylane import qnode +from pennylane.devices import DefaultQubit pytestmark = pytest.mark.tf @@ -27,9 +28,9 @@ shots_and_num_copies_hess = [((10, (5, 1)), 2)] qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", {"h": 10e-2}], - ["default.qubit.legacy", "parameter-shift", {}], - ["default.qubit.legacy", "spsa", {"h": 10e-2, "num_directions": 20}], + [DefaultQubit(), "finite-diff", {"h": 10e-2}], + [DefaultQubit(), "parameter-shift", {}], + [DefaultQubit(), "spsa", {"h": 10e-2, "num_directions": 20}], ] TOLS = { @@ -45,16 +46,15 @@ @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) @pytest.mark.parametrize( - "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnWithShotVectors: """Class to test the shape of the Grad/Jacobian/Hessian with different return types and shot vectors.""" def test_jac_single_measurement_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and one param, the gradient is a float.""" - dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -65,7 +65,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -74,10 +74,9 @@ def circuit(a): assert jac.shape == (num_copies,) def test_jac_single_measurement_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" - dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -89,7 +88,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -101,10 +100,9 @@ def circuit(a, b): assert j.shape == (num_copies,) def test_jacobian_single_measurement_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For one measurement and multiple param as a single array params, the gradient is an array.""" - dev = qml.device(dev_name, wires=1, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -115,7 +113,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -124,11 +122,10 @@ def circuit(a): assert jac.shape == (num_copies, 2) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -139,7 +136,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -148,11 +145,10 @@ def circuit(a): assert jac.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -164,7 +160,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, (a, b)) @@ -176,11 +172,10 @@ def circuit(a, b): assert j.shape == (num_copies, 4) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -191,7 +186,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack(res) jac = tape.jacobian(res, a) @@ -200,10 +195,9 @@ def circuit(a): assert jac.shape == (num_copies, 4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The gradient of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1) par_1 = tf.Variable(0.2) @@ -216,7 +210,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) with tf.GradientTape() as tape: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, (par_0, par_1)) @@ -228,10 +222,9 @@ def circuit(x, y): assert j.shape == (num_copies, 2) def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -243,7 +236,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2, 0.3]) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([qml.math.stack(r) for r in res]) jac = tape.jacobian(res, a) @@ -252,10 +245,9 @@ def circuit(a): assert jac.shape == (num_copies, 2, 3) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a single params return an array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -266,7 +258,7 @@ def circuit(a): a = tf.Variable(0.1) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -275,10 +267,9 @@ def circuit(a): assert jac.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a, b): @@ -290,7 +281,7 @@ def circuit(a, b): b = tf.Variable(0.2) with tf.GradientTape() as tape: - res = circuit(a, b) + res = circuit(a, b, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, (a, b)) @@ -302,10 +293,9 @@ def circuit(a, b): assert j.shape == (num_copies, 5) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The jacobian of multiple measurements with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) @qnode(dev, diff_method=diff_method, interface=interface, **gradient_kwargs) def circuit(a): @@ -316,7 +306,7 @@ def circuit(a): a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as tape: - res = circuit(a) + res = circuit(a, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape.jacobian(res, a) @@ -328,16 +318,15 @@ def circuit(a): @pytest.mark.slow @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies_hess) @pytest.mark.parametrize( - "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnShotVectorHessian: """Class to test the shape of the Hessian with different return types and shot vectors.""" def test_hessian_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of a single measurement with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1, dtype=tf.float64) par_1 = tf.Variable(0.2, dtype=tf.float64) @@ -351,7 +340,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack(res) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -366,10 +355,9 @@ def circuit(x, y): assert h.shape == (2, num_copies) def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of single measurement with a multiple params array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) params = tf.Variable([0.1, 0.2], dtype=tf.float64) @@ -382,7 +370,7 @@ def circuit(x): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(params) + res = circuit(params, shots=shots) res = qml.math.stack(res) jac = tape2.jacobian(res, params, experimental_use_pfor=False) @@ -393,10 +381,9 @@ def circuit(x): assert hess.shape == (num_copies, 2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) par_0 = tf.Variable(0.1, dtype=tf.float64) par_1 = tf.Variable(0.2, dtype=tf.float64) @@ -410,7 +397,7 @@ def circuit(x, y): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(par_0, par_1) + res = circuit(par_0, par_1, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape2.jacobian(res, (par_0, par_1), experimental_use_pfor=False) @@ -425,12 +412,10 @@ def circuit(x, y): assert h.shape == (2, num_copies, 3) def test_hessian_expval_probs_multiple_param_array( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """The hessian of multiple measurements with a multiple param array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) - params = tf.Variable([0.1, 0.2], dtype=tf.float64) @qnode(dev, diff_method=diff_method, interface=interface, max_diff=2, **gradient_kwargs) @@ -442,7 +427,7 @@ def circuit(x): with tf.GradientTape() as tape1: with tf.GradientTape(persistent=True) as tape2: - res = circuit(params) + res = circuit(params, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) jac = tape2.jacobian(res, params, experimental_use_pfor=False) @@ -453,22 +438,22 @@ def circuit(x): assert hess.shape == (num_copies, 3, 2, 2) -shots_and_num_copies = [((1000000, 900000, 800000), 3), ((1000000, (900000, 2)), 3)] +shots_and_num_copies = [((30000, 28000, 26000), 3), ((30000, (28000, 2)), 3)] @pytest.mark.parametrize("shots,num_copies", shots_and_num_copies) @pytest.mark.parametrize( - "interface,dev_name,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,gradient_kwargs", interface_and_qubit_device_and_diff_method ) class TestReturnShotVectorIntegration: """Tests for the integration of shots with the TF interface.""" def test_single_expectation_value( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """Tests correct output shape and evaluation for a tape with a single expval output""" - dev = qml.device(dev_name, wires=2, shots=shots) + x = tf.Variable(0.543) y = tf.Variable(-0.654) @@ -480,7 +465,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) with tf.GradientTape() as tape: - res = circuit(x, y) + res = circuit(x, y, shots=shots) res = qml.math.stack(res) all_res = tape.jacobian(res, (x, y)) @@ -497,11 +482,11 @@ def circuit(x, y): assert np.allclose(res, exp, atol=tol, rtol=0) def test_prob_expectation_values( - self, dev_name, diff_method, gradient_kwargs, shots, num_copies, interface + self, dev, diff_method, gradient_kwargs, shots, num_copies, interface ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - dev = qml.device(dev_name, wires=2, shots=shots) + x = tf.Variable(0.543) y = tf.Variable(-0.654) @@ -513,7 +498,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0)), qml.probs(wires=[0, 1]) with tf.GradientTape() as tape: - res = circuit(x, y) + res = circuit(x, y, shots=shots) res = qml.math.stack([tf.experimental.numpy.hstack(r) for r in res]) all_res = tape.jacobian(res, (x, y)) diff --git a/tests/interfaces/test_torch.py b/tests/interfaces/test_torch.py index 5aac97fb0e3..3cdcf5eae30 100644 --- a/tests/interfaces/test_torch.py +++ b/tests/interfaces/test_torch.py @@ -1,4 +1,4 @@ -# Copyright 2018-2022 Xanadu Quantum Technologies Inc. +# Copyright 2018-2023 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. @@ -11,254 +11,45 @@ # 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. -"""Unit tests for the Torch interface""" -# pylint: disable=protected-access,too-few-public-methods +"""Torch specific tests for execute and default qubit 2.""" import numpy as np import pytest +from param_shift_dev import ParamShiftDerivativesDevice import pennylane as qml from pennylane import execute -from pennylane.gradients import finite_diff, param_shift -from pennylane.typing import TensorLike +from pennylane.devices import DefaultQubit +from pennylane.gradients import param_shift +from pennylane.measurements import Shots -pytestmark = pytest.mark.torch torch = pytest.importorskip("torch") -torch_functional = pytest.importorskip("torch.autograd.functional") -torch_cuda = pytest.importorskip("torch.cuda") - - -@pytest.mark.parametrize("interface", ["torch", "auto"]) -class TestTorchExecuteUnitTests: - """Unit tests for torch execution""" - - def test_jacobian_options(self, interface, mocker): - """Test setting jacobian options""" - spy = mocker.spy(qml.gradients, "param_shift") - - a = torch.tensor([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute( - [tape], - dev, - gradient_fn=param_shift, - gradient_kwargs={"shifts": [(np.pi / 4,)] * 2}, - interface=interface, - )[0] - - res.backward() - - for args in spy.call_args_list: - assert args[1]["shift"] == [(np.pi / 4,)] * 2 - - def test_incorrect_grad_on_execution(self, interface): - """Test that an error is raised if a gradient transform - is used with grad_on_execution=True""" - a = torch.tensor([0.1, 0.2], requires_grad=True) - - dev = qml.device("default.qubit.legacy", wires=1) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - with pytest.raises( - ValueError, match="Gradient transforms cannot be used with grad_on_execution=True" - ): - execute( - [tape], dev, gradient_fn=param_shift, grad_on_execution=True, interface=interface - ) - - def test_grad_on_execution_reuse_state(self, interface, mocker): - """Test that grad_on_execution uses the `device.execute_and_gradients` pathway - while reusing the quantum state.""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(dev, "execute_and_gradients") - a = torch.tensor([0.1, 0.2], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - interface=interface, - ) - - # adjoint method only performs a single device execution, but gets both result and gradient - assert dev.num_executions == 1 - spy.assert_called() - - def test_grad_on_execution(self, interface, mocker): - """Test that grad on execution uses the `device.execute_and_gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(dev, "execute_and_gradients") - - a = torch.tensor([0.1, 0.2], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian"}, - interface=interface, - ) - - # two device executions; one for the value, one for the Jacobian - assert dev.num_executions == 2 - spy.assert_called() - - def test_no_grad_on_execution(self, interface, mocker): - """Test that no grad on execution uses the `device.batch_execute` and `device.gradients` pathway""" - dev = qml.device("default.qubit.legacy", wires=1) - spy_execute = mocker.spy(qml.devices.DefaultQubitLegacy, "batch_execute") - spy_gradients = mocker.spy(qml.devices.DefaultQubitLegacy, "gradients") - - a = torch.tensor([0.1, 0.2], requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute( - [tape], - dev, - gradient_fn="device", - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - interface=interface, - )[0] +pytestmark = pytest.mark.torch - assert dev.num_executions == 1 - spy_execute.assert_called() - spy_gradients.assert_not_called() - res.backward() - spy_gradients.assert_called() +@pytest.fixture(autouse=True) +def run_before_and_after_tests(): + torch.set_default_dtype(torch.float64) + yield + torch.set_default_dtype(torch.float32) +# pylint: disable=too-few-public-methods class TestCaching: - """Test for caching behaviour""" - - def test_cache_maxsize(self, mocker): - """Test the cachesize property of the cache""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cachesize): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute( - [tape], dev, gradient_fn=param_shift, cachesize=cachesize, interface="torch" - )[0][0] - - params = torch.tensor([0.1, 0.2], requires_grad=True) - res = cost(params, cachesize=2) - res.backward() - cache = spy.call_args.kwargs["cache"] - - assert cache.maxsize == 2 - assert cache.currsize == 2 - assert len(cache) == 2 - - def test_custom_cache(self, mocker): - """Test the use of a custom cache object""" - dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="torch")[0][ - 0 - ] - - custom_cache = {} - params = torch.tensor([0.1, 0.2], requires_grad=True) - res = cost(params, cache=custom_cache) - res.backward() - - cache = spy.call_args.kwargs["cache"] - assert cache is custom_cache - - def test_caching_param_shift(self): - """Test that, with the parameter-shift transform, - Torch always uses the optimum number of evals when computing the Jacobian.""" - dev = qml.device("default.qubit.legacy", wires=1) - - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, gradient_fn=param_shift, cache=cache, interface="torch")[0][ - 0 - ] - - # Without caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - params = torch.tensor([0.1, 0.2], requires_grad=True) - torch_functional.jacobian(lambda p: cost(p, cache=None), params) - assert dev.num_executions == 5 - - # With caching, 5 evaluations are required to compute - # the Jacobian: 1 (forward pass) + (2 shifts * 2 params) - dev._num_executions = 0 - torch_functional.jacobian(lambda p: cost(p, cache=True), params) - assert dev.num_executions == 5 + """Tests for caching behaviour""" + @pytest.mark.skip("caching is not implemented for torch") @pytest.mark.parametrize("num_params", [2, 3]) - def test_caching_param_shift_hessian(self, num_params, tol): - """Test that, with the parameter-shift transform, + def test_caching_param_shift_hessian(self, num_params): + """Test that, when using parameter-shift transform, caching reduces the number of evaluations to their optimum when computing Hessians.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor(np.arange(1, num_params + 1) / 10, requires_grad=True) + dev = DefaultQubit() + params = torch.arange(1, num_params + 1, requires_grad=True, dtype=torch.float64) N = len(params) - def cost(x, cache): + def get_cost_tape(x): with qml.queuing.AnnotatedQueue() as q: qml.RX(x[0], wires=[0]) qml.RY(x[1], wires=[1]) @@ -267,447 +58,447 @@ def cost(x, cache): qml.RZ(x[i], wires=[i % 2]) qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) + qml.var(qml.prod(qml.PauliZ(0), qml.PauliX(1))) - tape = qml.tape.QuantumScript.from_queue(q) + return qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], dev, gradient_fn=param_shift, cache=cache, interface="torch", max_diff=2 + def cost_no_cache(x): + return qml.execute( + [get_cost_tape(x)], + dev, + gradient_fn=qml.gradients.param_shift, + cache=False, + max_diff=2, + )[0] + + def cost_cache(x): + return qml.execute( + [get_cost_tape(x)], + dev, + gradient_fn=qml.gradients.param_shift, + cache=True, + max_diff=2, )[0] # No caching: number of executions is not ideal - hess1 = torch.autograd.functional.hessian(lambda x: cost(x, cache=None), params) + with qml.Tracker(dev) as tracker: + hess1 = torch.autograd.functional.hessian(cost_no_cache, params) if num_params == 2: # compare to theoretical result - x, y, *_ = params.detach() + x, y, *_ = params.clone().detach() expected = torch.tensor( [ - [2 * np.cos(2 * x) * np.sin(y) ** 2, np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + [2 * torch.cos(2 * x) * torch.sin(y) ** 2, torch.sin(2 * x) * torch.sin(2 * y)], + [ + torch.sin(2 * x) * torch.sin(2 * y), + -2 * torch.cos(x) ** 2 * torch.cos(2 * y), + ], ] ) - assert np.allclose(expected, hess1, atol=tol, rtol=0) + assert torch.allclose(expected, hess1) expected_runs = 1 # forward pass - expected_runs += 2 * N # Jacobian - expected_runs += 4 * N + 1 # Hessian diagonal - expected_runs += 4 * N**2 # Hessian off-diagonal - assert dev.num_executions == expected_runs + + # Jacobian of an involutory observable: + # ------------------------------------ + # + # 2 * N execs: evaluate the analytic derivative of + # 1 execs: Get , the expectation value of the tape with unshifted parameters. + num_shifted_evals = 2 * N + runs_for_jacobian = num_shifted_evals + 1 + expected_runs += runs_for_jacobian + + # Each tape used to compute the Jacobian is then shifted again + expected_runs += runs_for_jacobian * num_shifted_evals + assert tracker.totals["executions"] == expected_runs # Use caching: number of executions is ideal - dev._num_executions = 0 - hess2 = torch.autograd.functional.hessian(lambda x: cost(x, cache=True), params) - assert np.allclose(hess1, hess2, atol=tol, rtol=0) + + with qml.Tracker(dev) as tracker2: + hess2 = torch.autograd.functional.hessian(cost_cache, params) + assert torch.allclose(hess1, hess2) expected_runs_ideal = 1 # forward pass expected_runs_ideal += 2 * N # Jacobian expected_runs_ideal += N + 1 # Hessian diagonal expected_runs_ideal += 4 * N * (N - 1) // 2 # Hessian off-diagonal - assert dev.num_executions == expected_runs_ideal + assert tracker2.totals["executions"] == expected_runs_ideal assert expected_runs_ideal < expected_runs - def test_caching_adjoint_no_grad_on_execution(self): - """Test that caching reduces the number of adjoint evaluations - when grad_on_execution=False""" - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor([0.1, 0.2, 0.3]) - def cost(a, cache): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a[0], wires=0) - qml.RX(a[1], wires=0) - qml.RY(a[2], wires=0) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) +# add tests for lightning 2 when possible +# set rng for device when possible +test_matrix = [ + ({"gradient_fn": param_shift}, Shots(100000), DefaultQubit(seed=42)), + ({"gradient_fn": param_shift}, Shots((100000, 100000)), DefaultQubit(seed=42)), + ({"gradient_fn": param_shift}, Shots(None), DefaultQubit()), + ({"gradient_fn": "backprop"}, Shots(None), DefaultQubit()), + ( + {"gradient_fn": "adjoint", "grad_on_execution": True, "device_vjp": False}, + Shots(None), + DefaultQubit(), + ), + ( + { + "gradient_fn": "adjoint", + "grad_on_execution": False, + "device_vjp": False, + }, + Shots(None), + DefaultQubit(), + ), + ({"gradient_fn": "adjoint", "device_vjp": True}, Shots(None), DefaultQubit()), + ( + {"gradient_fn": "device", "device_vjp": False}, + Shots((100000, 100000)), + ParamShiftDerivativesDevice(), + ), + ( + {"gradient_fn": "device", "device_vjp": True}, + Shots((100000, 100000)), + ParamShiftDerivativesDevice(), + ), +] - tape = qml.tape.QuantumScript.from_queue(q) - return execute( - [tape], - dev, - gradient_fn="device", - cache=cache, - grad_on_execution=False, - gradient_kwargs={"method": "adjoint_jacobian"}, - interface="torch", - )[0] - - # Without caching, 2 evaluations are required. - # 1 for the forward pass, and one for the backward pass - torch_functional.jacobian(lambda x: cost(x, cache=None), params) - assert dev.num_executions == 2 - - # With caching, only 2 evaluations are required. One - # for the forward pass, and one for the backward pass. - dev._num_executions = 0 - torch_functional.jacobian(lambda x: cost(x, cache=True), params) - assert dev.num_executions == 2 - - -torch_devices = [None] - -if torch_cuda.is_available(): - torch_devices.append(torch.device("cuda")) - - -execute_kwargs_integration = [ - {"gradient_fn": param_shift, "interface": "torch"}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": False}, - "interface": "torch", - }, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - "interface": "torch", - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - "interface": "torch", - }, - {"gradient_fn": param_shift, "interface": "auto"}, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": False}, - "interface": "auto", - }, - { - "gradient_fn": "device", - "grad_on_execution": True, - "gradient_kwargs": {"method": "adjoint_jacobian", "use_device_state": True}, - "interface": "auto", - }, - { - "gradient_fn": "device", - "grad_on_execution": False, - "gradient_kwargs": {"method": "adjoint_jacobian"}, - "interface": "auto", - }, -] +def atol_for_shots(shots): + """Return higher tolerance if finite shots.""" + return 1e-2 if shots else 1e-6 -@pytest.mark.gpu -@pytest.mark.parametrize("torch_device", torch_devices) -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_integration) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) class TestTorchExecuteIntegration: """Test the torch interface execute function integrates well for both forward and backward execution""" - def test_execution(self, torch_device, execute_kwargs): - """Test that the execute function produces results with the expected shapes""" - dev = qml.device("default.qubit.legacy", wires=1) - a = torch.tensor(0.1, requires_grad=True, device=torch_device) - b = torch.tensor(0.2, requires_grad=False, device=torch_device) - - with qml.queuing.AnnotatedQueue() as q1: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) + def test_execution(self, execute_kwargs, shots, device): + """Test execution""" - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.expval(qml.PauliZ(0)) + def cost(a, b): + ops1 = [qml.RY(a, wires=0), qml.RX(b, wires=0)] + tape1 = qml.tape.QuantumScript(ops1, [qml.expval(qml.PauliZ(0))], shots=shots) - tape2 = qml.tape.QuantumScript.from_queue(q2) + ops2 = [qml.RY(a, wires="a"), qml.RX(b, wires="a")] + tape2 = qml.tape.QuantumScript(ops2, [qml.expval(qml.PauliZ("a"))], shots=shots) - res = execute([tape1, tape2], dev, **execute_kwargs) + return execute([tape1, tape2], device, **execute_kwargs) - assert isinstance(res, TensorLike) - assert len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () + a = torch.tensor(0.1, requires_grad=True) + b = torch.tensor(0.2, requires_grad=False) - def test_scalar_jacobian(self, torch_device, execute_kwargs, tol): - """Test scalar jacobian calculation by comparing two types of pipelines""" - a = torch.tensor(0.1, requires_grad=True, dtype=torch.float64, device=torch_device) - dev = qml.device("default.qubit.legacy", wires=2) + with device.tracker: + res = cost(a, b) - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + if execute_kwargs.get("grad_on_execution", False): + assert device.tracker.totals["execute_and_derivative_batches"] == 1 + else: + assert device.tracker.totals["batches"] == 1 + assert device.tracker.totals["executions"] == 2 # different wires so different hashes - tape = qml.tape.QuantumScript.from_queue(q) + assert len(res) == 2 + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () + exp = torch.cos(a) * torch.cos(b) + if shots.has_partitioned_shots: + for shot in range(2): + for wire in range(2): + assert qml.math.allclose(res[shot][wire], exp, atol=atol_for_shots(shots)) + else: + for wire in range(2): + assert qml.math.allclose(res[wire], exp, atol=atol_for_shots(shots)) - res = execute([tape], dev, **execute_kwargs)[0] - res.backward() + def test_scalar_jacobian(self, execute_kwargs, shots, device): + """Test scalar jacobian calculation""" + a = torch.tensor(0.1, requires_grad=True) - # compare to backprop gradient def cost(a): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - dev = qml.device("default.qubit.autograd", wires=2) - return dev.batch_execute([tape])[0] - - expected = qml.grad(cost, argnum=0)(0.1) - assert torch.allclose(a.grad, torch.tensor(expected, device=torch_device), atol=tol, rtol=0) - - def test_jacobian(self, torch_device, execute_kwargs, tol): - """Test jacobian calculation by checking against analytic values""" - a_val = 0.1 - b_val = 0.2 - - a = torch.tensor(a_val, requires_grad=True, device=torch_device) - b = torch.tensor(b_val, requires_grad=True, device=torch_device) - - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.RZ(torch.tensor(0.543, device=torch_device), wires=0) - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [1, 2] - - assert isinstance(res, tuple) - - assert isinstance(res[0], torch.Tensor) - assert res[0].shape == () + tape = qml.tape.QuantumScript([qml.RY(a, 0)], [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] + + res = torch.autograd.functional.jacobian(cost, a) + if not shots.has_partitioned_shots: + assert res.shape == () # pylint: disable=no-member + + # compare to standard tape jacobian + tape = qml.tape.QuantumScript([qml.RY(a, wires=0)], [qml.expval(qml.PauliZ(0))]) + tape.trainable_params = [0] + tapes, fn = param_shift(tape) + expected = fn(device.execute(tapes)) + + assert expected.shape == () + if shots.has_partitioned_shots: + for i in range(shots.num_copies): + assert torch.allclose(res[i], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[i], -torch.sin(a), atol=atol_for_shots(shots)) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res, -torch.sin(a), atol=atol_for_shots(shots)) + + def test_jacobian(self, execute_kwargs, shots, device): + """Test jacobian calculation""" + a = torch.tensor(0.1, requires_grad=True) + b = torch.tensor(0.2, requires_grad=True) + + def cost(a, b): + ops = [qml.RY(a, wires=0), qml.RX(b, wires=1), qml.CNOT(wires=[0, 1])] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + [res] = execute([tape], device, **execute_kwargs) + if shots.has_partitioned_shots: + return torch.hstack(res[0] + res[1]) + return torch.hstack(res) + + res = cost(a, b) + expected = torch.tensor([torch.cos(a), -torch.cos(a) * torch.sin(b)]) + if shots.has_partitioned_shots: + assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert isinstance(res[1], torch.Tensor) - assert res[1].shape == () + res = torch.autograd.functional.jacobian(cost, (a, b)) + assert isinstance(res, tuple) and len(res) == 2 - expected = torch.tensor( - [np.cos(a_val), -np.cos(a_val) * np.sin(b_val)], device=torch_device + expected = ( + torch.tensor([-torch.sin(a), torch.sin(a) * torch.sin(b)]), + torch.tensor([0, -torch.cos(a) * torch.cos(b)]), ) - assert torch.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) - assert torch.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert res[0].shape == (4,) + assert res[1].shape == (4,) - loss = res[0] + res[1] + for _r, _e in zip(res, expected): + assert torch.allclose(_r[:2], _e, atol=atol_for_shots(shots)) + assert torch.allclose(_r[2:], _e, atol=atol_for_shots(shots)) - loss.backward() - expected = torch.tensor( - [-np.sin(a_val) + np.sin(a_val) * np.sin(b_val), -np.cos(a_val) * np.cos(b_val)], - dtype=a.dtype, - device=torch_device, - ) - assert torch.allclose(a.grad, expected[0], atol=tol, rtol=0) - assert torch.allclose(b.grad, expected[1], atol=tol, rtol=0) + else: + assert res[0].shape == (2,) + assert res[1].shape == (2,) + + for _r, _e in zip(res, expected): + assert torch.allclose(_r, _e, atol=atol_for_shots(shots)) - def test_tape_no_parameters(self, torch_device, execute_kwargs, tol): + def test_tape_no_parameters(self, execute_kwargs, shots, device): """Test that a tape with no parameters is correctly ignored during the gradient computation""" - dev = qml.device("default.qubit.legacy", wires=1) - params = torch.tensor([0.1, 0.2], requires_grad=True, device=torch_device) - x, y = params.detach() - with qml.queuing.AnnotatedQueue() as q1: - qml.Hadamard(0) - qml.expval(qml.PauliX(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.Hadamard(0)], [qml.expval(qml.PauliX(0))], shots=shots + ) - with qml.queuing.AnnotatedQueue() as q2: - qml.RY(0.5, wires=0) - qml.expval(qml.PauliZ(0)) + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), wires=0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - tape2 = qml.tape.QuantumScript.from_queue(q2) + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - with qml.queuing.AnnotatedQueue() as q3: - qml.RY(params[0], wires=0) - qml.RX(params[1], wires=0) - qml.expval(qml.PauliZ(0)) + tape4 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), 0)], + [qml.probs(wires=[0, 1])], + shots=shots, + ) + res = execute([tape1, tape2, tape3, tape4], device, **execute_kwargs) + if shots.has_partitioned_shots: + res = [qml.math.asarray(ri, like="torch") for r in res for ri in r] + else: + res = [qml.math.asarray(r, like="torch") for r in res] + return sum(torch.hstack(res)) - tape3 = qml.tape.QuantumScript.from_queue(q3) + params = torch.tensor([0.1, 0.2], requires_grad=True) + x, y = params.clone().detach() - res = sum(execute([tape1, tape2, tape3], dev, **execute_kwargs)) - expected = torch.tensor(1 + np.cos(0.5), dtype=res.dtype) + torch.cos(x) * torch.cos(y) - expected = expected.to(device=res.device) + res = cost(params) + expected = 2 + np.cos(0.5) + np.cos(x) * np.cos(y) - assert torch.allclose(res, expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert torch.allclose(res, 2 * expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) res.backward() - grad = params.grad.detach() - expected = torch.tensor( - [-torch.cos(y) * torch.sin(x), -torch.cos(x) * torch.sin(y)], - dtype=grad.dtype, - device=grad.device, - ) - assert torch.allclose(grad, expected, atol=tol, rtol=0) - - def test_reusing_quantum_tape(self, torch_device, execute_kwargs, tol): - """Test re-using a quantum tape by passing new parameters""" - a = torch.tensor(0.1, requires_grad=True, device=torch_device) - b = torch.tensor(0.2, requires_grad=True, device=torch_device) - - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliY(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - assert tape.trainable_params == [0, 1] + expected = torch.tensor([-torch.cos(y) * torch.sin(x), -torch.cos(x) * torch.sin(y)]) + if shots.has_partitioned_shots: + assert torch.allclose(params.grad, 2 * expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(params.grad, expected, atol=atol_for_shots(shots), rtol=0) - res = execute([tape], dev, **execute_kwargs)[0] - loss = res[0] + res[1] - loss.backward() + @pytest.mark.skip("torch cannot reuse tensors in various computations") + def test_tapes_with_different_return_size(self, execute_kwargs, shots, device): + """Test that tapes wit different can be executed and differentiated.""" - a_val = 0.54 - b_val = 0.8 - a = torch.tensor(a_val, requires_grad=True, device=torch_device) - b = torch.tensor(b_val, requires_grad=True, device=torch_device) + if execute_kwargs["gradient_fn"] == "backprop": + pytest.xfail("backprop is not compatible with something about this situation.") - tape = tape.bind_new_parameters([2 * a, b], [0, 1]) - res2 = execute([tape], dev, **execute_kwargs)[0] + def cost(params): + tape1 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))], + shots=shots, + ) - expected = torch.tensor( - [np.cos(2 * a_val), -np.cos(2 * a_val) * np.sin(b_val)], - device=torch_device, - dtype=res2[0].dtype, - ) - assert torch.allclose(res2[0].detach(), expected[0], atol=tol, rtol=0) - assert torch.allclose(res2[1].detach(), expected[1], atol=tol, rtol=0) + tape2 = qml.tape.QuantumScript( + [qml.RY(np.array(0.5), 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) - loss = res2[0] + res2[1] - loss.backward() + tape3 = qml.tape.QuantumScript( + [qml.RY(params[0], 0), qml.RX(params[1], 0)], + [qml.expval(qml.PauliZ(0))], + shots=shots, + ) + res = execute([tape1, tape2, tape3], device, **execute_kwargs) + return torch.hstack([qml.math.asarray(r, like="torch") for r in res]) - expected = torch.tensor( - [ - -2 * np.sin(2 * a_val) + 2 * np.sin(2 * a_val) * np.sin(b_val), - -np.cos(2 * a_val) * np.cos(b_val), - ], - dtype=a.dtype, - device=torch_device, - ) + params = torch.tensor([0.1, 0.2], requires_grad=True) + x, y = params.clone().detach() - assert torch.allclose(a.grad, expected[0], atol=tol, rtol=0) - assert torch.allclose(b.grad, expected[1], atol=tol, rtol=0) + res = cost(params) + assert isinstance(res, torch.Tensor) + assert res.shape == (4,) - def test_classical_processing(self, torch_device, execute_kwargs, tol): - """Test the classical processing of gate parameters within the quantum tape""" - p_val = [0.1, 0.2] - params = torch.tensor(p_val, requires_grad=True, device=torch_device) + assert torch.allclose(res[0], torch.cos(x) * torch.cos(y), atol=atol_for_shots(shots)) + assert torch.allclose(res[1], torch.tensor(1.0), atol=atol_for_shots(shots)) + assert torch.allclose(res[2], torch.cos(torch.tensor(0.5)), atol=atol_for_shots(shots)) + assert torch.allclose(res[3], torch.cos(x) * torch.cos(y), atol=atol_for_shots(shots)) - dev = qml.device("default.qubit.legacy", wires=1) + jac = torch.autograd.functional.jacobian(cost, params) + assert isinstance(jac, torch.Tensor) + assert jac.shape == (4, 2) # pylint: disable=no-member - with qml.queuing.AnnotatedQueue() as q: - qml.RY(params[0] * params[1], wires=0) - qml.RZ(0.2, wires=0) - qml.RX(params[1] + params[1] ** 2 + torch.sin(params[0]), wires=0) - qml.expval(qml.PauliZ(0)) + assert torch.allclose(jac[1:3], torch.tensor(0.0), atol=atol_for_shots(shots)) - tape = qml.tape.QuantumScript.from_queue(q) + d1 = -torch.sin(x) * torch.cos(y) + assert torch.allclose(jac[0, 0], d1, atol=atol_for_shots(shots)) # fails for torch + assert torch.allclose(jac[3, 0], d1, atol=atol_for_shots(shots)) - res = execute([tape], dev, **execute_kwargs)[0] + d2 = -torch.cos(x) * torch.sin(y) + assert torch.allclose(jac[0, 1], d2, atol=atol_for_shots(shots)) # fails for torch + assert torch.allclose(jac[3, 1], d2, atol=atol_for_shots(shots)) - assert tape.trainable_params == [0, 2] + def test_reusing_quantum_tape(self, execute_kwargs, shots, device): + """Test re-using a quantum tape by passing new parameters""" + a = torch.tensor(0.1, requires_grad=True) + b = torch.tensor(0.2, requires_grad=True) - tape_params = torch.tensor([i.detach() for i in tape.get_parameters()], device=torch_device) - expected = torch.tensor( - [p_val[0] * p_val[1], p_val[1] + p_val[1] ** 2 + np.sin(p_val[0])], - dtype=tape_params.dtype, - device=torch_device, + tape = qml.tape.QuantumScript( + [qml.RY(a, 0), qml.RX(b, 1), qml.CNOT((0, 1))], + [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1))], ) + assert tape.trainable_params == [0, 1] - assert torch.allclose( - tape_params, - expected, - atol=tol, - rtol=0, - ) + def cost(a, b): + new_tape = tape.bind_new_parameters([a, b], [0, 1]) + return torch.hstack(execute([new_tape], device, **execute_kwargs)[0]) + + jac = torch.autograd.functional.jacobian(cost, (a, b)) + + a = torch.tensor(0.54, requires_grad=True) + b = torch.tensor(0.8, requires_grad=True) + + # check that the cost function continues to depend on the + # values of the parameters for subsequent calls + res2 = cost(2 * a, b) + expected = torch.tensor([torch.cos(2 * a), -torch.cos(2 * a) * torch.sin(b)]) + assert torch.allclose(res2, expected, atol=atol_for_shots(shots), rtol=0) + + jac = torch.autograd.functional.jacobian(lambda a, b: cost(2 * a, b), (a, b)) + expected = ( + torch.tensor([-2 * torch.sin(2 * a), 2 * torch.sin(2 * a) * torch.sin(b)]), + torch.tensor([0, -torch.cos(2 * a) * torch.cos(b)]), + ) + assert isinstance(jac, tuple) and len(jac) == 2 + for _j, _e in zip(jac, expected): + assert torch.allclose(_j, _e, atol=atol_for_shots(shots), rtol=0) + + def test_classical_processing(self, execute_kwargs, device, shots): + """Test classical processing within the quantum tape""" + a = torch.tensor(0.1, requires_grad=True) + b = torch.tensor(0.2, requires_grad=False) + c = torch.tensor(0.3, requires_grad=True) + + def cost(a, b, c): + ops = [ + qml.RY(a * c, wires=0), + qml.RZ(b, wires=0), + qml.RX(c + c**2 + torch.sin(a), wires=0), + ] - res.backward() + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))], shots=shots) + return execute([tape], device, **execute_kwargs)[0] - assert isinstance(params.grad, torch.Tensor) - assert params.shape == (2,) + # PyTorch docs suggest a lambda for cost functions with some non-trainable args + # See for more: https://pytorch.org/docs/stable/autograd.html#functional-higher-level-api + res = torch.autograd.functional.jacobian(lambda _a, _c: cost(_a, b, _c), (a, c)) - def test_no_trainable_parameters(self, torch_device, execute_kwargs): - """Test evaluation and Jacobian if there are no trainable parameters""" - dev = qml.device("default.qubit.legacy", wires=2) + # Only two arguments are trainable + assert isinstance(res, tuple) and len(res) == 2 + if not shots.has_partitioned_shots: + assert res[0].shape == () + assert res[1].shape == () - with qml.queuing.AnnotatedQueue() as q: - qml.RY(0.2, wires=0) - qml.RX(torch.tensor(0.1, device=torch_device), wires=0) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.expval(qml.PauliZ(1)) + # I tried getting analytic results for this circuit but I kept being wrong and am giving up - tape = qml.tape.QuantumScript.from_queue(q) + @pytest.mark.skip("torch handles gradients and jacobians differently") + def test_no_trainable_parameters(self, execute_kwargs, shots, device): + """Test evaluation and Jacobian if there are no trainable parameters""" + a = torch.tensor(0.1, requires_grad=False) + b = torch.tensor(0.2, requires_grad=False) - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [] + def cost(a, b): + ops = [qml.RY(a, 0), qml.RX(b, 0), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))] + tape = qml.tape.QuantumScript(ops, m, shots=shots) + return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - assert isinstance(res, tuple) - assert len(res) == 2 + res = cost(a, b) + assert res.shape == (2,) - assert isinstance(res[0], torch.Tensor) - assert res[0].shape == () + res = torch.autograd.functional.jacobian(cost, (a, b)) + assert len(res) == 0 - assert isinstance(res[1], torch.Tensor) - assert res[1].shape == () + def loss(a, b): + return torch.sum(cost(a, b)) - with pytest.raises( - RuntimeError, - match="element 0 of tensors does not require grad and does not have a grad_fn", - ): - res[0].backward() + res = loss(a, b) + res.backward() - with pytest.raises( - RuntimeError, - match="element 0 of tensors does not require grad and does not have a grad_fn", - ): - res[1].backward() + assert torch.allclose(torch.tensor([a.grad, b.grad]), 0) - @pytest.mark.parametrize( - "U", [torch.tensor([[0.0, 1.0], [1.0, 0.0]]), np.array([[0.0, 1.0], [1.0, 0.0]])] - ) - def test_matrix_parameter(self, torch_device, U, execute_kwargs, tol): + def test_matrix_parameter(self, execute_kwargs, device, shots): """Test that the torch interface works correctly with a matrix parameter""" - a_val = 0.1 - a = torch.tensor(a_val, requires_grad=True, device=torch_device) - - if isinstance(U, torch.Tensor) and torch_device is not None: - U = U.to(torch_device) - - dev = qml.device("default.qubit.legacy", wires=2) - - with qml.queuing.AnnotatedQueue() as q: - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - qml.expval(qml.PauliZ(0)) + U = torch.tensor([[0, 1], [1, 0]], requires_grad=False, dtype=torch.float64) + a = torch.tensor(0.1, requires_grad=True) - tape = qml.tape.QuantumScript.from_queue(q) + def cost(a, U): + ops = [qml.QubitUnitary(U, wires=0), qml.RY(a, wires=0)] + tape = qml.tape.QuantumScript(ops, [qml.expval(qml.PauliZ(0))]) + return execute([tape], device, **execute_kwargs)[0] - res = execute([tape], dev, **execute_kwargs)[0] - assert tape.trainable_params == [1] + res = cost(a, U) + assert torch.allclose(res, -torch.cos(a), atol=atol_for_shots(shots)) - expected = torch.tensor(-np.cos(a_val), dtype=res.dtype, device=torch_device) - assert torch.allclose(res.detach(), expected, atol=tol, rtol=0) + jac = torch.autograd.functional.jacobian(lambda y: cost(y, U), a) + assert isinstance(jac, torch.Tensor) + assert torch.allclose(jac, torch.sin(a), atol=atol_for_shots(shots), rtol=0) - res.backward() - expected = torch.tensor([np.sin(a_val)], dtype=a.grad.dtype, device=torch_device) - assert torch.allclose(a.grad, expected, atol=tol, rtol=0) - - def test_differentiable_expand(self, torch_device, execute_kwargs, tol): - """Test that operation and nested tape expansion + def test_differentiable_expand(self, execute_kwargs, device, shots): + """Test that operation and nested tapes expansion is differentiable""" class U3(qml.U3): + """Dummy operator.""" + def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -716,519 +507,244 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - dev = qml.device("default.qubit.legacy", wires=1) - a = np.array(0.1) - p_val = [0.1, 0.2, 0.3] - p = torch.tensor(p_val, requires_grad=True, device=torch_device) - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(a, wires=0) - U3(p[0], p[1], p[2], wires=0) - qml.expval(qml.PauliX(0)) - - tape = qml.tape.QuantumScript.from_queue(q) + def cost_fn(a, p): + tape = qml.tape.QuantumScript( + [qml.RX(a, wires=0), U3(*p, wires=0)], [qml.expval(qml.PauliX(0))] + ) + gradient_fn = execute_kwargs["gradient_fn"] + if gradient_fn is None: + _gradient_method = None + elif isinstance(gradient_fn, str): + _gradient_method = gradient_fn + else: + _gradient_method = "gradient-transform" + config = qml.devices.ExecutionConfig( + interface="autograd", + gradient_method=_gradient_method, + grad_on_execution=execute_kwargs.get("grad_on_execution", None), + ) + program, _ = device.preprocess(execution_config=config) + return execute([tape], device, **execute_kwargs, transform_program=program)[0] - res = execute([tape], dev, **execute_kwargs)[0] + a = torch.tensor(0.1, requires_grad=False) + p = torch.tensor([0.1, 0.2, 0.3], requires_grad=True) - expected = torch.tensor( - np.cos(a) * np.cos(p_val[1]) * np.sin(p_val[0]) - + np.sin(a) - * ( - np.cos(p_val[2]) * np.sin(p_val[1]) - + np.cos(p_val[0]) * np.cos(p_val[1]) * np.sin(p_val[2]) - ), - dtype=res.dtype, - device=torch_device, + res = cost_fn(a, p) + expected = torch.cos(a) * torch.cos(p[1]) * torch.sin(p[0]) + torch.sin(a) * ( + torch.cos(p[2]) * torch.sin(p[1]) + torch.cos(p[0]) * torch.cos(p[1]) * torch.sin(p[2]) ) - assert torch.allclose(res.detach(), expected, atol=tol, rtol=0) + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res.backward() + res = torch.autograd.functional.jacobian(lambda _p: cost_fn(a, _p), p) expected = torch.tensor( [ - np.cos(p_val[1]) - * (np.cos(a) * np.cos(p_val[0]) - np.sin(a) * np.sin(p_val[0]) * np.sin(p_val[2])), - np.cos(p_val[1]) * np.cos(p_val[2]) * np.sin(a) - - np.sin(p_val[1]) - * (np.cos(a) * np.sin(p_val[0]) + np.cos(p_val[0]) * np.sin(a) * np.sin(p_val[2])), - np.sin(a) + torch.cos(p[1]) + * ( + torch.cos(a) * torch.cos(p[0]) + - torch.sin(a) * torch.sin(p[0]) * torch.sin(p[2]) + ), + torch.cos(p[1]) * torch.cos(p[2]) * torch.sin(a) + - torch.sin(p[1]) + * ( + torch.cos(a) * torch.sin(p[0]) + + torch.cos(p[0]) * torch.sin(a) * torch.sin(p[2]) + ), + torch.sin(a) * ( - np.cos(p_val[0]) * np.cos(p_val[1]) * np.cos(p_val[2]) - - np.sin(p_val[1]) * np.sin(p_val[2]) + torch.cos(p[0]) * torch.cos(p[1]) * torch.cos(p[2]) + - torch.sin(p[1]) * torch.sin(p[2]) ), - ], - dtype=p.grad.dtype, - device=torch_device, + ] ) - assert torch.allclose(p.grad, expected, atol=tol, rtol=0) + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - def test_probability_differentiation(self, torch_device, execute_kwargs, tol): + def test_probability_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") - - dev = qml.device("default.qubit.legacy", wires=2) - x_val = 0.543 - y_val = -0.654 - x = torch.tensor(x_val, requires_grad=True, device=torch_device) - y = torch.tensor(y_val, requires_grad=True, device=torch_device) - - def circuit(x, y): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=[0]) - qml.probs(wires=[1]) + def cost(x, y): + ops = [qml.RX(x, 0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.probs(wires=0), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, **execute_kwargs)[0] - - res = circuit(x, y) - - expected_0 = torch.tensor( - [np.cos(x_val / 2) ** 2, np.sin(x_val / 2) ** 2], - dtype=res[0].dtype, - device=torch_device, - ) + x = torch.tensor(0.543, requires_grad=True) + y = torch.tensor(-0.654, requires_grad=True) - expected_1 = torch.tensor( + res = cost(x, y) + expected = torch.tensor( [ - (1 + np.cos(x_val) * np.cos(y_val)) / 2, - (1 - np.cos(x_val) * np.cos(y_val)) / 2, - ], - dtype=res[0].dtype, - device=torch_device, + [ + torch.cos(x / 2) ** 2, + torch.sin(x / 2) ** 2, + (1 + torch.cos(x) * torch.cos(y)) / 2, + (1 - torch.cos(x) * torch.cos(y)) / 2, + ], + ] ) + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert torch.allclose(res[0], expected_0, atol=tol, rtol=0) - assert torch.allclose(res[1], expected_1, atol=tol, rtol=0) - - jac = torch_functional.jacobian(circuit, (x, y)) - dtype_jac = jac[0][0].dtype + res = torch.autograd.functional.jacobian(cost, (x, y)) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (4,) + assert res[1].shape == (4,) - res_0 = torch.tensor( - [-np.sin(x_val) / 2, np.sin(x_val) / 2], dtype=dtype_jac, device=torch_device - ) - res_1 = torch.tensor([0.0, 0.0], dtype=dtype_jac, device=torch_device) - res_2 = torch.tensor( - [-np.sin(x_val) * np.cos(y_val) / 2, np.cos(y_val) * np.sin(x_val) / 2], - dtype=dtype_jac, - device=torch_device, - ) - res_3 = torch.tensor( - [-np.cos(x_val) * np.sin(y_val) / 2, +np.cos(x_val) * np.sin(y_val) / 2], - dtype=dtype_jac, - device=torch_device, + expected = ( + torch.tensor( + [ + [ + -torch.sin(x) / 2, + torch.sin(x) / 2, + -torch.sin(x) * torch.cos(y) / 2, + torch.sin(x) * torch.cos(y) / 2, + ], + ] + ), + torch.tensor( + [ + [0, 0, -torch.cos(x) * torch.sin(y) / 2, torch.cos(x) * torch.sin(y) / 2], + ] + ), ) - assert torch.allclose(jac[0][0], res_0, atol=tol, rtol=0) - assert torch.allclose(jac[0][1], res_1, atol=tol, rtol=0) - assert torch.allclose(jac[1][0], res_2, atol=tol, rtol=0) - assert torch.allclose(jac[1][1], res_3, atol=tol, rtol=0) + assert torch.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) - def test_ragged_differentiation(self, torch_device, execute_kwargs, tol): + def test_ragged_differentiation(self, execute_kwargs, device, shots): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" - if execute_kwargs["gradient_fn"] == "device": - pytest.skip("Adjoint differentiation does not yet support probabilities") - dev = qml.device("default.qubit.legacy", wires=2) - x_val = 0.543 - y_val = -0.654 - x = torch.tensor(x_val, requires_grad=True, device=torch_device) - y = torch.tensor(y_val, requires_grad=True, device=torch_device) + def cost(x, y): + ops = [qml.RX(x, wires=0), qml.RY(y, 1), qml.CNOT((0, 1))] + m = [qml.expval(qml.PauliZ(0)), qml.probs(wires=1)] + tape = qml.tape.QuantumScript(ops, m) + return torch.hstack(execute([tape], device, **execute_kwargs)[0]) - def circuit(x, y): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - qml.probs(wires=[1]) + x = torch.tensor(0.543, requires_grad=True) + y = torch.tensor(-0.654, requires_grad=True) - tape = qml.tape.QuantumScript.from_queue(q) - - return execute([tape], dev, **execute_kwargs)[0] - - res = circuit(x, y) - - res_0 = torch.tensor(np.cos(x_val), dtype=res[0].dtype, device=torch_device) - res_1 = torch.tensor( - [(1 + np.cos(x_val) * np.cos(y_val)) / 2, (1 - np.cos(x_val) * np.cos(y_val)) / 2], - dtype=res[0].dtype, - device=torch_device, + res = cost(x, y) + expected = torch.tensor( + [ + torch.cos(x), + (1 + torch.cos(x) * torch.cos(y)) / 2, + (1 - torch.cos(x) * torch.cos(y)) / 2, + ] ) + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - assert isinstance(res, tuple) - assert len(res) == 2 - - assert torch.allclose(res[0], res_0, atol=tol, rtol=0) - assert torch.allclose(res[1], res_1, atol=tol, rtol=0) - - jac = torch_functional.jacobian(circuit, (x, y)) - dtype_jac = jac[0][0].dtype + res = torch.autograd.functional.jacobian(cost, (x, y)) + assert isinstance(res, tuple) and len(res) == 2 + assert res[0].shape == (3,) + assert res[1].shape == (3,) - res_0 = torch.tensor( - -np.sin(x_val), - dtype=dtype_jac, - device=torch_device, - ) - res_1 = torch.tensor( - 0.0, - dtype=dtype_jac, - device=torch_device, - ) - res_2 = torch.tensor( - [-np.sin(x_val) * np.cos(y_val) / 2, np.cos(y_val) * np.sin(x_val) / 2], - dtype=dtype_jac, - device=torch_device, - ) - res_3 = torch.tensor( - [-np.cos(x_val) * np.sin(y_val) / 2, +np.cos(x_val) * np.sin(y_val) / 2], - dtype=dtype_jac, - device=torch_device, + expected = ( + torch.tensor( + [-torch.sin(x), -torch.sin(x) * torch.cos(y) / 2, torch.sin(x) * torch.cos(y) / 2] + ), + torch.tensor([0, -torch.cos(x) * torch.sin(y) / 2, torch.cos(x) * torch.sin(y) / 2]), ) - - assert torch.allclose(jac[0][0], res_0, atol=tol, rtol=0) - assert torch.allclose(jac[0][1], res_1, atol=tol, rtol=0) - assert torch.allclose(jac[1][0], res_2, atol=tol, rtol=0) - assert torch.allclose(jac[1][1], res_3, atol=tol, rtol=0) - - def test_sampling(self, torch_device, execute_kwargs): - """Test sampling works as expected""" - # pylint: disable=unused-argument - if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True - ): - pytest.skip("Adjoint differentiation does not support samples") - if execute_kwargs["interface"] == "auto": - pytest.skip("Can't detect interface without a parametrized gate in the tape") - - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - - with qml.queuing.AnnotatedQueue() as q: - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.sample(qml.PauliZ(0)) - qml.sample(qml.PauliX(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, **execute_kwargs)[0] - - assert isinstance(res, tuple) - assert len(res) == 2 - - assert isinstance(res[0], torch.Tensor) - assert res[0].shape == (10,) - - assert isinstance(res[1], torch.Tensor) - assert res[1].shape == (10,) - - def test_sampling_expval(self, torch_device, execute_kwargs): - """Test sampling works as expected if combined with expectation values""" - # pylint: disable=unused-argument - if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True - ): - pytest.skip("Adjoint differentiation does not support samples") - if execute_kwargs["interface"] == "auto": - pytest.skip("Can't detect interface without a parametrized gate in the tape") - - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - - with qml.queuing.AnnotatedQueue() as q: - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - qml.sample(qml.PauliZ(0)) - qml.expval(qml.PauliX(1)) - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, **execute_kwargs)[0] - - assert len(res) == 2 - assert isinstance(res, tuple) - assert res[0].shape == (10,) - assert res[1].shape == () - assert isinstance(res[0], torch.Tensor) - assert isinstance(res[1], torch.Tensor) - - def test_sampling_gradient_error(self, torch_device, execute_kwargs): - """Test differentiating a tape with sampling results in an error""" - # pylint: disable=unused-argument - if ( - execute_kwargs["gradient_fn"] == "device" - and execute_kwargs["grad_on_execution"] is True - ): - pytest.skip("Adjoint differentiation does not support samples") - - dev = qml.device("default.qubit.legacy", wires=1, shots=10) - - x = torch.tensor(0.65, requires_grad=True) - - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x, wires=[0]) - qml.sample() - - tape = qml.tape.QuantumScript.from_queue(q) - - res = execute([tape], dev, **execute_kwargs)[0] - - with pytest.raises( - RuntimeError, - match="element 0 of tensors does not require grad and does not have a grad_fn", - ): - res.backward() - - def test_repeated_application_after_expand(self, torch_device, execute_kwargs): - """Test that the Torch interface continues to work after - tape expansions""" - # pylint: disable=unused-argument - n_qubits = 2 - dev = qml.device("default.qubit.legacy", wires=n_qubits) - - weights = torch.ones((3,)) - - with qml.queuing.AnnotatedQueue() as q: - qml.U3(*weights, wires=0) - qml.expval(qml.PauliZ(wires=0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - tape = tape.expand() - execute([tape], dev, **execute_kwargs) + assert torch.allclose(res[0], expected[0], atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[1], expected[1], atol=atol_for_shots(shots), rtol=0) -@pytest.mark.parametrize("torch_device", torch_devices) class TestHigherOrderDerivatives: """Test that the torch execute function can be differentiated""" @pytest.mark.parametrize( "params", [ - torch.tensor([0.543, -0.654], requires_grad=True), - torch.tensor([0, -0.654], requires_grad=True), - torch.tensor([-2.0, 0], requires_grad=True), + torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64), + torch.tensor([0, -0.654], requires_grad=True, dtype=torch.float64), + torch.tensor([-2.0, 0], requires_grad=True, dtype=torch.float64), ], ) - def test_parameter_shift_hessian(self, torch_device, params, tol): + def test_parameter_shift_hessian(self, params, tol): """Tests that the output of the parameter-shift transform can be differentiated using torch, yielding second derivatives.""" - # pylint: disable=unused-argument - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) + dev = DefaultQubit() def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[0], wires=0) - qml.RY(x[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) - - tape2 = qml.tape.QuantumScript.from_queue(q2) + ops1 = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops1, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - result = execute( - [tape1, tape2], dev, gradient_fn=param_shift, interface="torch", max_diff=2 - ) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=2) return result[0] + result[1][0] res = cost_fn(params) - x, y = params.detach() - expected = torch.as_tensor(0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))) + x, y = params.clone().detach() + expected = 0.5 * (3 + torch.cos(x) ** 2 * torch.cos(2 * y)) assert torch.allclose(res, expected, atol=tol, rtol=0) res.backward() expected = torch.tensor( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] + [-torch.cos(x) * torch.cos(2 * y) * torch.sin(x), -torch.cos(x) ** 2 * torch.sin(2 * y)] ) - assert torch.allclose(params.grad.detach(), expected, atol=tol, rtol=0) + assert torch.allclose(params.grad, expected, atol=tol, rtol=0) res = torch.autograd.functional.hessian(cost_fn, params) expected = torch.tensor( [ - [-np.cos(2 * x) * np.cos(2 * y), np.sin(2 * x) * np.sin(2 * y)], - [np.sin(2 * x) * np.sin(2 * y), -2 * np.cos(x) ** 2 * np.cos(2 * y)], + [-torch.cos(2 * x) * torch.cos(2 * y), torch.sin(2 * x) * torch.sin(2 * y)], + [torch.sin(2 * x) * torch.sin(2 * y), -2 * torch.cos(x) ** 2 * torch.cos(2 * y)], ] ) assert torch.allclose(res, expected, atol=tol, rtol=0) - def test_hessian_vector_valued(self, torch_device, tol): - """Test hessian calculation of a vector valued tape""" - dev = qml.device("default.qubit.legacy", wires=1) - - def circuit(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - qml.probs(wires=0) - - tape = qml.tape.QuantumScript.from_queue(q) - - return torch.stack( - execute([tape], dev, gradient_fn=param_shift, interface="torch", max_diff=2) - ) - - x = torch.tensor([1.0, 2.0], requires_grad=True, device=torch_device) - res = circuit(x) - - if torch_device is not None: - a, b = x.detach().cpu().numpy() - else: - a, b = x.detach().numpy() - - expected_res = torch.tensor( - [ - 0.5 + 0.5 * np.cos(a) * np.cos(b), - 0.5 - 0.5 * np.cos(a) * np.cos(b), - ], - dtype=res.dtype, - device=torch_device, - ) - assert torch.allclose(res.detach(), expected_res, atol=tol, rtol=0) - - def jac_fn(x): - return torch_functional.jacobian(circuit, x, create_graph=True) - - g = jac_fn(x) - - hess = torch_functional.jacobian(jac_fn, x) - - expected_g = torch.tensor( - [ - [-0.5 * np.sin(a) * np.cos(b), -0.5 * np.cos(a) * np.sin(b)], - [0.5 * np.sin(a) * np.cos(b), 0.5 * np.cos(a) * np.sin(b)], - ], - dtype=g.dtype, - device=torch_device, - ) - assert torch.allclose(g.detach(), expected_g, atol=tol, rtol=0) - - expected_hess = torch.tensor( - [ - [ - [-0.5 * np.cos(a) * np.cos(b), 0.5 * np.sin(a) * np.sin(b)], - [0.5 * np.sin(a) * np.sin(b), -0.5 * np.cos(a) * np.cos(b)], - ], - [ - [0.5 * np.cos(a) * np.cos(b), -0.5 * np.sin(a) * np.sin(b)], - [-0.5 * np.sin(a) * np.sin(b), 0.5 * np.cos(a) * np.cos(b)], - ], - ], - dtype=hess.dtype, - device=torch_device, - ) - assert torch.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - - def test_adjoint_hessian(self, torch_device, tol): - """Since the adjoint hessian is not a differentiable transform, - higher-order derivatives are not supported.""" - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor( - [0.543, -0.654], requires_grad=True, dtype=torch.float64, device=torch_device - ) - - def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - - tape = qml.tape.QuantumScript.from_queue(q) - - return execute( - [tape], - dev, - gradient_fn="device", - gradient_kwargs={"method": "adjoint_jacobian", "use_device_state": True}, - interface="torch", - )[0] - - res = torch.autograd.functional.hessian(cost_fn, params) - expected = torch.zeros([2, 2], dtype=torch.float64, device=torch_device) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - def test_max_diff(self, torch_device, tol): + def test_max_diff(self, tol): """Test that setting the max_diff parameter blocks higher-order derivatives""" - dev = qml.device("default.qubit.legacy", wires=2) - params = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) + dev = DefaultQubit() + params = torch.tensor([0.543, -0.654], requires_grad=True) def cost_fn(x): - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(x[0], wires=[0]) - qml.RY(x[1], wires=[1]) - qml.CNOT(wires=[0, 1]) - qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(x[0], wires=0) - qml.RY(x[0], wires=1) - qml.CNOT(wires=[0, 1]) - qml.probs(wires=1) + ops = [qml.RX(x[0], 0), qml.RY(x[1], 1), qml.CNOT((0, 1))] + tape1 = qml.tape.QuantumScript(ops, [qml.var(qml.PauliZ(0) @ qml.PauliX(1))]) - tape2 = qml.tape.QuantumScript.from_queue(q2) + ops2 = [qml.RX(x[0], 0), qml.RY(x[0], 1), qml.CNOT((0, 1))] + tape2 = qml.tape.QuantumScript(ops2, [qml.probs(wires=1)]) - result = execute( - [tape1, tape2], dev, gradient_fn=param_shift, max_diff=1, interface="torch" - ) + result = execute([tape1, tape2], dev, gradient_fn=param_shift, max_diff=1) return result[0] + result[1][0] res = cost_fn(params) - x, y = params.detach() - expected = torch.as_tensor(0.5 * (3 + np.cos(x) ** 2 * np.cos(2 * y))) - assert torch.allclose(res.to(torch_device), expected.to(torch_device), atol=tol, rtol=0) + x, y = params.clone().detach() + expected = 0.5 * (3 + torch.cos(x) ** 2 * torch.cos(2 * y)) + assert torch.allclose(res, expected, atol=tol, rtol=0) res.backward() expected = torch.tensor( - [-np.cos(x) * np.cos(2 * y) * np.sin(x), -np.cos(x) ** 2 * np.sin(2 * y)] - ) - assert torch.allclose( - params.grad.detach().to(torch_device), expected.to(torch_device), atol=tol, rtol=0 + [-torch.cos(x) * torch.cos(2 * y) * torch.sin(x), -torch.cos(x) ** 2 * torch.sin(2 * y)] ) + assert torch.allclose(params.grad, expected, atol=tol, rtol=0) res = torch.autograd.functional.hessian(cost_fn, params) - expected = torch.zeros([2, 2], dtype=torch.float64) - assert torch.allclose(res.to(torch_device), expected.to(torch_device), atol=tol, rtol=0) - - -execute_kwargs_hamiltonian = [ - {"gradient_fn": param_shift, "interface": "torch"}, - {"gradient_fn": finite_diff, "interface": "torch"}, -] + expected = torch.zeros([2, 2]) + assert torch.allclose(res, expected, atol=tol, rtol=0) -@pytest.mark.parametrize("execute_kwargs", execute_kwargs_hamiltonian) +@pytest.mark.parametrize("execute_kwargs, shots, device", test_matrix) +@pytest.mark.usefixtures("use_legacy_and_new_opmath") class TestHamiltonianWorkflows: """Test that tapes ending with expectations of Hamiltonians provide correct results and gradients""" @pytest.fixture - def cost_fn(self, execute_kwargs): + def cost_fn(self, execute_kwargs, shots, device): """Cost function for gradient tests""" - def _cost_fn(weights, coeffs1, coeffs2, dev=None): + def _cost_fn(weights, coeffs1, coeffs2): obs1 = [qml.PauliZ(0), qml.PauliZ(0) @ qml.PauliX(1), qml.PauliY(0)] H1 = qml.Hamiltonian(coeffs1, obs1) + if qml.operation.active_new_opmath(): + H1 = qml.pauli.pauli_sentence(H1).operation() obs2 = [qml.PauliZ(0)] H2 = qml.Hamiltonian(coeffs2, obs2) + if qml.operation.active_new_opmath(): + H2 = qml.pauli.pauli_sentence(H2).operation() with qml.queuing.AnnotatedQueue() as q: qml.RX(weights[0], wires=0) @@ -1237,72 +753,97 @@ def _cost_fn(weights, coeffs1, coeffs2, dev=None): qml.expval(H1) qml.expval(H2) - tape = qml.tape.QuantumScript.from_queue(q) - - return torch.hstack(execute([tape], dev, **execute_kwargs)[0]) + tape = qml.tape.QuantumScript.from_queue(q, shots=shots) + res = execute([tape], device, **execute_kwargs)[0] + if shots.has_partitioned_shots: + return torch.hstack(res[0] + res[1]) + return torch.hstack(res) return _cost_fn @staticmethod def cost_fn_expected(weights, coeffs1, coeffs2): """Analytic value of cost_fn above""" - a, b, c = coeffs1.detach().numpy() - d = coeffs2.detach().numpy()[0] - x, y = weights.detach().numpy() - return [-c * np.sin(x) * np.sin(y) + np.cos(x) * (a + b * np.sin(y)), d * np.cos(x)] + a, b, c = coeffs1.clone().detach() + d = coeffs2[0].clone().detach() + x, y = weights.clone().detach() + return torch.tensor( + [ + -c * torch.sin(x) * torch.sin(y) + torch.cos(x) * (a + b * torch.sin(y)), + d * torch.cos(x), + ] + ) @staticmethod def cost_fn_jacobian(weights, coeffs1, coeffs2): """Analytic jacobian of cost_fn above""" - a, b, c = coeffs1.detach().numpy() - d = coeffs2.detach().numpy()[0] - x, y = weights.detach().numpy() - return np.array( + a, b, c = coeffs1.clone().detach() + d = coeffs2[0].clone().detach() + x, y = weights.clone().detach() + return torch.tensor( [ [ - -c * np.cos(x) * np.sin(y) - np.sin(x) * (a + b * np.sin(y)), - b * np.cos(x) * np.cos(y) - c * np.cos(y) * np.sin(x), - np.cos(x), - np.cos(x) * np.sin(y), - -(np.sin(x) * np.sin(y)), + -c * torch.cos(x) * torch.sin(y) - torch.sin(x) * (a + b * torch.sin(y)), + b * torch.cos(x) * torch.cos(y) - c * torch.cos(y) * torch.sin(x), + torch.cos(x), + torch.cos(x) * torch.sin(y), + -(torch.sin(x) * torch.sin(y)), 0, ], - [-d * np.sin(x), 0, 0, 0, 0, np.cos(x)], + [-d * torch.sin(x), 0, 0, 0, 0, torch.cos(x)], ] ) - def test_multiple_hamiltonians_not_trainable(self, cost_fn, execute_kwargs, tol): - # pylint: disable=unused-argument - coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False, dtype=torch.float64) - coeffs2 = torch.tensor([0.7], requires_grad=False, dtype=torch.float64) - weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64) - dev = qml.device("default.qubit.legacy", wires=2) + def test_multiple_hamiltonians_not_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with no trainable parameters.""" + + if execute_kwargs["gradient_fn"] == "adjoint" and not qml.operation.active_new_opmath(): + pytest.skip("adjoint differentiation does not suppport hamiltonians.") + + coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=False) + coeffs2 = torch.tensor([0.7], requires_grad=False) + weights = torch.tensor([0.4, 0.5], requires_grad=True) - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) - assert np.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = torch.hstack( - torch_functional.jacobian(lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)) - ) - expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - assert np.allclose(res.detach(), expected, atol=tol, rtol=0) + res = torch.autograd.functional.jacobian(lambda w: cost_fn(w, coeffs1, coeffs2), weights) + expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2)[:, :2] + if shots.has_partitioned_shots: + assert torch.allclose(res[:2, :], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[2:, :], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) + + def test_multiple_hamiltonians_trainable(self, execute_kwargs, cost_fn, shots): + """Test hamiltonian with trainable parameters.""" + if execute_kwargs["gradient_fn"] == "adjoint": + pytest.skip("trainable hamiltonians not supported with adjoint") + if qml.operation.active_new_opmath(): + pytest.skip("parameter shift derivatives do not yet support sums.") - def test_multiple_hamiltonians_trainable(self, cost_fn, execute_kwargs, tol): - # pylint: disable=unused-argument - coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True, dtype=torch.float64) - coeffs2 = torch.tensor([0.7], requires_grad=True, dtype=torch.float64) - weights = torch.tensor([0.4, 0.5], requires_grad=True, dtype=torch.float64) - dev = qml.device("default.qubit.legacy", wires=2) + coeffs1 = torch.tensor([0.1, 0.2, 0.3], requires_grad=True) + coeffs2 = torch.tensor([0.7], requires_grad=True) + weights = torch.tensor([0.4, 0.5], requires_grad=True) - res = cost_fn(weights, coeffs1, coeffs2, dev=dev) + res = cost_fn(weights, coeffs1, coeffs2) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) - assert np.allclose(res[0].detach(), expected[0], atol=tol, rtol=0) - assert np.allclose(res[1].detach(), expected[1], atol=tol, rtol=0) + if shots.has_partitioned_shots: + assert torch.allclose(res[:2], expected, atol=atol_for_shots(shots), rtol=0) + assert torch.allclose(res[2:], expected, atol=atol_for_shots(shots), rtol=0) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) - res = torch.hstack( - torch_functional.jacobian(lambda *x: cost_fn(*x, dev=dev), (weights, coeffs1, coeffs2)) - ) + res = torch.hstack(torch.autograd.functional.jacobian(cost_fn, (weights, coeffs1, coeffs2))) expected = self.cost_fn_jacobian(weights, coeffs1, coeffs2) - assert np.allclose(res.detach(), expected, atol=tol, rtol=0) + if shots.has_partitioned_shots: + pytest.xfail( + "multiple hamiltonians with shot vectors does not seem to be differentiable." + ) + else: + assert torch.allclose(res, expected, atol=atol_for_shots(shots), rtol=0) diff --git a/tests/interfaces/test_torch_qnode.py b/tests/interfaces/test_torch_qnode.py index 3965255a4a4..2f1c85f6275 100644 --- a/tests/interfaces/test_torch_qnode.py +++ b/tests/interfaces/test_torch_qnode.py @@ -1,4 +1,4 @@ -# Copyright 2022 Xanadu Quantum Technologies Inc. +# Copyright 2023 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. @@ -12,13 +12,15 @@ # See the License for the specific language governing permissions and # limitations under the License. """Integration tests for using the Torch interface with a QNode""" -# pylint: disable=too-many-arguments,too-many-public-methods,too-few-public-methods, -# pylint: disable=use-dict-literal, use-implicit-booleaness-not-comparison, unnecessary-lambda-assignment +# pylint: disable=too-many-arguments,unexpected-keyword-arg,no-member,comparison-with-callable, no-name-in-module +# pylint: disable=use-implicit-booleaness-not-comparison, unnecessary-lambda-assignment, use-dict-literal import numpy as np import pytest import pennylane as qml from pennylane import qnode +from pennylane.devices import DefaultQubit +from tests.param_shift_dev import ParamShiftDerivativesDevice pytestmark = pytest.mark.torch @@ -26,14 +28,25 @@ jacobian = torch.autograd.functional.jacobian hessian = torch.autograd.functional.hessian +# device, diff_method, grad_on_execution, device_vjp qubit_device_and_diff_method = [ - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "spsa", False], - ["default.qubit.legacy", "hadamard", False], + [DefaultQubit(), "finite-diff", False, False], + [DefaultQubit(), "parameter-shift", False, False], + [DefaultQubit(), "backprop", True, False], + [DefaultQubit(), "adjoint", True, False], + [DefaultQubit(), "adjoint", False, False], + [DefaultQubit(), "adjoint", True, True], + [DefaultQubit(), "adjoint", False, True], + [DefaultQubit(), "spsa", False, False], + [DefaultQubit(), "hadamard", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", False, True], + [qml.device("lightning.qubit", wires=5), "adjoint", True, True], + [qml.device("lightning.qubit", wires=5), "adjoint", False, False], + [qml.device("lightning.qubit", wires=5), "adjoint", True, False], + [ParamShiftDerivativesDevice(), "parameter-shift", False, False], + [ParamShiftDerivativesDevice(), "best", False, False], + [ParamShiftDerivativesDevice(), "parameter-shift", True, False], + [ParamShiftDerivativesDevice(), "parameter-shift", False, True], ] interface_and_qubit_device_and_diff_method = [ @@ -46,25 +59,25 @@ @pytest.mark.parametrize( - "interface, dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface, dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQNode: """Test that using the QNode with Torch integrates with the PennyLane stack""" - def test_execution_with_interface(self, interface, dev_name, diff_method, grad_on_execution): + def test_execution_with_interface( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): """Test execution works with the interface""" if diff_method == "backprop": pytest.skip("Test does not support backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -79,7 +92,7 @@ def circuit(a): # with the interface, the tape returns torch tensors assert isinstance(res, torch.Tensor) - assert res.shape == tuple() + assert res.shape == () # the tape is able to deduce trainable parameters assert circuit.qtape.trainable_params == [0] @@ -89,20 +102,18 @@ def circuit(a): grad = a.grad assert isinstance(grad, torch.Tensor) - assert grad.shape == tuple() + assert grad.shape == () - def test_interface_swap(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_interface_swap(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test that the Torch interface can be applied to a QNode with a pre-existing interface""" - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, diff_method=diff_method, interface="autograd", grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface="autograd", + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a, wires=0) @@ -128,22 +139,19 @@ def circuit(a): assert np.allclose(res1, res2.detach().numpy(), atol=tol, rtol=0) assert np.allclose(grad1, grad2, atol=tol, rtol=0) - def test_drawing(self, interface, dev_name, diff_method, grad_on_execution): + def test_drawing(self, interface, dev, diff_method, grad_on_execution, device_vjp): """Test circuit drawing when using the torch interface""" x = torch.tensor(0.1, requires_grad=True) y = torch.tensor([0.2, 0.3], requires_grad=True) z = torch.tensor(0.4, requires_grad=True) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(p1, p2=y, **kwargs): qml.RX(p1, wires=0) @@ -159,14 +167,17 @@ def circuit(p1, p2=y, **kwargs): assert result == expected - def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_jacobian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test jacobian calculation""" kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) - if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA a_val = 0.1 @@ -175,13 +186,6 @@ def test_jacobian(self, interface, dev_name, diff_method, grad_on_execution, tol a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode(dev, **kwargs) def circuit(a, b): qml.RY(a, wires=0) @@ -217,19 +221,29 @@ def circuit(a, b): assert np.allclose(b.grad, expected[1], atol=tol, rtol=0) # TODO: fix this behavior with float: already present before return type. - @pytest.mark.xfail - def test_jacobian_dtype(self, interface, dev_name, diff_method, grad_on_execution): + def test_jacobian_dtype( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test calculating the jacobian with a different datatype""" + if not "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("Failing unless lightning.qubit") if diff_method == "backprop": pytest.skip("Test does not support backprop") a = torch.tensor(0.1, dtype=torch.float32, requires_grad=True) b = torch.tensor(0.2, dtype=torch.float32, requires_grad=True) - dev = qml.device(dev_name, wires=2) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -253,15 +267,20 @@ def circuit(a, b): assert a.grad.dtype is torch.float32 assert b.grad.dtype is torch.float32 - def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execution): + def test_jacobian_options( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test setting jacobian options""" if diff_method not in {"finite-diff", "spsa"}: pytest.skip("Test only works with finite-diff and spsa") a = torch.tensor([0.1, 0.2], requires_grad=True) - dev = qml.device(dev_name, wires=1) - @qnode( dev, diff_method=diff_method, @@ -269,6 +288,7 @@ def test_jacobian_options(self, interface, dev_name, diff_method, grad_on_execut interface=interface, h=1e-8, approx_order=2, + device_vjp=device_vjp, ) def circuit(a): qml.RY(a[0], wires=0) @@ -278,7 +298,9 @@ def circuit(a): res = circuit(a) res.backward() - def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_changing_trainability( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that changing the trainability of parameters changes the number of differentiation requests made""" if diff_method != "parameter-shift": @@ -290,10 +312,12 @@ def test_changing_trainability(self, interface, dev_name, diff_method, grad_on_e a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) b = torch.tensor(b_val, dtype=torch.float64, requires_grad=True) - dev = qml.device(dev_name, wires=2) - @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -342,21 +366,25 @@ def circuit(a, b): expected = -np.sin(a_val) + np.sin(a_val) * np.sin(b_val) assert np.allclose(a.grad, expected, atol=tol, rtol=0) - def test_classical_processing(self, interface, dev_name, diff_method, grad_on_execution): + def test_classical_processing( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test classical processing within the quantum tape""" a = torch.tensor(0.1, dtype=torch.float64, requires_grad=True) b = torch.tensor(0.2, dtype=torch.float64, requires_grad=False) c = torch.tensor(0.3, dtype=torch.float64, requires_grad=True) - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(a, b, c): qml.RY(a * c, wires=0) @@ -376,17 +404,22 @@ def circuit(a, b, c): assert b.grad is None assert isinstance(c.grad, torch.Tensor) - def test_no_trainable_parameters(self, interface, dev_name, diff_method, grad_on_execution): + def test_no_trainable_parameters( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test evaluation and Jacobian if there are no trainable parameters""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(a, b): qml.RY(a, wires=0) @@ -430,21 +463,20 @@ def circuit(a, b): np.array([[0, 1], [1, 0]]), ], ) - def test_matrix_parameter(self, interface, dev_name, diff_method, grad_on_execution, U, tol): + def test_matrix_parameter( + self, interface, dev, diff_method, grad_on_execution, device_vjp, U, tol + ): """Test that the Torch interface works correctly with a matrix parameter""" a_val = 0.1 a = torch.tensor(a_val, dtype=torch.float64, requires_grad=True) - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( - dev, diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + dev, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) def circuit(U, a): qml.QubitUnitary(U, wires=0) @@ -461,18 +493,23 @@ def circuit(U, a): res.backward() assert np.allclose(a.grad, np.sin(a_val), atol=tol, rtol=0) - def test_differentiable_expand(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_differentiable_expand( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that operation and nested tapes expansion is differentiable""" kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) if diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - class U3(qml.U3): + class U3(qml.U3): # pylint:disable=too-few-public-methods def decomposition(self): theta, phi, lam = self.data wires = self.wires @@ -481,12 +518,6 @@ def decomposition(self): qml.PhaseShift(phi + lam, wires=wires), ] - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) a = np.array(0.1) p_val = [0.1, 0.2, 0.3] p = torch.tensor(p_val, dtype=torch.float64, requires_grad=True) @@ -527,9 +558,9 @@ class TestShotsIntegration: """Test that the QNode correctly changes shot value, and differentiates it.""" - def test_changing_shots(self, mocker, tol): + def test_changing_shots(self): """Test that changing shots works on execution""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = torch.tensor([0.543, -0.654], requires_grad=True, dtype=torch.float64) @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) @@ -537,33 +568,22 @@ def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliY(1)) - - spy = mocker.spy(dev, "sample") + return qml.sample(wires=(0, 1)) # execute with device default shots (None) - res = circuit(a, b) - assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0) - spy.assert_not_called() + with pytest.raises(qml.DeviceError): + circuit(a, b) # execute with shots=100 - res = circuit(a, b, shots=100) # pylint: disable=unexpected-keyword-arg - spy.assert_called_once() - assert spy.spy_return.shape == (100,) - - # device state has been unaffected - assert dev.shots is None - res = circuit(a, b) - assert torch.allclose(res, -torch.cos(a) * torch.sin(b), atol=tol, rtol=0) - spy.assert_called_once() # only same call as above + res = circuit(a, b, shots=100) + assert res.shape == (100, 2) # TODO: add this test after shot vectors addition @pytest.mark.xfail def test_gradient_integration(self): """Test that temporarily setting the shots works for gradient computations""" - # pylint: disable=unexpected-keyword-arg - dev = qml.device("default.qubit.legacy", wires=2, shots=None) + dev = DefaultQubit() a, b = torch.tensor([0.543, -0.654], requires_grad=True) @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) @@ -585,11 +605,10 @@ def test_multiple_gradient_integration(self, tol): """Test that temporarily setting the shots works for gradient computations, even if the QNode has been re-evaluated with a different number of shots in the meantime.""" - dev = qml.device("default.qubit.legacy", wires=2, shots=None) weights = torch.tensor([0.543, -0.654], requires_grad=True) a, b = weights - @qnode(dev, interface="torch", diff_method=qml.gradients.param_shift) + @qnode(DefaultQubit(), interface="torch", diff_method=qml.gradients.param_shift) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) @@ -597,7 +616,7 @@ def circuit(a, b): return qml.expval(qml.PauliY(1)) res1 = circuit(*weights) - assert qml.math.shape(res1) == tuple() + assert qml.math.shape(res1) == () res2 = circuit(*weights, shots=[(1, 1000)]) assert len(res2) == 1000 @@ -609,120 +628,47 @@ def circuit(a, b): def test_update_diff_method(self, mocker): """Test that temporarily setting the shots updates the diff method""" - dev = qml.device("default.qubit.legacy", wires=2, shots=100) a, b = torch.tensor([0.543, -0.654], requires_grad=True) spy = mocker.spy(qml, "execute") - @qnode(dev, interface="torch") + @qnode(DefaultQubit(), interface="torch") def cost_fn(a, b): qml.RY(a, wires=0) qml.RX(b, wires=1) qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) + cost_fn(a, b, shots=100) # since we are using finite shots, parameter-shift will # be chosen - assert cost_fn.gradient_fn is qml.gradients.param_shift - - cost_fn(a, b) + assert cost_fn.gradient_fn == qml.gradients.param_shift assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - # if we set the shots to None, backprop can now be used - cost_fn(a, b, shots=None) # pylint: disable=unexpected-keyword-arg - assert spy.call_args[1]["gradient_fn"] == "backprop" - assert cost_fn.gradient_fn == "backprop" - - # original QNode settings are unaffected - + # if we use the default shots value of None, backprop can now be used cost_fn(a, b) - assert cost_fn.gradient_fn is qml.gradients.param_shift - assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - - -class TestAdjoint: - """Specific integration tests for the adjoint method""" - - def test_reuse_state(self, mocker): - """Tests that the Torch interface reuses the device state for adjoint differentiation""" - dev = qml.device("default.qubit.legacy", wires=2) - - @qnode(dev, diff_method="adjoint", interface="torch") - def circ(x): - qml.RX(x[0], wires=0) - qml.RY(x[1], wires=1) - qml.CNOT(wires=(0, 1)) - return qml.expval(qml.PauliZ(0)) - - expected_grad = lambda x: torch.tensor([-torch.sin(x[0]), torch.cos(x[1])]) - - spy = mocker.spy(dev, "adjoint_jacobian") - - x1 = torch.tensor([1.0, 1.0], requires_grad=True) - res = circ(x1) - res.backward() - - assert np.allclose(x1.grad[0], expected_grad(x1)[0]) - assert circ.device.num_executions == 1 - spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) - - def test_resuse_state_multiple_evals(self, mocker, tol): - """Tests that the Torch interface reuses the device state for adjoint differentiation, - even where there are intermediate evaluations.""" - dev = qml.device("default.qubit.legacy", wires=2) - - x_val = 0.543 - y_val = -0.654 - x = torch.tensor(x_val, requires_grad=True) - y = torch.tensor(y_val, requires_grad=True) - - @qnode(dev, diff_method="adjoint", interface="torch") - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)) - - spy = mocker.spy(dev, "adjoint_jacobian") - - res1 = circuit(x, y) - assert np.allclose(res1.detach(), np.cos(x_val), atol=tol, rtol=0) - - # intermediate evaluation with different values - circuit(torch.tan(x), torch.cosh(y)) - - # the adjoint method will continue to compute the correct derivative - res1.backward() - assert np.allclose(x.grad.detach(), -np.sin(x_val), atol=tol, rtol=0) - assert dev.num_executions == 2 - spy.assert_called_with(mocker.ANY, use_device_state=mocker.ANY) + assert spy.call_args[1]["gradient_fn"] == "backprop" @pytest.mark.parametrize( - "interface,dev_name,diff_method,grad_on_execution", interface_and_qubit_device_and_diff_method + "interface,dev,diff_method,grad_on_execution, device_vjp", + interface_and_qubit_device_and_diff_method, ) class TestQubitIntegration: """Tests that ensure various qubit circuits integrate correctly""" def test_probability_differentiation( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measureing probabilities with adjoint.") kwargs = {} - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - x_val = 0.543 y_val = -0.654 x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64) @@ -733,6 +679,7 @@ def test_probability_differentiation( diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface, + device_vjp=device_vjp, **kwargs, ) def circuit(x, y): @@ -768,25 +715,24 @@ def circuit(x, y): assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0) assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0) - def test_ragged_differentiation(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_ragged_differentiation( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Tests correct output shape and evaluation for a tape with prob and expval outputs""" + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("lightning does not support measureing probabilities with adjoint.") kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + diff_method=diff_method, + grad_on_execution=grad_on_execution, + interface=interface, + device_vjp=device_vjp, ) - if diff_method == "adjoint": - pytest.skip("The adjoint method does not currently support returning probabilities") - elif diff_method == "spsa": + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - x_val = 0.543 y_val = -0.654 x = torch.tensor(x_val, requires_grad=True, dtype=torch.float64) @@ -824,17 +770,22 @@ def circuit(x, y): assert np.allclose(jac[1][0], res_2, atol=tol, rtol=0) assert np.allclose(jac[1][1], res_3, atol=tol, rtol=0) - def test_chained_qnodes(self, interface, dev_name, diff_method, grad_on_execution): + def test_chained_qnodes( + self, + interface, + dev, + diff_method, + grad_on_execution, + device_vjp, + ): """Test that the gradient of chained QNodes works without error""" - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) @qnode( - dev, interface=interface, diff_method=diff_method, grad_on_execution=grad_on_execution + dev, + interface=interface, + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit1(weights): qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1]) @@ -866,18 +817,11 @@ def cost(weights): loss = cost(weights) loss.backward() - def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test hessian calculation of a scalar valued QNode""" - if diff_method in {"adjoint", "spsa"}: + if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": pytest.skip("Adjoint and SPSA do not support second derivative.") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -888,6 +832,7 @@ def test_hessian(self, interface, dev_name, diff_method, grad_on_execution, tol) grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, **options, ) def circuit(x): @@ -923,18 +868,13 @@ def circuit(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian_vector_valued( + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test hessian calculation of a vector valued QNode""" - if diff_method in {"adjoint", "spsa"}: + if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": pytest.skip("Adjoint and SPSA do not support second derivative.") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -945,6 +885,7 @@ def test_hessian_vector_valued(self, interface, dev_name, diff_method, grad_on_e grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, **options, ) def circuit(x): @@ -991,18 +932,11 @@ def circuit(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_hessian_ragged(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test hessian calculation of a ragged QNode""" - if diff_method in {"adjoint", "spsa"}: + if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": pytest.skip("Adjoint and SPSA do not support second derivative.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires) - options = {} if diff_method == "finite-diff": options = {"h": 1e-6} @@ -1013,6 +947,7 @@ def test_hessian_ragged(self, interface, dev_name, diff_method, grad_on_executio grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, **options, ) def circuit(x): @@ -1071,29 +1006,23 @@ def circuit_stack(x): assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) def test_hessian_vector_valued_postprocessing( - self, interface, dev_name, diff_method, grad_on_execution, tol + self, interface, dev, diff_method, grad_on_execution, device_vjp, tol ): """Test hessian calculation of a vector valued QNode with post-processing""" - if diff_method in {"adjoint", "spsa"}: + if diff_method in {"adjoint", "spsa"} or dev.name == "param_shift.qubit": pytest.skip("Adjoint and SPSA do not support second derivative.") options = {} if diff_method == "finite-diff": options = {"h": 1e-6} - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) - @qnode( dev, diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, interface=interface, + device_vjp=device_vjp, **options, ) def circuit(x): @@ -1138,23 +1067,23 @@ def cost_fn(x): else: assert np.allclose(hess.detach(), expected_hess, atol=tol, rtol=0) - def test_state(self, interface, dev_name, diff_method, grad_on_execution, tol): + def test_state(self, interface, dev, diff_method, grad_on_execution, device_vjp, tol): """Test that the state can be returned and differentiated""" - if diff_method == "adjoint": - pytest.skip("Adjoint does not support states") - - num_wires = 2 - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires) + if dev.name == "param_shift.qubit": + pytest.skip("parameter shift does not support measuring the state.") + if "lightning" in getattr(dev, "name", "").lower(): + pytest.xfail("Lightning devices do not support state with adjoint diff.") x = torch.tensor(0.543, requires_grad=True) y = torch.tensor(-0.654, requires_grad=True) @qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -1164,7 +1093,7 @@ def circuit(x, y): def cost_fn(x, y): res = circuit(x, y) - assert res.dtype is torch.complex128 + assert torch.is_complex(res) probs = torch.abs(res) ** 2 return probs[0] + probs[2] @@ -1181,20 +1110,25 @@ def cost_fn(x, y): assert torch.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.parametrize("state", [[1], [0, 1]]) # Basis state and state vector - def test_projector(self, state, interface, dev_name, diff_method, grad_on_execution, tol): + def test_projector( + self, state, interface, dev, diff_method, grad_on_execution, device_vjp, tol + ): """Test that the variance of a projector is correctly returned""" kwargs = dict( - diff_method=diff_method, grad_on_execution=grad_on_execution, interface=interface + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) if diff_method == "adjoint": - pytest.skip("Adjoint does not support projectors") - elif diff_method == "spsa": + pytest.skip("adjoint supports either all expvals or all diagonal measurements") + if diff_method == "spsa": kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("Hadamard does not support variances.") - dev = qml.device(dev_name, wires=2) P = torch.tensor(state, requires_grad=False) x, y = 0.765, -0.654 @@ -1222,99 +1156,64 @@ def circuit(x, y): ) assert np.allclose(weights.grad.detach(), expected, atol=tol, rtol=0) + def test_postselection_differentiation( + self, interface, dev, diff_method, grad_on_execution, device_vjp + ): + """Test that when postselecting with default.qubit, differentiation works correctly.""" -@pytest.mark.parametrize( - "diff_method,kwargs", - [ - ["finite-diff", {}], - ["spsa", {"num_directions": 100, "h": 0.05}], - ("parameter-shift", {}), - ("parameter-shift", {"force_order2": True}), - ], -) -class TestCV: - """Tests for CV integration""" - - def test_first_order_observable(self, diff_method, kwargs, tol): - """Test variance of a first order CV observable""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA - - r = torch.tensor(0.543, dtype=torch.float64, requires_grad=True) - phi = torch.tensor(-0.654, dtype=torch.float64, requires_grad=True) - - @qnode(dev, interface="torch", diff_method=diff_method, **kwargs) - def circuit(r, phi): - qml.Squeezing(r, 0, wires=0) - qml.Rotation(phi, wires=0) - return qml.var(qml.QuadX(0)) - - res = circuit(r, phi) - expected = torch.exp(2 * r) * torch.sin(phi) ** 2 + torch.exp(-2 * r) * torch.cos(phi) ** 2 - assert torch.allclose(res, expected, atol=tol, rtol=0) + if diff_method in ["adjoint", "spsa", "hadamard"]: + pytest.skip("Diff method does not support postselection.") - # circuit jacobians - res.backward() - res = torch.tensor([r.grad, phi.grad]) - expected = torch.tensor( - [ - [ - 2 * torch.exp(2 * r) * torch.sin(phi) ** 2 - - 2 * torch.exp(-2 * r) * torch.cos(phi) ** 2, - 2 * torch.sinh(2 * r) * torch.sin(2 * phi), - ] - ] + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) - assert torch.allclose(res, expected, atol=tol, rtol=0) - - def test_second_order_observable(self, diff_method, kwargs, tol): - """Test variance of a second order CV expectation value""" - dev = qml.device("default.gaussian", wires=1) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) - tol = TOL_FOR_SPSA + def circuit(phi, theta): + qml.RX(phi, wires=0) + qml.CNOT([0, 1]) + qml.measure(wires=0, postselect=1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) + + @qml.qnode( + dev, + diff_method=diff_method, + interface=interface, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) + def expected_circuit(theta): + qml.PauliX(1) + qml.RX(theta, wires=1) + return qml.expval(qml.PauliZ(1)) - n = torch.tensor(0.12, dtype=torch.float64, requires_grad=True) - a = torch.tensor(0.765, dtype=torch.float64, requires_grad=True) + phi = torch.tensor(1.23, requires_grad=True) + theta = torch.tensor(4.56, requires_grad=True) - @qnode(dev, interface="torch", diff_method=diff_method, **kwargs) - def circuit(n, a): - qml.ThermalState(n, wires=0) - qml.Displacement(a, 0, wires=0) - return qml.var(qml.NumberOperator(0)) + assert qml.math.allclose(circuit(phi, theta), expected_circuit(theta)) - res = circuit(n, a) - expected = n**2 + n + torch.abs(a) ** 2 * (1 + 2 * n) - assert torch.allclose(res, expected, atol=tol, rtol=0) + gradient = torch.autograd.grad(circuit(phi, theta), [phi, theta]) + exp_theta_grad = torch.autograd.grad(expected_circuit(theta), theta)[0] + assert qml.math.allclose(gradient, [0.0, exp_theta_grad]) - # circuit jacobians - res.backward() - res = torch.tensor([n.grad, a.grad]) - expected = torch.tensor([[2 * a**2 + 2 * n + 1, 2 * a * (2 * n + 1)]]) - assert torch.allclose(res, expected, atol=tol, rtol=0) - -@pytest.mark.parametrize("dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method) +@pytest.mark.parametrize( + "dev,diff_method,grad_on_execution, device_vjp", qubit_device_and_diff_method +) class TestTapeExpansion: """Test that tape expansion within the QNode integrates correctly with the Torch interface""" - def test_gradient_expansion(self, dev_name, diff_method, grad_on_execution): + def test_gradient_expansion(self, dev, diff_method, grad_on_execution, device_vjp): """Test that a *supported* operation with no gradient recipe is expanded for both parameter-shift and finite-differences, but not for execution.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - class PhaseShift(qml.PhaseShift): + class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods grad_method = None def decomposition(self): @@ -1325,6 +1224,7 @@ def decomposition(self): diff_method=diff_method, grad_on_execution=grad_on_execution, max_diff=2, + device_vjp=device_vjp, interface="torch", ) def circuit(x): @@ -1335,33 +1235,32 @@ def circuit(x): x = torch.tensor(0.5, requires_grad=True, dtype=torch.float64) loss = circuit(x) + loss.backward() res = x.grad assert torch.allclose(res, -3 * torch.sin(3 * x)) - if diff_method == "parameter-shift": + if diff_method == "parameter-shift" and dev.name != "param_shift.qubit": # test second order derivatives res = torch.autograd.functional.hessian(circuit, x) assert torch.allclose(res, -9 * torch.cos(3 * x)) @pytest.mark.parametrize("max_diff", [1, 2]) def test_gradient_expansion_trainable_only( - self, dev_name, diff_method, grad_on_execution, max_diff + self, + dev, + diff_method, + grad_on_execution, + device_vjp, + max_diff, ): """Test that a *supported* operation with no gradient recipe is only expanded for parameter-shift and finite-differences when it is trainable.""" if diff_method not in ("parameter-shift", "finite-diff", "spsa", "hadamard"): pytest.skip("Only supports gradient transforms") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires) - - class PhaseShift(qml.PhaseShift): + class PhaseShift(qml.PhaseShift): # pylint:disable=too-few-public-methods grad_method = None def decomposition(self): @@ -1373,6 +1272,7 @@ def decomposition(self): grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", + device_vjp=device_vjp, ) def circuit(x, y): qml.Hadamard(wires=0) @@ -1388,7 +1288,7 @@ def circuit(x, y): @pytest.mark.parametrize("max_diff", [1, 2]) def test_hamiltonian_expansion_analytic( - self, dev_name, diff_method, grad_on_execution, max_diff, tol + self, dev, diff_method, grad_on_execution, max_diff, device_vjp, tol ): """Test that if there are non-commuting groups and the number of shots is None @@ -1398,17 +1298,17 @@ def test_hamiltonian_expansion_analytic( grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", + device_vjp=device_vjp, ) if diff_method == "adjoint": pytest.skip("The adjoint method does not yet support Hamiltonians") elif diff_method == "spsa": - spsa_kwargs = dict(sampler_rng=np.random.default_rng(SEED_FOR_SPSA), num_directions=10) - kwargs = {**kwargs, **spsa_kwargs} + kwargs["sampler_rng"] = np.random.default_rng(SEED_FOR_SPSA) + kwargs["num_directions"] = 20 tol = TOL_FOR_SPSA elif diff_method == "hadamard": pytest.skip("The hadamard method does not yet support Hamiltonians") - dev = qml.device(dev_name, wires=3, shots=None) obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode(dev, **kwargs) @@ -1448,7 +1348,11 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method in ("parameter-shift", "backprop") and max_diff == 2: + if ( + diff_method in ("parameter-shift", "backprop") + and max_diff == 2 + and dev.name != "param_shift.qubit" + ): hessians = torch.autograd.functional.hessian(circuit, (d, w, c)) grad2_c = hessians[2][2] @@ -1468,14 +1372,14 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad2_w_c, expected, atol=tol) @pytest.mark.parametrize("max_diff", [1, 2]) - def test_hamiltonian_expansion_finite_shots( - self, dev_name, diff_method, grad_on_execution, max_diff, mocker + def test_hamiltonian_finite_shots( + self, dev, diff_method, device_vjp, grad_on_execution, max_diff ): - """Test that the Hamiltonian is expanded if there + """Test that the Hamiltonian is correctly measured if there are non-commuting groups and the number of shots is finite and the first and second order gradients are correctly evaluated""" gradient_kwargs = {} - tol = 0.3 + tol = 0.1 if diff_method in ("adjoint", "backprop"): pytest.skip("The adjoint and backprop methods do not yet support sampling") elif diff_method == "spsa": @@ -1487,11 +1391,10 @@ def test_hamiltonian_expansion_finite_shots( tol = TOL_FOR_SPSA elif diff_method == "finite-diff": gradient_kwargs = {"h": 0.05} + tol = 0.15 elif diff_method == "hadamard": pytest.skip("The hadamard method does not yet support Hamiltonians") - dev = qml.device(dev_name, wires=3, shots=50000) - spy = mocker.spy(qml.transforms, "split_non_commuting") obs = [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1), qml.PauliZ(0) @ qml.PauliZ(1)] @qnode( @@ -1500,6 +1403,7 @@ def test_hamiltonian_expansion_finite_shots( grad_on_execution=grad_on_execution, max_diff=max_diff, interface="torch", + device_vjp=device_vjp, **gradient_kwargs, ) def circuit(data, weights, coeffs): @@ -1514,15 +1418,17 @@ def circuit(data, weights, coeffs): c = torch.tensor([-0.6543, 0.24, 0.54], requires_grad=True, dtype=torch.float64) # test output - res = circuit(d, w, c) + res = circuit(d, w, c, shots=50000) expected = c[2] * torch.cos(d[1] + w[1]) - c[1] * torch.sin(d[0] + w[0]) * torch.sin( d[1] + w[1] ) assert torch.allclose(res, expected, atol=tol) - spy.assert_called() # test gradients + if diff_method in ["finite-diff", "spsa"]: + pytest.skip(f"{diff_method} not compatible") + res.backward() grad = (w.grad, c.grad) @@ -1541,8 +1447,10 @@ def circuit(data, weights, coeffs): assert torch.allclose(grad[1], expected_c, atol=tol) # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2: - hessians = torch.autograd.functional.hessian(circuit, (d, w, c)) + if diff_method == "parameter-shift" and max_diff == 2 and dev.name != "param_shift.qubit": + hessians = torch.autograd.functional.hessian( + lambda _d, _w, _c: circuit(_d, _w, _c, shots=50000), (d, w, c) + ) grad2_c = hessians[2][2] assert torch.allclose(grad2_c, torch.zeros([3, 3], dtype=torch.float64), atol=tol) @@ -1566,15 +1474,14 @@ class TestSample: def test_sample_dimension(self): """Test sampling works as expected""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert isinstance(res, tuple) assert len(res) == 2 @@ -1587,37 +1494,33 @@ def circuit(): def test_sampling_expval(self): """Test sampling works as expected if combined with expectation values""" - shots = 10 - dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert len(res) == 2 assert isinstance(res, tuple) assert isinstance(res[0], torch.Tensor) - assert res[0].shape == (shots,) + assert res[0].shape == (10,) assert isinstance(res[1], torch.Tensor) assert res[1].shape == () def test_counts_expval(self): """Test counts works as expected if combined with expectation values""" - shots = 10 - dev = qml.device("default.qubit.legacy", wires=2, shots=shots) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(qml.device("default.qubit"), diff_method="parameter-shift", interface="torch") def circuit(): qml.Hadamard(wires=[0]) qml.CNOT(wires=[0, 1]) return qml.counts(qml.PauliZ(0)), qml.expval(qml.PauliX(1)) - res = circuit() + res = circuit(shots=10) assert len(res) == 2 assert isinstance(res, tuple) @@ -1628,22 +1531,19 @@ def circuit(): def test_sample_combination(self): """Test the output of combining expval, var and sample""" - n_sample = 10 - - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)), qml.expval(qml.PauliX(1)), qml.var(qml.PauliY(2)) - result = circuit() + result = circuit(shots=10) assert isinstance(result, tuple) assert len(result) == 3 - assert np.array_equal(result[0].shape, (n_sample,)) + assert np.array_equal(result[0].shape, (10,)) assert result[1].shape == () assert isinstance(result[1], torch.Tensor) assert result[2].shape == () @@ -1652,73 +1552,73 @@ def circuit(): def test_single_wire_sample(self): """Test the return type and shape of sampling a single wire""" - n_sample = 10 - dev = qml.device("default.qubit.legacy", wires=1, shots=n_sample) - - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): qml.RX(0.54, wires=0) return qml.sample(qml.PauliZ(0)) - result = circuit() + result = circuit(shots=10) assert isinstance(result, torch.Tensor) - assert np.array_equal(result.shape, (n_sample,)) + assert np.array_equal(result.shape, (10,)) def test_multi_wire_sample_regular_shape(self): """Test the return type and shape of sampling multiple wires where a rectangular array is expected""" - n_sample = 10 - - dev = qml.device("default.qubit.legacy", wires=3, shots=n_sample) - @qnode(dev, diff_method="parameter-shift", interface="torch") + @qnode(DefaultQubit(), diff_method="parameter-shift", interface="torch") def circuit(): return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - result = circuit() + result = circuit(shots=10) # If all the dimensions are equal the result will end up to be a proper rectangular array assert isinstance(result, tuple) - assert tuple(result[0].shape) == (n_sample,) - assert tuple(result[1].shape) == (n_sample,) - assert tuple(result[2].shape) == (n_sample,) + assert tuple(result[0].shape) == (10,) + assert tuple(result[1].shape) == (10,) + assert tuple(result[2].shape) == (10,) assert result[0].dtype == torch.float64 assert result[1].dtype == torch.float64 assert result[2].dtype == torch.float64 qubit_device_and_diff_method_and_grad_on_execution = [ - ["default.qubit.legacy", "backprop", True], - ["default.qubit.legacy", "finite-diff", False], - ["default.qubit.legacy", "parameter-shift", False], - ["default.qubit.legacy", "adjoint", True], - ["default.qubit.legacy", "adjoint", False], - ["default.qubit.legacy", "hadamard", False], + [DefaultQubit(), "backprop", True, False], + [DefaultQubit(), "finite-diff", False, False], + [DefaultQubit(), "parameter-shift", False, False], + [DefaultQubit(), "adjoint", True, False], + [DefaultQubit(), "adjoint", False, False], + [DefaultQubit(), "adjoint", True, True], + [DefaultQubit(), "adjoint", False, True], + [DefaultQubit(), "hadamard", False, False], ] @pytest.mark.parametrize( - "dev_name,diff_method,grad_on_execution", qubit_device_and_diff_method_and_grad_on_execution + "dev,diff_method,grad_on_execution, device_vjp", + qubit_device_and_diff_method_and_grad_on_execution, ) @pytest.mark.parametrize("shots", [None, 10000]) class TestReturn: """Class to test the shape of the Grad/Jacobian/Hessian with different return types.""" - def test_grad_single_measurement_param(self, dev_name, diff_method, grad_on_execution, shots): + # pylint:disable=too-many-public-methods + + def test_grad_single_measurement_param( + self, dev, diff_method, grad_on_execution, device_vjp, shots + ): """For one measurement and one param, the gradient is a float.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1726,7 +1626,7 @@ def circuit(a): a = torch.tensor(0.1, requires_grad=True) - res = circuit(a) + res = circuit(a, shots=shots) assert isinstance(res, torch.Tensor) assert res.shape == () @@ -1738,20 +1638,19 @@ def circuit(a): assert grad.shape == () def test_grad_single_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, shots, device_vjp ): """For one measurement and multiple param, the gradient is a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1760,7 +1659,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - res = circuit(a, b) + res = circuit(a, b, shots=shots) # gradient res.backward() @@ -1771,20 +1670,19 @@ def circuit(a, b): assert grad_b.shape == () def test_grad_single_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """For one measurement and multiple param as a single array params, the gradient is an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 1 - - if diff_method == "hadamard": - num_wires = 2 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1792,7 +1690,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - res = circuit(a) + res = circuit(a, shots=shots) # gradient res.backward() @@ -1802,7 +1700,7 @@ def circuit(a): assert grad.shape == (2,) def test_jacobian_single_measurement_param_probs( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """For a multi dimensional measurement (probs), check that a single array is returned with the correct dimension""" @@ -1812,14 +1710,13 @@ def test_jacobian_single_measurement_param_probs( if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -1833,7 +1730,7 @@ def circuit(a): assert jac.shape == (4,) def test_jacobian_single_measurement_probs_multiple_param( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1842,14 +1739,13 @@ def test_jacobian_single_measurement_probs_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -1858,7 +1754,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - jac = jacobian(circuit, (a, b)) + jac = jacobian(lambda _a, _b: circuit(_a, _b, shots=shots), (a, b)) assert isinstance(jac, tuple) @@ -1869,7 +1765,7 @@ def circuit(a, b): assert jac[1].shape == (4,) def test_jacobian_single_measurement_probs_multiple_param_single_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """For a multi dimensional measurement (probs), check that a single tuple is returned containing arrays with the correct dimension""" @@ -1878,39 +1774,31 @@ def test_jacobian_single_measurement_probs_multiple_param_single_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) return qml.probs(wires=[0, 1]) a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, torch.Tensor) assert jac.shape == (4, 2) def test_jacobian_expval_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - par_0 = torch.tensor(0.1, requires_grad=True) par_1 = torch.tensor(0.2, requires_grad=True) @@ -1919,6 +1807,7 @@ def test_jacobian_expval_expval_multiple_params( interface="torch", diff_method=diff_method, max_diff=1, + device_vjp=device_vjp, grad_on_execution=grad_on_execution, ) def circuit(x, y): @@ -1927,7 +1816,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.expval(qml.PauliZ(0)) - jac = jacobian(circuit, (par_0, par_1)) + jac = jacobian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) assert isinstance(jac, tuple) @@ -1946,20 +1835,19 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_expval_expval_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -1967,7 +1855,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -1979,7 +1867,7 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_var_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, device_vjp, grad_on_execution, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -1989,8 +1877,6 @@ def test_jacobian_var_var_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - par_0 = torch.tensor(0.1, requires_grad=True) par_1 = torch.tensor(0.2, requires_grad=True) @@ -1998,6 +1884,7 @@ def test_jacobian_var_var_multiple_params( dev, interface="torch", diff_method=diff_method, + device_vjp=device_vjp, max_diff=1, grad_on_execution=grad_on_execution, ) @@ -2007,7 +1894,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - jac = jacobian(circuit, (par_0, par_1)) + jac = jacobian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2027,7 +1914,7 @@ def circuit(x, y): assert jac[1][1].shape == () def test_jacobian_var_var_multiple_params_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, device_vjp, grad_on_execution, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2037,9 +1924,13 @@ def test_jacobian_var_var_multiple_params_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2047,7 +1938,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2059,22 +1950,22 @@ def circuit(a): assert jac[1].shape == (2,) def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The jacobian of multiple measurements with a single params return an array.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because of probabilities.") - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a): qml.RY(a, wires=0) qml.RX(0.2, wires=0) @@ -2082,7 +1973,7 @@ def circuit(a): a = torch.tensor(0.1, requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2094,7 +1985,7 @@ def circuit(a): assert jac[1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The jacobian of multiple measurements with a multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2102,14 +1993,13 @@ def test_jacobian_multiple_measurement_multiple_param( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + device_vjp=device_vjp, + grad_on_execution=grad_on_execution, + ) def circuit(a, b): qml.RY(a, wires=0) qml.RX(b, wires=0) @@ -2118,7 +2008,7 @@ def circuit(a, b): a = torch.tensor(0.1, requires_grad=True) b = torch.tensor(0.2, requires_grad=True) - jac = jacobian(circuit, (a, b)) + jac = jacobian(lambda _a, _b: circuit(_a, _b, shots=shots), (a, b)) assert isinstance(jac, tuple) assert len(jac) == 2 @@ -2138,7 +2028,7 @@ def circuit(a, b): assert jac[1][1].shape == (4,) def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The jacobian of multiple measurements with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2146,14 +2036,13 @@ def test_jacobian_multiple_measurement_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 3 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - @qnode(dev, interface="torch", diff_method=diff_method, grad_on_execution=grad_on_execution) + @qnode( + dev, + interface="torch", + diff_method=diff_method, + grad_on_execution=grad_on_execution, + device_vjp=device_vjp, + ) def circuit(a): qml.RY(a[0], wires=0) qml.RX(a[1], wires=0) @@ -2161,7 +2050,7 @@ def circuit(a): a = torch.tensor([0.1, 0.2], requires_grad=True) - jac = jacobian(circuit, a) + jac = jacobian(lambda _a: circuit(_a, shots=shots), a) assert isinstance(jac, tuple) assert len(jac) == 2 # measurements @@ -2172,23 +2061,18 @@ def circuit(a): assert isinstance(jac[1], torch.Tensor) assert jac[1].shape == (4, 2) - def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_execution, shots): + def test_hessian_expval_multiple_params( + self, dev, diff_method, grad_on_execution, shots, device_vjp + ): """The hessian of single a measurement with multiple params return a tuple of arrays.""" if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") - par_0 = torch.tensor(0.1, requires_grad=True) - par_1 = torch.tensor(0.2, requires_grad=True) + par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) + par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2196,6 +2080,7 @@ def test_hessian_expval_multiple_params(self, dev_name, diff_method, grad_on_exe diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2203,7 +2088,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(circuit, (par_0, par_1)) + hess = hessian(lambda a, b: circuit(a, b, shots=shots), (par_0, par_1)) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2223,7 +2108,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_expval_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2231,13 +2116,6 @@ def test_hessian_expval_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - params = torch.tensor([0.1, 0.2], requires_grad=True) @qnode( @@ -2246,6 +2124,7 @@ def test_hessian_expval_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2253,12 +2132,14 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(circuit, params) + hess = hessian(lambda _a: circuit(_a, shots=shots), params) assert isinstance(hess, torch.Tensor) assert hess.shape == (2, 2) - def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execution, shots): + def test_hessian_var_multiple_params( + self, dev, diff_method, grad_on_execution, device_vjp, shots + ): """The hessian of a single measurement with multiple params returns a tuple of arrays.""" if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") @@ -2267,15 +2148,8 @@ def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execut if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - par_0 = torch.tensor(0.1, requires_grad=True) - par_1 = torch.tensor(0.2, requires_grad=True) + par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) + par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2283,6 +2157,7 @@ def test_hessian_var_multiple_params(self, dev_name, diff_method, grad_on_execut diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2290,7 +2165,7 @@ def circuit(x, y): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(circuit, (par_0, par_1)) + hess = hessian(lambda _a, _b: circuit(_a, _b, shots=shots), (par_0, par_1)) assert isinstance(hess, tuple) assert len(hess) == 2 @@ -2310,7 +2185,7 @@ def circuit(x, y): assert hess[1][1].shape == () def test_hessian_var_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of single measurement with a multiple params array return a single array.""" if diff_method == "adjoint": @@ -2320,9 +2195,7 @@ def test_hessian_var_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - dev = qml.device(dev_name, wires=2, shots=shots) - - params = torch.tensor([0.1, 0.2], requires_grad=True) + params = torch.tensor([0.1, 0.2], requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2330,6 +2203,7 @@ def test_hessian_var_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2337,13 +2211,13 @@ def circuit(x): qml.CNOT(wires=[0, 1]) return qml.var(qml.PauliZ(0) @ qml.PauliX(1)) - hess = hessian(circuit, params) + hess = hessian(lambda _a: circuit(_a, shots=shots), params) assert isinstance(hess, torch.Tensor) assert hess.shape == (2, 2) def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" if diff_method == "adjoint": @@ -2354,15 +2228,8 @@ def test_hessian_probs_expval_multiple_params( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - num_wires = 2 - - if diff_method == "hadamard": - num_wires = 4 - - dev = qml.device(dev_name, wires=num_wires, shots=shots) - - par_0 = torch.tensor(0.1, requires_grad=True) - par_1 = torch.tensor(0.2, requires_grad=True) + par_0 = torch.tensor(0.1, requires_grad=True, dtype=torch.float64) + par_1 = torch.tensor(0.2, requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2370,6 +2237,7 @@ def test_hessian_probs_expval_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2378,7 +2246,7 @@ def circuit(x, y): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x, y): - return torch.hstack(circuit(x, y)) + return torch.hstack(circuit(x, y, shots=shots)) jac_fn = lambda x, y: jacobian(circuit_stack, (x, y), create_graph=True) @@ -2402,10 +2270,9 @@ def circuit_stack(x, y): assert tuple(hess[1][1].shape) == (3,) def test_hessian_expval_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with a multiple param array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2414,7 +2281,7 @@ def test_hessian_expval_probs_multiple_param_array( if shots is not None and diff_method in ("backprop", "adjoint"): pytest.skip("Test does not support finite shots and adjoint/backprop") - par = torch.tensor([0.1, 0.2], requires_grad=True) + par = torch.tensor([0.1, 0.2], requires_grad=True, dtype=torch.float64) @qnode( dev, @@ -2422,6 +2289,7 @@ def test_hessian_expval_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2430,7 +2298,7 @@ def circuit(x): return qml.expval(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x): - return torch.hstack(circuit(x)) + return torch.hstack(circuit(x, shots=shots)) jac_fn = lambda x: jacobian(circuit_stack, x, create_graph=True) @@ -2440,10 +2308,9 @@ def circuit_stack(x): assert tuple(hess.shape) == (3, 2, 2) def test_hessian_probs_var_multiple_params( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with multiple params return a tuple of arrays.""" - dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2461,6 +2328,7 @@ def test_hessian_probs_var_multiple_params( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x, y): qml.RX(x, wires=[0]) @@ -2469,7 +2337,7 @@ def circuit(x, y): return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x, y): - return torch.hstack(circuit(x, y)) + return torch.hstack(circuit(x, y, shots=shots)) jac_fn = lambda x, y: jacobian(circuit_stack, (x, y), create_graph=True) @@ -2493,10 +2361,9 @@ def circuit_stack(x, y): assert tuple(hess[1][1].shape) == (3,) def test_hessian_var_probs_multiple_param_array( - self, dev_name, diff_method, grad_on_execution, shots + self, dev, diff_method, grad_on_execution, device_vjp, shots ): """The hessian of multiple measurements with a multiple param array return a single array.""" - dev = qml.device(dev_name, wires=2, shots=shots) if diff_method == "adjoint": pytest.skip("Test does not supports adjoint because second order diff.") elif diff_method == "hadamard": @@ -2513,6 +2380,7 @@ def test_hessian_var_probs_multiple_param_array( diff_method=diff_method, max_diff=2, grad_on_execution=grad_on_execution, + device_vjp=device_vjp, ) def circuit(x): qml.RX(x[0], wires=[0]) @@ -2521,7 +2389,7 @@ def circuit(x): return qml.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.probs(wires=[0]) def circuit_stack(x): - return torch.hstack(circuit(x)) + return torch.hstack(circuit(x, shots=shots)) jac_fn = lambda x: jacobian(circuit_stack, x, create_graph=True) @@ -2531,14 +2399,11 @@ def circuit_stack(x): assert tuple(hess.shape) == (3, 2, 2) -@pytest.mark.parametrize("dev_name", ["default.qubit.legacy", "default.mixed"]) -def test_no_ops(dev_name): +def test_no_ops(): """Test that the return value of the QNode matches in the interface even if there are no ops""" - dev = qml.device(dev_name, wires=1) - - @qml.qnode(dev, interface="torch") + @qml.qnode(DefaultQubit(), interface="torch") def circuit(): qml.Hadamard(wires=0) return qml.state() diff --git a/tests/interfaces/test_transform_program_integration.py b/tests/interfaces/test_transform_program_integration.py index 8b6bb3a837e..ef492b52035 100644 --- a/tests/interfaces/test_transform_program_integration.py +++ b/tests/interfaces/test_transform_program_integration.py @@ -25,11 +25,12 @@ from pennylane.tape import QuantumTapeBatch from pennylane.typing import PostprocessingFn -device_suite = ( - qml.device("default.qubit.legacy", wires=5), - qml.devices.DefaultQubit(), - qml.device("lightning.qubit", wires=5), -) +with pytest.warns(qml.PennyLaneDeprecationWarning): + device_suite = ( + qml.device("default.qubit.legacy", wires=5), + qml.devices.DefaultQubit(), + qml.device("lightning.qubit", wires=5), + ) @pytest.mark.all_interfaces @@ -169,7 +170,7 @@ def split_sum_terms(tape): def test_chained_preprocessing(self): """Test a transform program with two transforms where their order affects the output.""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) def null_postprocessing(results): return results[0] diff --git a/tests/math/test_functions.py b/tests/math/test_functions.py index ec0eb076419..789bf5e2568 100644 --- a/tests/math/test_functions.py +++ b/tests/math/test_functions.py @@ -2189,7 +2189,7 @@ def circuit(weights): def test_jax(self, tol): """Test that the covariance matrix computes the correct result, and is differentiable, using the JAX interface""" - dev = qml.device("default.qubit.jax", wires=3) + dev = qml.device("default.qubit", wires=3) @qml.qnode(dev, interface="jax", diff_method="backprop") def circuit(weights): diff --git a/tests/measurements/test_state.py b/tests/measurements/test_state.py index 07c8b8d5fe6..6dce16479d2 100644 --- a/tests/measurements/test_state.py +++ b/tests/measurements/test_state.py @@ -17,7 +17,7 @@ import pennylane as qml from pennylane import numpy as pnp -from pennylane.devices import DefaultQubitLegacy +from pennylane.devices import DefaultMixed from pennylane.math.matrix_manipulation import _permute_dense_matrix from pennylane.math.quantum import reduce_dm, reduce_statevector from pennylane.measurements import ( @@ -361,7 +361,7 @@ def func(x): def test_no_state_capability(self, monkeypatch): """Test if an error is raised for devices that are not capable of returning the state. This is tested by changing the capability of default.qubit""" - dev = qml.device("default.qubit.legacy", wires=1) + dev = qml.device("default.mixed", wires=1) capabilities = dev.capabilities().copy() capabilities["returns_state"] = False @@ -370,7 +370,7 @@ def func(): return state() with monkeypatch.context() as m: - m.setattr(DefaultQubitLegacy, "capabilities", lambda *args, **kwargs: capabilities) + m.setattr(DefaultMixed, "capabilities", lambda *args, **kwargs: capabilities) with pytest.raises(qml.QuantumFunctionError, match="The current device is not capable"): func() @@ -410,9 +410,9 @@ def test_default_qubit_tf(self, diff_method): """Test that the returned state is equal to the expected returned state for all of PennyLane's built in statevector devices""" - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit", wires=4) - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, interface="tf", diff_method=diff_method) def func(): for i in range(4): qml.Hadamard(i) @@ -429,9 +429,9 @@ def test_default_qubit_autograd(self, diff_method): """Test that the returned state is equal to the expected returned state for all of PennyLane's built in statevector devices""" - dev = qml.device("default.qubit.autograd", wires=4) + dev = qml.device("default.qubit", wires=4) - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, interface="autograd", diff_method=diff_method) def func(): for i in range(4): qml.Hadamard(i) @@ -444,11 +444,11 @@ def func(): @pytest.mark.tf def test_gradient_with_passthru_tf(self): - """Test that the gradient of the state is accessible when using default.qubit.tf with the - backprop diff_method.""" + """Test that the gradient of the state is accessible when using default.qubit with the + tf interface and backprop diff_method.""" import tensorflow as tf - dev = qml.device("default.qubit.tf", wires=1) + dev = qml.device("default.qubit", wires=1) @qml.qnode(dev, interface="tf", diff_method="backprop") def func(x): @@ -466,10 +466,10 @@ def func(x): @pytest.mark.autograd def test_gradient_with_passthru_autograd(self): - """Test that the gradient of the state is accessible when using default.qubit.autograd - with the backprop diff_method.""" + """Test that the gradient of the state is accessible when using default.qubit + with autograd interface and the backprop diff_method.""" - dev = qml.device("default.qubit.autograd", wires=1) + dev = qml.device("default.qubit", wires=1) @qml.qnode(dev, interface="autograd", diff_method="backprop") def func(x): @@ -1014,7 +1014,7 @@ def func(): def test_no_state_capability(self, monkeypatch): """Test if an error is raised for devices that are not capable of returning the density matrix. This is tested by changing the capability of default.qubit""" - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.mixed", wires=2) capabilities = dev.capabilities().copy() capabilities["returns_state"] = False @@ -1023,7 +1023,7 @@ def func(): return density_matrix(0) with monkeypatch.context() as m: - m.setattr(DefaultQubitLegacy, "capabilities", lambda *args, **kwargs: capabilities) + m.setattr(DefaultMixed, "capabilities", lambda *args, **kwargs: capabilities) with pytest.raises( qml.QuantumFunctionError, match="The current device is not capable" " of returning the state", diff --git a/tests/ops/op_math/test_evolution.py b/tests/ops/op_math/test_evolution.py index 6143e10226c..8033c180f64 100644 --- a/tests/ops/op_math/test_evolution.py +++ b/tests/ops/op_math/test_evolution.py @@ -172,7 +172,7 @@ def circ_param_shift(x): Evolution(base, -0.5 * x) return qml.expval(qml.PauliZ(0)) - @qml.qnode(qml.device("default.qubit.jax", wires=1), interface="jax") + @qml.qnode(qml.device("default.qubit"), interface="jax") def circ(x): Evolution(qml.PauliX(0), -0.5 * x) return qml.expval(qml.PauliZ(0)) diff --git a/tests/ops/op_math/test_linear_combination.py b/tests/ops/op_math/test_linear_combination.py index c6e1584d19c..3c9b065bc41 100644 --- a/tests/ops/op_math/test_linear_combination.py +++ b/tests/ops/op_math/test_linear_combination.py @@ -53,13 +53,15 @@ def test_isinstance_Hamiltonian(self): assert isinstance(H, qml.Hamiltonian) -@pytest.mark.filterwarnings( - "ignore:Using 'qml.ops.Hamiltonian' with new operator arithmetic is deprecated" -) def test_mixed_legacy_warning_Hamiltonian(): """Test that mixing legacy ops and LinearCombination.compare raises a warning""" op1 = qml.ops.LinearCombination([0.5, 0.5], [X(0) @ X(1), qml.Hadamard(0)]) - op2 = qml.ops.Hamiltonian([0.5, 0.5], [qml.operation.Tensor(X(0), X(1)), qml.Hadamard(0)]) + + with pytest.warns( + qml.PennyLaneDeprecationWarning, + match="Using 'qml.ops.Hamiltonian' with new operator arithmetic is deprecated", + ): + op2 = qml.ops.Hamiltonian([0.5, 0.5], [qml.operation.Tensor(X(0), X(1)), qml.Hadamard(0)]) with pytest.warns(UserWarning, match="Attempting to compare a legacy operator class instance"): res = op1.compare(op2) diff --git a/tests/ops/qubit/test_observables.py b/tests/ops/qubit/test_observables.py index 50aa733fa68..448080e8e25 100644 --- a/tests/ops/qubit/test_observables.py +++ b/tests/ops/qubit/test_observables.py @@ -19,6 +19,7 @@ import numpy as np import pytest from gate_data import H, I, X, Y, Z +from scipy.sparse import csr_matrix import pennylane as qml from pennylane.ops.qubit.observables import BasisStateProjector, StateVectorProjector @@ -540,6 +541,75 @@ def test_serialization(self): qml.assert_equal(new_proj, proj) assert new_proj.id == proj.id # Ensure they are identical + def test_single_qubit_basis_state_0(self): + """Tests the function with a single-qubit basis state |0>.""" + basis_state = [0] + data = [1] + row_indices = [0] + col_indices = [0] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(2, 2)) + + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_single_qubit_basis_state_1(self): + """Tests the function with a single-qubit basis state |1>.""" + basis_state = [1] + data = [1] + row_indices = [1] + col_indices = [1] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(2, 2)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_two_qubit_basis_state_10(self): + """Tests the function with a two-qubits basis state |10>.""" + basis_state = [1, 0] + data = [1] + row_indices = [2] + col_indices = [2] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(4, 4)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_two_qubit_basis_state_01(self): + """Tests the function with a two-qubits basis state |01>.""" + basis_state = [0, 1] + data = [1] + row_indices = [1] + col_indices = [1] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(4, 4)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_two_qubit_basis_state_11(self): + """Tests the function with a two-qubits basis state |11>.""" + basis_state = [1, 1] + data = [1] + row_indices = [3] + col_indices = [3] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(4, 4)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_three_qubit_basis_state_101(self): + """Tests the function with a three-qubits basis state |101>.""" + basis_state = [1, 0, 1] + data = [1] + row_indices = [5] + col_indices = [5] + expected_matrix = csr_matrix((data, (row_indices, col_indices)), shape=(8, 8)) + actual_matrix = BasisStateProjector.compute_sparse_matrix(basis_state) + assert np.array_equal(expected_matrix.toarray(), actual_matrix.toarray()) + + def test_invalid_basis_state(self): + """Tests the function with an invalid state.""" + basis_state = [0, 2] # Invalid basis state + with pytest.raises(ValueError): + BasisStateProjector.compute_sparse_matrix(basis_state) + @pytest.mark.jax def test_jit_measurement(self): """Test that the measurement of a projector can be jitted.""" @@ -676,10 +746,10 @@ def test_matrix_representation(self, basis_state, expected, n_wires, tol): assert np.allclose(res_dynamic, expected, atol=tol) assert np.allclose(res_static, expected, atol=tol) - @pytest.mark.parametrize( - "dev", (qml.device("default.qubit"), qml.device("default.qubit.legacy", wires=1)) - ) - def test_integration_batched_state(self, dev): + @pytest.mark.parametrize("dev_name", ("default.qubit", "default.qubit.legacy")) + def test_integration_batched_state(self, dev_name): + dev = qml.device(dev_name, wires=1) + @qml.qnode(dev) def circuit(x): qml.RX(x, wires=0) diff --git a/tests/ops/qubit/test_parametric_ops.py b/tests/ops/qubit/test_parametric_ops.py index 6780cfe21a1..b20b5cc94bf 100644 --- a/tests/ops/qubit/test_parametric_ops.py +++ b/tests/ops/qubit/test_parametric_ops.py @@ -3464,9 +3464,9 @@ def test_simplify_rotations_grad_jax(self, op): import jax import jax.numpy as jnp - dev = qml.device("default.qubit.jax", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev) + @qml.qnode(dev, interface="jax") def circuit(simplify, wires, *params, **hyperparams): if simplify: qml.simplify(op(*params, wires=wires, **hyperparams)) diff --git a/tests/ops/qubit/test_qchem_ops.py b/tests/ops/qubit/test_qchem_ops.py index ce3be78b948..f652da4c5a0 100644 --- a/tests/ops/qubit/test_qchem_ops.py +++ b/tests/ops/qubit/test_qchem_ops.py @@ -273,10 +273,10 @@ def test_autograd(self, excitation): """Tests that operations are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = qml.device("default.qubit") state = np.array([0, -1 / np.sqrt(2), 1 / np.sqrt(2), 0]) - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) excitation(phi, wires=[0, 1]) @@ -298,9 +298,9 @@ def test_autograd_grad(self, diff_method, excitation, phi): """Tests that gradients are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="autograd") def circuit(phi): qml.PauliX(wires=0) excitation(phi, wires=[0, 1]) @@ -324,9 +324,9 @@ def test_tf(self, excitation, phi, diff_method): import tensorflow as tf - dev = qml.device("default.qubit.tf", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="tf") def circuit(phi): qml.PauliX(wires=0) excitation(phi, wires=[0, 1]) @@ -355,9 +355,9 @@ def test_jax(self, excitation, phi, diff_method): import jax - dev = qml.device("default.qubit.jax", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="jax") def circuit(phi): qml.PauliX(wires=0) excitation(phi, wires=[0, 1]) @@ -505,12 +505,12 @@ def test_autograd(self, excitation): """Tests that operations are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=4) + dev = qml.device("default.qubit") state = np.array( [0, 0, 0, -1 / np.sqrt(2), 0, 0, 0, 0, 0, 0, 0, 0, 1 / np.sqrt(2), 0, 0, 0] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -528,12 +528,12 @@ def test_tf(self, excitation): """Tests that operations are computed correctly using the tensorflow interface""" - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") state = np.array( [0, 0, 0, -1 / np.sqrt(2), 0, 0, 0, 0, 0, 0, 0, 0, 1 / np.sqrt(2), 0, 0, 0] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="tf") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -551,12 +551,12 @@ def test_jax(self, excitation): """Tests that operations are computed correctly using the jax interface""" - dev = qml.device("default.qubit.jax", wires=4) + dev = qml.device("default.qubit") state = np.array( [0, 0, 0, -1 / np.sqrt(2), 0, 0, 0, 0, 0, 0, 0, 0, 1 / np.sqrt(2), 0, 0, 0] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="jax") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -579,9 +579,9 @@ def test_autograd_grad(self, excitation, phi): """Tests that gradients are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -607,9 +607,9 @@ def test_tf_grad(self, excitation, phi, diff_method): import tensorflow as tf - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="tf") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -639,9 +639,9 @@ def test_jax_grad(self, excitation, phi, diff_method): import jax - dev = qml.device("default.qubit.jax", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="jax") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -774,7 +774,7 @@ def test_autograd(self): """Tests that operations are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=4) + dev = qml.device("default.qubit") state = np.array( [ 0.0 + 0.0j, @@ -796,7 +796,7 @@ def test_autograd(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -811,7 +811,7 @@ def test_tf(self): """Tests that operations are computed correctly using the tensorflow interface""" - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") state = np.array( [ 0.0 + 0.0j, @@ -833,7 +833,7 @@ def test_tf(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="tf") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -848,7 +848,7 @@ def test_jax(self): """Tests that operations are computed correctly using the jax interface""" - dev = qml.device("default.qubit.jax", wires=4) + dev = qml.device("default.qubit") state = np.array( [ 0.0 + 0.0j, @@ -870,7 +870,7 @@ def test_jax(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="jax") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -885,7 +885,7 @@ def test_torch(self): """Tests that operations are computed correctly using the torch interface""" - dev = qml.device("default.qubit.torch", wires=4) + dev = qml.device("default.qubit") state = np.array( [ 0.0 + 0.0j, @@ -907,7 +907,7 @@ def test_torch(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="torch") def circuit(phi): qml.PauliX(wires=0) qml.PauliX(wires=1) @@ -930,10 +930,14 @@ def test_autograd_grad(self, phi, diff_method): """Tests that gradients are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=4) + dev = qml.device("default.qubit") - circuit_0 = qml.QNode(self.grad_circuit_0, dev, diff_method=diff_method) - circuit_1 = qml.QNode(self.grad_circuit_1, dev, diff_method=diff_method) + circuit_0 = qml.QNode( + self.grad_circuit_0, dev, interface="autograd", diff_method=diff_method + ) + circuit_1 = qml.QNode( + self.grad_circuit_1, dev, interface="autograd", diff_method=diff_method + ) total = lambda phi: 1.1 * circuit_0(phi) + 0.7 * circuit_1(phi) assert np.allclose(qml.grad(total)(phi), self.expected_grad_fn(phi)) @@ -950,10 +954,10 @@ def test_tf_grad(self, phi, diff_method): import tensorflow as tf - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") - circuit_0 = qml.QNode(self.grad_circuit_0, dev, diff_method=diff_method) - circuit_1 = qml.QNode(self.grad_circuit_1, dev, diff_method=diff_method) + circuit_0 = qml.QNode(self.grad_circuit_0, dev, interface="tf", diff_method=diff_method) + circuit_1 = qml.QNode(self.grad_circuit_1, dev, interface="tf", diff_method=diff_method) total = lambda phi: 1.1 * circuit_0(phi) + 0.7 * circuit_1(phi) phi_t = tf.Variable(phi, dtype=tf.float64) @@ -976,10 +980,10 @@ def test_jax_grad(self, phi, diff_method): import jax - dev = qml.device("default.qubit.jax", wires=4) + dev = qml.device("default.qubit") - circuit_0 = qml.QNode(self.grad_circuit_0, dev, diff_method=diff_method) - circuit_1 = qml.QNode(self.grad_circuit_1, dev, diff_method=diff_method) + circuit_0 = qml.QNode(self.grad_circuit_0, dev, interface="jax", diff_method=diff_method) + circuit_1 = qml.QNode(self.grad_circuit_1, dev, interface="jax", diff_method=diff_method) total = lambda phi: 1.1 * circuit_0(phi) + 0.7 * circuit_1(phi) phi_j = jax.numpy.array(phi) @@ -998,10 +1002,10 @@ def test_torch_grad(self, phi, diff_method): import torch - dev = qml.device("default.qubit.torch", wires=4) + dev = qml.device("default.qubit") - circuit_0 = qml.QNode(self.grad_circuit_0, dev, diff_method=diff_method) - circuit_1 = qml.QNode(self.grad_circuit_1, dev, diff_method=diff_method) + circuit_0 = qml.QNode(self.grad_circuit_0, dev, interface="torch", diff_method=diff_method) + circuit_1 = qml.QNode(self.grad_circuit_1, dev, interface="torch", diff_method=diff_method) total = lambda phi: 1.1 * circuit_0(phi) + 0.7 * circuit_1(phi) phi_t = torch.tensor(phi, dtype=torch.complex128, requires_grad=True) @@ -1105,7 +1109,7 @@ def test_autograd(self): """Tests that operations are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = qml.device("default.qubit") state = np.array( [ 0, @@ -1115,7 +1119,7 @@ def test_autograd(self): ] ) - @qml.qnode(dev) + @qml.qnode(dev, interface="autograd") def circuit(phi): qml.PauliX(wires=0) qml.FermionicSWAP(phi, wires=[0, 1]) @@ -1137,9 +1141,9 @@ def test_autograd_grad(self, diff_method, phi): """Tests that gradients are computed correctly using the autograd interface""" - dev = qml.device("default.qubit.autograd", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, interface="autograd", diff_method=diff_method) def circuit(phi): qml.PauliX(wires=0) qml.FermionicSWAP(phi, wires=[0, 1]) @@ -1163,9 +1167,9 @@ def test_tf(self, phi, diff_method): import tensorflow as tf - dev = qml.device("default.qubit.tf", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, diff_method=diff_method, interface="tf") def circuit(phi): qml.PauliX(wires=0) qml.FermionicSWAP(phi, wires=[0, 1]) @@ -1194,9 +1198,9 @@ def test_jax(self, phi, diff_method): import jax - dev = qml.device("default.qubit.jax", wires=2) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method=diff_method) + @qml.qnode(dev, interface="jax", diff_method=diff_method) def circuit(phi): qml.PauliX(wires=0) qml.FermionicSWAP(phi, wires=[0, 1]) diff --git a/tests/pulse/test_transmon.py b/tests/pulse/test_transmon.py index eb280f04bab..6a0631d3abc 100644 --- a/tests/pulse/test_transmon.py +++ b/tests/pulse/test_transmon.py @@ -510,7 +510,7 @@ def fb(p, t): Hd = transmon_drive(amplitude=fa, phase=fb, freq=0.5, wires=[0]) H = Hi + Hd - dev = qml.device("default.qubit.jax", wires=wires) + dev = qml.device("default.qubit") ts = jnp.array([0.0, 3.0]) H_obj = sum(qml.PauliZ(i) for i in range(2)) diff --git a/tests/qnn/test_keras.py b/tests/qnn/test_keras.py index af8d457cd9e..f4f9769edc2 100644 --- a/tests/qnn/test_keras.py +++ b/tests/qnn/test_keras.py @@ -503,7 +503,7 @@ def test_gradients(self, get_circuit, output_dim, n_qubits): # pylint: disable= def test_backprop_gradients(self, mocker): # pylint: disable=no-self-use """Test if KerasLayer is compatible with the backprop diff method.""" - dev = qml.device("default.qubit.tf", wires=2) + dev = qml.device("default.qubit") @qml.qnode(dev, interface="tf", diff_method="backprop") def f(inputs, weights): @@ -813,9 +813,9 @@ def test_no_attribute(): @pytest.mark.tf def test_batch_input_single_measure(tol): """Test input batching in keras""" - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method="parameter-shift") + @qml.qnode(dev, interface="tf", diff_method="parameter-shift") def circuit(x, weights): qml.AngleEmbedding(x, wires=range(4), rotation="Y") qml.RY(weights[0], wires=0) @@ -830,7 +830,6 @@ def circuit(x, weights): res = layer(x) assert res.shape == (10, 2) - assert dev.num_executions == 1 for x_, r in zip(x, res): assert qml.math.allclose(r, circuit(x_, layer.qnode_weights["weights"]), atol=tol) @@ -842,9 +841,9 @@ def circuit(x, weights): @pytest.mark.tf def test_batch_input_multi_measure(tol): """Test input batching in keras for multiple measurements""" - dev = qml.device("default.qubit.tf", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method="parameter-shift") + @qml.qnode(dev, interface="tf", diff_method="parameter-shift") def circuit(x, weights): qml.AngleEmbedding(x, wires=range(4), rotation="Y") qml.RY(weights[0], wires=0) @@ -859,7 +858,6 @@ def circuit(x, weights): res = layer(x) assert res.shape == (10, 5) - assert dev.num_executions == 1 for x_, r in zip(x, res): exp = tf.experimental.numpy.hstack(circuit(x_, layer.qnode_weights["weights"])) diff --git a/tests/qnn/test_qnn_torch.py b/tests/qnn/test_qnn_torch.py index a25ee7d6949..64aeb9b1a9c 100644 --- a/tests/qnn/test_qnn_torch.py +++ b/tests/qnn/test_qnn_torch.py @@ -805,9 +805,9 @@ def circ(inputs, w0): # pylint: disable=unused-argument @pytest.mark.torch def test_batch_input_single_measure(tol): """Test input batching in torch""" - dev = qml.device("default.qubit.torch", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method="parameter-shift") + @qml.qnode(dev, interface="torch", diff_method="parameter-shift") def circuit(x, weights): qml.AngleEmbedding(x, wires=range(4), rotation="Y") qml.RY(weights[0], wires=0) @@ -820,7 +820,6 @@ def circuit(x, weights): res = layer(x) assert res.shape == (10, 2) - assert dev.num_executions == 1 for x_, r in zip(x, res): assert qml.math.allclose(r, circuit(x_, layer.qnode_weights["weights"]), atol=tol) @@ -832,9 +831,9 @@ def circuit(x, weights): @pytest.mark.torch def test_batch_input_multi_measure(tol): """Test input batching in torch for multiple measurements""" - dev = qml.device("default.qubit.torch", wires=4) + dev = qml.device("default.qubit") - @qml.qnode(dev, diff_method="parameter-shift") + @qml.qnode(dev, interface="torch", diff_method="parameter-shift") def circuit(x, weights): qml.AngleEmbedding(x, wires=range(4), rotation="Y") qml.RY(weights[0], wires=0) @@ -847,7 +846,6 @@ def circuit(x, weights): res = layer(x) assert res.shape == (10, 5) - assert dev.num_executions == 1 for x_, r in zip(x, res): exp = torch.hstack(circuit(x_, layer.qnode_weights["weights"])) diff --git a/tests/resource/test_specs.py b/tests/resource/test_specs.py index ef9592e5cd9..ac2744d50e1 100644 --- a/tests/resource/test_specs.py +++ b/tests/resource/test_specs.py @@ -22,6 +22,13 @@ from pennylane.tape import QuantumTapeBatch from pennylane.typing import PostprocessingFn +with pytest.warns(qml.PennyLaneDeprecationWarning): + devices_list = [ + (qml.device("default.qubit"), 1), + (qml.device("default.qubit", wires=2), 2), + (qml.device("default.qubit.legacy", wires=2), 2), + ] + class TestSpecsTransform: """Tests for the transform specs using the QNode""" @@ -329,11 +336,7 @@ def circuit(): @pytest.mark.parametrize( "device,num_wires", - [ - (qml.device("default.qubit"), 1), - (qml.device("default.qubit", wires=2), 2), - (qml.device("default.qubit.legacy", wires=2), 2), - ], + devices_list, ) def test_num_wires_source_of_truth(self, device, num_wires): """Tests that num_wires behaves differently on old and new devices.""" diff --git a/tests/test_compiler.py b/tests/test_compiler.py index 4b66bf2d061..92130d5f3fd 100644 --- a/tests/test_compiler.py +++ b/tests/test_compiler.py @@ -353,6 +353,45 @@ def inner(j): assert circuit(5, 6) == 30 # 5 * 6 assert circuit(4, 7) == 28 # 4 * 7 + def test_while_loop_python_fallback(self): + """Test that qml.while_loop fallsback to + Python without qjit""" + + def f(n, m): + @qml.while_loop(lambda i, _: i < n) + def outer(i, sm): + @qml.while_loop(lambda j: j < m) + def inner(j): + return j + 1 + + return i + 1, sm + inner(0) + + return outer(0, 0)[1] + + assert f(5, 6) == 30 # 5 * 6 + assert f(4, 7) == 28 # 4 * 7 + + def test_fallback_while_loop_qnode(self): + """Test that qml.while_loop inside a qnode fallsback to + Python without qjit""" + dev = qml.device("lightning.qubit", wires=1) + + @qml.qnode(dev) + def circuit(n): + @qml.while_loop(lambda v: v[0] < v[1]) + def loop(v): + qml.PauliX(wires=0) + return v[0] + 1, v[1] + + loop((0, n)) + return qml.expval(qml.PauliZ(0)) + + assert jnp.allclose(circuit(1), -1.0) + + res = circuit.tape.operations + expected = [qml.PauliX(0) for i in range(4)] + _ = [qml.assert_equal(i, j) for i, j in zip(res, expected)] + def test_dynamic_wires_for_loops(self): """Test for loops with iteration index-dependant wires.""" dev = qml.device("lightning.qubit", wires=6) @@ -405,6 +444,57 @@ def inner(j): assert jnp.allclose(circuit(4), jnp.eye(2**4)[0]) + def test_for_loop_python_fallback(self): + """Test that qml.for_loop fallsback to Python + interpretation if Catalyst is not available""" + dev = qml.device("lightning.qubit", wires=3) + + @qml.qnode(dev) + def circuit(x, n): + + # for loop with dynamic bounds + @qml.for_loop(0, n, 1) + def loop_fn(i): + qml.Hadamard(wires=i) + + # nested for loops. + # outer for loop updates x + @qml.for_loop(0, n, 1) + def loop_fn_returns(i, x): + qml.RX(x, wires=i) + + # inner for loop + @qml.for_loop(i + 1, n, 1) + def inner(j): + qml.CRY(x**2, [i, j]) + + inner() + + return x + 0.1 + + loop_fn() + loop_fn_returns(x) + + return qml.expval(qml.PauliZ(0)) + + x = 0.5 + assert jnp.allclose(circuit(x, 3), qml.qjit(circuit)(x, 3)) + + res = circuit.tape.operations + expected = [ + qml.Hadamard(wires=[0]), + qml.Hadamard(wires=[1]), + qml.Hadamard(wires=[2]), + qml.RX(0.5, wires=[0]), + qml.CRY(0.25, wires=[0, 1]), + qml.CRY(0.25, wires=[0, 2]), + qml.RX(0.6, wires=[1]), + qml.CRY(0.36, wires=[1, 2]), + qml.RX(0.7, wires=[2]), + ] + + _ = [qml.assert_equal(i, j) for i, j in zip(res, expected)] + def test_cond(self): """Test condition with simple true_fn""" dev = qml.device("lightning.qubit", wires=1) diff --git a/tests/test_debugging.py b/tests/test_debugging.py index 67544dc5634..bdd10a4cfd5 100644 --- a/tests/test_debugging.py +++ b/tests/test_debugging.py @@ -181,7 +181,12 @@ def circuit(): return qml.expval(qml.PauliZ(0)) - qml.snapshots(circuit)(shots=200) + with ( + pytest.warns(UserWarning, match="Requested state or density matrix with finite shots") + if isinstance(dev, qml.devices.default_qutrit.DefaultQutrit) + else nullcontext() + ): + qml.snapshots(circuit)(shots=200) @pytest.mark.parametrize("diff_method", [None, "parameter-shift"]) def test_all_state_measurement_snapshot_pure_qubit_dev(self, dev, diff_method): @@ -275,7 +280,8 @@ def circuit(): @pytest.mark.parametrize("diff_method", [None, "backprop", "parameter-shift", "adjoint"]) def test_default_qubit_legacy_only_supports_state(self, diff_method): - dev = qml.device("default.qubit.legacy", wires=2) + with pytest.warns(qml.PennyLaneDeprecationWarning, match="Use of 'default.qubit"): + dev = qml.device("default.qubit.legacy", wires=2) assert qml.debugging.snapshot._is_snapshot_compatible(dev) diff --git a/tests/test_operation.py b/tests/test_operation.py index 5fae0e13aaf..3062803df0c 100644 --- a/tests/test_operation.py +++ b/tests/test_operation.py @@ -1088,7 +1088,7 @@ def test_all_wires_defined_but_init_with_one(self): """Test that an exception is raised if the class is defined with ALL wires, but then instantiated with only one""" - dev1 = qml.device("default.qubit.legacy", wires=2) + dev1 = qml.device("default.qubit", wires=2) class DummyOp(qml.operation.Operation): r"""Dummy custom operator""" diff --git a/tests/test_qnode.py b/tests/test_qnode.py index 207ce42e1ec..52d9155e448 100644 --- a/tests/test_qnode.py +++ b/tests/test_qnode.py @@ -887,10 +887,7 @@ def circuit(x, y): assert np.allclose(res, expected, atol=tol, rtol=0) - @pytest.mark.parametrize( - "dev", - [qml.device("default.qubit", wires=3), qml.device("default.qubit.legacy", wires=3)], - ) + @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.legacy"]) @pytest.mark.parametrize("first_par", np.linspace(0.15, np.pi - 0.3, 3)) @pytest.mark.parametrize("sec_par", np.linspace(0.15, np.pi - 0.3, 3)) @pytest.mark.parametrize( @@ -905,11 +902,12 @@ def circuit(x, y): ], ) def test_defer_meas_if_mcm_unsupported( - self, dev, first_par, sec_par, return_type, mv_return, mv_res, mocker + self, dev_name, first_par, sec_par, return_type, mv_return, mv_res, mocker ): # pylint: disable=too-many-arguments """Tests that the transform using the deferred measurement principle is applied if the device doesn't support mid-circuit measurements natively.""" + dev = qml.device(dev_name, wires=3) @qml.qnode(dev) def cry_qnode(x, y): diff --git a/tests/test_qnode_legacy.py b/tests/test_qnode_legacy.py index b21eaa305aa..d8edc0890ee 100644 --- a/tests/test_qnode_legacy.py +++ b/tests/test_qnode_legacy.py @@ -491,13 +491,14 @@ def test_adjoint_finite_shots(self): dev = qml.device("default.qubit.legacy", wires=1, shots=1) - @qnode(dev, diff_method="adjoint") - def circ(): - return qml.expval(qml.PauliZ(0)) - with pytest.warns( UserWarning, match="Requested adjoint differentiation to be computed with finite shots." ): + + @qnode(dev, diff_method="adjoint") + def circ(): + return qml.expval(qml.PauliZ(0)) + circ() @pytest.mark.autograd @@ -621,7 +622,8 @@ def circuit(params): } def test_autograd_interface_device_switched_no_warnings(self): - """Test that checks that no warning is raised for device switch when you define an interface.""" + """Test that checks that no warning is raised for device switch when you define an interface, + except for the deprecation warnings which will be caught by the fixture.""" dev = qml.device("default.qubit.legacy", wires=1) @qml.qnode(dev, interface="autograd") @@ -629,17 +631,16 @@ def circuit(params): qml.RX(params, wires=0) return qml.expval(qml.PauliZ(0)) - with warnings.catch_warnings(record=True) as record: - circuit(qml.numpy.array(0.1, requires_grad=True)) - - assert len(record) == 0 + circuit(qml.numpy.array(0.1, requires_grad=True)) def test_not_giving_mode_kwarg_does_not_raise_warning(self): - """Test that not providing a value for mode does not raise a warning.""" + """Test that not providing a value for mode does not raise a warning + except for the deprecation warning.""" with warnings.catch_warnings(record=True) as record: - _ = qml.QNode(lambda f: f, qml.device("default.qubit.legacy", wires=1)) + qml.QNode(lambda f: f, qml.device("default.qubit.legacy", wires=1)) - assert len(record) == 0 + assert len(record) == 1 + assert record[0].category == qml.PennyLaneDeprecationWarning class TestTapeConstruction: @@ -1089,9 +1090,7 @@ def circuit(): assert len(circuit.tape.operations) == 2 assert isinstance(circuit.tape.operations[1], qml.measurements.MidMeasureMP) - @pytest.mark.parametrize( - "dev", [qml.device("default.qubit", wires=3), qml.device("default.qubit.legacy", wires=3)] - ) + @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.legacy"]) @pytest.mark.parametrize("first_par", np.linspace(0.15, np.pi - 0.3, 3)) @pytest.mark.parametrize("sec_par", np.linspace(0.15, np.pi - 0.3, 3)) @pytest.mark.parametrize( @@ -1106,12 +1105,14 @@ def circuit(): ], ) def test_defer_meas_if_mcm_unsupported( - self, dev, first_par, sec_par, return_type, mv_return, mv_res, mocker + self, dev_name, first_par, sec_par, return_type, mv_return, mv_res, mocker ): # pylint: disable=too-many-arguments """Tests that the transform using the deferred measurement principle is applied if the device doesn't support mid-circuit measurements natively.""" + dev = qml.device(dev_name, wires=3) + @qml.qnode(dev) def cry_qnode(x, y): """QNode where we apply a controlled Y-rotation.""" diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py index 2b2387ff6ed..787ed809fbf 100644 --- a/tests/test_qubit_device.py +++ b/tests/test_qubit_device.py @@ -1211,7 +1211,7 @@ def test_device_executions(self): """Test the number of times a qubit device is executed over a QNode's lifetime is tracked by `num_executions`""" - dev_1 = qml.device("default.qubit.legacy", wires=2) + dev_1 = qml.device("default.mixed", wires=2) def circuit_1(x, y): qml.RX(x, wires=[0]) @@ -1227,7 +1227,7 @@ def circuit_1(x, y): assert dev_1.num_executions == num_evals_1 # test a second instance of a default qubit device - dev_2 = qml.device("default.qubit.legacy", wires=2) + dev_2 = qml.device("default.mixed", wires=2) def circuit_2(x): qml.RX(x, wires=[0]) @@ -1272,7 +1272,7 @@ def test_device_executions(self): """Test the number of times a qubit device is executed over a QNode's lifetime is tracked by `num_executions`""" - dev_1 = qml.device("default.qubit.legacy", wires=2) + dev_1 = qml.device("default.mixed", wires=2) def circuit_1(x, y): qml.RX(x, wires=[0]) @@ -1285,10 +1285,10 @@ def circuit_1(x, y): for _ in range(num_evals_1): node_1(0.432, np.array([0.12, 0.5, 3.2])) - assert dev_1.num_executions == num_evals_1 + assert dev_1.num_executions == num_evals_1 * 3 # test a second instance of a default qubit device - dev_2 = qml.device("default.qubit.legacy", wires=2) + dev_2 = qml.device("default.mixed", wires=2) assert dev_2.num_executions == 0 @@ -1302,7 +1302,7 @@ def circuit_2(x, y): for _ in range(num_evals_2): node_2(np.array([0.432, 0.61, 8.2]), 0.12) - assert dev_2.num_executions == num_evals_2 + assert dev_2.num_executions == num_evals_2 * 3 # test a new circuit on an existing instance of a qubit device def circuit_3(x, y): @@ -1315,7 +1315,7 @@ def circuit_3(x, y): for _ in range(num_evals_3): node_3(np.array([0.432, 0.2]), np.array([0.12, 1.214])) - assert dev_1.num_executions == num_evals_1 + num_evals_3 + assert dev_1.num_executions == num_evals_1 * 3 + num_evals_3 * 2 class TestBatchExecution: @@ -1504,7 +1504,7 @@ def test_tracker_multi_execution(self, dev_name): @pytest.mark.autograd def test_tracker_grad(self): """Test that the tracker can track resources through a gradient computation""" - dev = qml.device("default.qubit.legacy", wires=1, shots=100) + dev = qml.device("default.qubit", wires=1, shots=100) @qml.qnode(dev, diff_method="parameter-shift") def circuit(x): @@ -1540,8 +1540,9 @@ def test_samples_to_counts_with_nan(self): """Test that the counts function disregards failed measurements (samples including NaN values) when totalling counts""" # generate 1000 samples for 2 wires, randomly distributed between 0 and 1 - device = qml.device("default.qubit.legacy", wires=2, shots=1000) - device._state = [0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j] + device = qml.device("default.mixed", wires=2, shots=1000) + sv = [0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j, 0.5 + 0.0j] + device._state = np.outer(sv, sv) device._samples = device.generate_samples() samples = device.sample(qml.measurements.CountsMP()) @@ -1568,9 +1569,12 @@ def test_samples_to_counts_with_many_wires(self, all_outcomes): # generate 1000 samples for 10 wires, randomly distributed between 0 and 1 n_wires = 10 shots = 100 - device = qml.device("default.qubit.legacy", wires=n_wires, shots=shots) - state = np.random.rand(*([2] * n_wires)) - device._state = state / np.linalg.norm(state) + device = qml.device("default.mixed", wires=n_wires, shots=shots) + + sv = np.random.rand(*([2] * n_wires)) + state = sv / np.linalg.norm(sv) + + device._state = np.outer(state, state) device._samples = device.generate_samples() samples = device.sample(qml.measurements.CountsMP(all_outcomes=all_outcomes)) diff --git a/tests/transforms/test_hamiltonian_expand.py b/tests/transforms/test_hamiltonian_expand.py index 894bf4b65a9..b3f2ebad99d 100644 --- a/tests/transforms/test_hamiltonian_expand.py +++ b/tests/transforms/test_hamiltonian_expand.py @@ -15,6 +15,7 @@ Unit tests for the ``hamiltonian_expand`` transform. """ import functools +import warnings import numpy as np import pytest @@ -91,11 +92,15 @@ class TestHamiltonianExpand: """Tests for the hamiltonian_expand transform""" @pytest.fixture(scope="function", autouse=True) - def capture_warnings(self, recwarn): - yield - if len(recwarn) > 0: - for w in recwarn: - assert isinstance(w.message, qml.PennyLaneDeprecationWarning) + def capture_warnings(self): + with pytest.warns(qml.PennyLaneDeprecationWarning) as record: + yield + + for w in record: + assert isinstance(w.message, qml.PennyLaneDeprecationWarning) + if "qml.transforms.hamiltonian_expand is deprecated" not in str(w.message): + warnings.warn(w.message, w.category) + else: assert "qml.transforms.hamiltonian_expand is deprecated" in str(w.message) def test_ham_with_no_terms_raises(self): @@ -527,11 +532,15 @@ class TestSumExpand: """Tests for the sum_expand transform""" @pytest.fixture(scope="function", autouse=True) - def capture_warnings(self, recwarn): - yield - if len(recwarn) > 0: - for w in recwarn: - assert isinstance(w.message, qml.PennyLaneDeprecationWarning) + def capture_warnings(self): + with pytest.warns(qml.PennyLaneDeprecationWarning) as record: + yield + + for w in record: + assert isinstance(w.message, qml.PennyLaneDeprecationWarning) + if "qml.transforms.sum_expand is deprecated" not in str(w.message): + warnings.warn(w.message, w.category) + else: assert "qml.transforms.sum_expand is deprecated" in str(w.message) def test_observables_on_same_wires(self): @@ -560,6 +569,7 @@ def test_sums(self, qscript, output): assert all(qml.math.allclose(o, e) for o, e in zip(output, expval)) @pytest.mark.parametrize(("qscript", "output"), zip(SUM_QSCRIPTS, SUM_OUTPUTS)) + @pytest.mark.filterwarnings("ignore:Use of 'default.qubit.legacy' is deprecated") def test_sums_legacy_opmath(self, qscript, output): """Tests that the sum_expand transform returns the correct value""" dev_old = qml.device("default.qubit.legacy", wires=4) diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py index fe8d8bc3b46..1a78c1ab4e2 100644 --- a/tests/transforms/test_qcut.py +++ b/tests/transforms/test_qcut.py @@ -4070,9 +4070,7 @@ def test_simple_cut_circuit_torch_trace(self, mocker, use_opt_einsum): import torch - # TODO: this passes with default.qubit locally, but fails on CI - # possibly an architecture-specific issue - dev = qml.device("default.qubit.legacy", wires=2) + dev = qml.device("default.qubit", wires=2) @qml.qnode(dev, interface="torch") def circuit(x): diff --git a/tests/transforms/test_tape_expand.py b/tests/transforms/test_tape_expand.py index 5d7b59a671a..4bf16a061d4 100644 --- a/tests/transforms/test_tape_expand.py +++ b/tests/transforms/test_tape_expand.py @@ -926,7 +926,7 @@ def test_custom_decomp_used_twice(self): custom_decomps = {"MultiRZ": qml.MultiRZ.compute_decomposition} dev = qml.device("lightning.qubit", wires=2, custom_decomps=custom_decomps) - @qml.qnode(dev, diff_method="adjoint", expansion_strategy="device") + @qml.qnode(dev, diff_method="adjoint", expansion_strategy="gradient") def cost(theta): qml.Hadamard(wires=0) qml.Hadamard(wires=1)