diff --git a/.github/workflows/interface-unit-tests.yml b/.github/workflows/interface-unit-tests.yml index 601e762fc92..70104146499 100644 --- a/.github/workflows/interface-unit-tests.yml +++ b/.github/workflows/interface-unit-tests.yml @@ -379,8 +379,7 @@ jobs: # catalyst requires the latest version of pennylane that is about to be released. # Installing catalyst after pennylane to make sure that the latest catalyst is used. install_catalyst_nightly: true - # using lightning master does not work for the tests with external libraries - install_pennylane_lightning_master: false + install_pennylane_lightning_master: true pytest_coverage_flags: ${{ inputs.pytest_coverage_flags }} pytest_markers: external additional_pip_packages: pyzx matplotlib stim quimb mitiq pennylane-qiskit ply diff --git a/Makefile b/Makefile index 7f36c24c698..8b3bec9be42 100644 --- a/Makefile +++ b/Makefile @@ -60,12 +60,10 @@ clean-docs: test: $(PYTHON) $(TESTRUNNER) - $(PYTHON) $(PLUGIN_TESTRUNNER) --device=default.qubit.autograd coverage: @echo "Generating coverage report..." $(PYTHON) $(TESTRUNNER) $(COVERAGE) - $(PYTHON) $(PLUGIN_TESTRUNNER) --device=default.qubit.autograd $(COVERAGE) --cov-append .PHONY:format format: diff --git a/doc/development/deprecations.rst b/doc/development/deprecations.rst index 3954c633d00..c91979963a9 100644 --- a/doc/development/deprecations.rst +++ b/doc/development/deprecations.rst @@ -9,17 +9,18 @@ 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. +* ``Device``, ``QubitDevice``, and ``QutritDevice`` will no longer be imported top level in v0.40. They instead + we be available as ``qml.devices.LegacyDevice``, ``qml.devices.QubitDevice``, and ``qml.devices.QutritDevice`` + respectively. - - Deprecated in v0.38 - - Will be removed in v0.39 + - Deprecated top level access in v0.39 + - Top level access removed in v0.40 -* The logic for internally switching a device for a different backpropagation - compatible device is now deprecated, as it was in place for the deprecated ``default.qubit.legacy``. +* `QNode.gradient_fn` is deprecated. Please use `QNode.diff_method` instead. `QNode.get_gradient_fn` can also be used to + process the diff method. - - Deprecated in v0.38 - - Will be removed in v0.39 + - Deprecated in v0.39 + - Will be removed in v0.40 * The ``decomp_depth`` argument in ``qml.device`` is deprecated. @@ -82,6 +83,23 @@ Other deprecations Completed deprecation cycles ---------------------------- +* All of the legacy devices (any with the name ``default.qubit.{autograd,torch,tf,jax,legacy}``) are removed. Use ``default.qubit`` instead, + as it supports backpropagation for the many backends the legacy devices support. + + - Deprecated in v0.38 + - Removed in v0.39 + +* The logic for internally switching a device for a different backpropagation + compatible device is removed, as it was in place for removed ``default.qubit.legacy``. + + - Deprecated in v0.38 + - Removed in v0.39 + +* `Operator.expand` is now removed. Use `qml.tape.QuantumScript(op.deocomposition())` instead. + + - Deprecated in v0.38 + - Removed in v0.39 + * The ``expansion_strategy`` attribute of ``qml.QNode`` is removed. Users should make use of ``qml.workflow.construct_batch``, should they require fine control over the output tape(s). diff --git a/doc/development/guide/logging.rst b/doc/development/guide/logging.rst index 54f0a0cc349..9af2e34520c 100644 --- a/doc/development/guide/logging.rst +++ b/doc/development/guide/logging.rst @@ -298,7 +298,7 @@ process, and surrounding operations: # Get logger for use by this script only. logger = logging.getLogger(__name__) - dev_name = "default.qubit.jax" + dev_name = "default.qubit" num_wires = 2 num_shots = None diff --git a/doc/development/release_notes.md b/doc/development/release_notes.md index a1f5587597c..07ad0d279f7 100644 --- a/doc/development/release_notes.md +++ b/doc/development/release_notes.md @@ -5,6 +5,8 @@ This page contains the release notes for PennyLane. .. mdinclude:: ../releases/changelog-dev.md +.. mdinclude:: ../releases/changelog-0.38.1.md + .. mdinclude:: ../releases/changelog-0.38.0.md .. mdinclude:: ../releases/changelog-0.37.0.md diff --git a/doc/introduction/interfaces/jax.rst b/doc/introduction/interfaces/jax.rst index baf643a5716..04cd22703b5 100644 --- a/doc/introduction/interfaces/jax.rst +++ b/doc/introduction/interfaces/jax.rst @@ -54,7 +54,7 @@ a JAX-capable QNode in PennyLane. Simply specify the ``interface='jax'`` keyword .. code-block:: python - dev = qml.device('default.qubit.jax', wires=2) + dev = qml.device('default.qubit', wires=2) @qml.qnode(dev, interface='jax') def circuit1(phi, theta): @@ -85,7 +85,7 @@ For example: .. code-block:: python - dev = qml.device('default.qubit.jax', wires=2) + dev = qml.device('default.qubit', wires=2) @qml.qnode(dev, interface='jax') def circuit3(phi, theta): @@ -119,7 +119,7 @@ the ``@jax.jit`` decorator can be directly applied to the QNode. .. code-block:: python - dev = qml.device('default.qubit.jax', wires=2) + dev = qml.device('default.qubit', wires=2) @jax.jit # QNode calls will now be jitted, and should run faster. @qml.qnode(dev, interface='jax') @@ -176,7 +176,7 @@ Example: # Device construction should happen inside a `jax.jit` decorated # method when using a PRNGKey. - dev = qml.device('default.qubit.jax', wires=2, prng_key=key, shots=100) + dev = qml.device('default.qubit', wires=2, prng_key=key, shots=100) @qml.qnode(dev, interface='jax', diff_method=None) diff --git a/doc/releases/changelog-0.38.0.md b/doc/releases/changelog-0.38.0.md index 391f9dbf36d..eb817d4275e 100644 --- a/doc/releases/changelog-0.38.0.md +++ b/doc/releases/changelog-0.38.0.md @@ -1,6 +1,6 @@ :orphan: -# Release 0.38.0 (current release) +# Release 0.38.0

New features since last release

diff --git a/doc/releases/changelog-0.38.1.md b/doc/releases/changelog-0.38.1.md new file mode 100644 index 00000000000..aab0ef38311 --- /dev/null +++ b/doc/releases/changelog-0.38.1.md @@ -0,0 +1,14 @@ +:orphan: + +# Release 0.38.1 (current release) + +

Bug fixes 🐛

+ +* Fix float-to-complex casting in various places across PennyLane. + [(#6260)](https://github.com/PennyLaneAI/pennylane/pull/6260) + +

Contributors ✍️

+ +This release contains contributions from (in alphabetical order): + +Mudit Pandey \ No newline at end of file diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index a2809a619bf..488fd8d37e8 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -36,6 +36,9 @@ `from pennylane.capture.primitives import *`. [(#6129)](https://github.com/PennyLaneAI/pennylane/pull/6129) +* `FermiWord` and `FermiSentence` classes now have methods to compute adjoints. + [(#6166)](https://github.com/PennyLaneAI/pennylane/pull/6166) + * The `SampleMP.process_samples` method is updated to support using JAX tracers for samples, allowing compatiblity with Catalyst workflows. [(#6211)](https://github.com/PennyLaneAI/pennylane/pull/6211) @@ -56,9 +59,12 @@ * Remove support for Python 3.9. [(#6223)](https://github.com/PennyLaneAI/pennylane/pull/6223) -* `DefaultQubitTF` and `DefaultQubitTorch` are removed. Please use `default.qubit` for all interfaces. +* `DefaultQubitTF`, `DefaultQubitTorch`, `DefaultQubitJax`, and `DefaultQubitAutograd` are removed. + Please use `default.qubit` for all interfaces. [(#6207)](https://github.com/PennyLaneAI/pennylane/pull/6207) [(#6208)](https://github.com/PennyLaneAI/pennylane/pull/6208) + [(#6209)](https://github.com/PennyLaneAI/pennylane/pull/6209) + [(#6210)](https://github.com/PennyLaneAI/pennylane/pull/6210) * `expand_fn`, `max_expansion`, `override_shots`, and `device_batch_transform` are removed from the signature of `qml.execute`. @@ -74,12 +80,27 @@ Please use `qml.transforms.split_non_commuting` instead. [(#6204)](https://github.com/PennyLaneAI/pennylane/pull/6204) +* `Operator.expand` is now removed. Use `qml.tape.QuantumScript(op.deocomposition())` instead. + [(#6227)](https://github.com/PennyLaneAI/pennylane/pull/6227) + +

Deprecations 👋

+* `Device`, `QubitDevice`, and `QutritDevice` will no longer be accessible via top-level import in v0.40. + They will still be accessible as `qml.devices.LegacyDevice`, `qml.devices.QubitDevice`, and `qml.devices.QutritDevice` + respectively. + [(#6238)](https://github.com/PennyLaneAI/pennylane/pull/6238/) + +* `QNode.gradient_fn` is deprecated. Please use `QNode.diff_method` and `QNode.get_gradient_fn` instead. + [(#6244)](https://github.com/PennyLaneAI/pennylane/pull/6244) +

Documentation 📝

Bug fixes 🐛

+* Fix a bug where zero-valued JVPs were calculated wrongly in the presence of shot vectors. + [(#6219)](https://github.com/PennyLaneAI/pennylane/pull/6219) + * Fix `qml.PrepSelPrep` template to work with `torch`: [(#6191)](https://github.com/PennyLaneAI/pennylane/pull/6191) @@ -92,6 +113,9 @@ * The ``qml.Qubitization`` template now orders the ``control`` wires first and the ``hamiltonian`` wires second, which is the expected according to other templates. [(#6229)](https://github.com/PennyLaneAI/pennylane/pull/6229) +* The ``qml.FABLE`` template now returns the correct value when JIT is enabled. + [(#6263)](https://github.com/PennyLaneAI/pennylane/pull/6263) + *

Contributors ✍️

This release contains contributions from (in alphabetical order): diff --git a/pennylane/__init__.py b/pennylane/__init__.py index 7f029461dea..3db66c953d2 100644 --- a/pennylane/__init__.py +++ b/pennylane/__init__.py @@ -16,8 +16,6 @@ PennyLane can be directly imported. """ -import numpy as _np - from pennylane.boolean_fn import BooleanFn import pennylane.numpy @@ -180,13 +178,30 @@ def __getattr__(name): if name == "plugin_devices": return pennylane.devices.device_constructor.plugin_devices + from warnings import warn # pylint: disable=import-outside-toplevel + if name == "QubitDevice": + warn( + "QubitDevice will no longer be accessible top level. Please access " + " the class as pennylane.devices.QubitDevice", + PennyLaneDeprecationWarning, + ) return pennylane.devices._qubit_device.QubitDevice # pylint:disable=protected-access if name == "QutritDevice": + warn( + "QutritDevice will no longer be accessible top level. Please access " + " the class as pennylane.devices.QutritDevice", + PennyLaneDeprecationWarning, + ) return pennylane.devices._qutrit_device.QutritDevice # pylint:disable=protected-access if name == "Device": + warn( + "Device will no longer be accessible top level. Please access " + " the class as pennylane.devices.LegacyDevice", + PennyLaneDeprecationWarning, + ) return pennylane.devices._legacy_device.Device # pylint:disable=protected-access raise AttributeError(f"module 'pennylane' has no attribute '{name}'") diff --git a/pennylane/_version.py b/pennylane/_version.py index 4ddbd563982..a6ead820881 100644 --- a/pennylane/_version.py +++ b/pennylane/_version.py @@ -16,4 +16,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "0.39.0-dev10" +__version__ = "0.39.0-dev14" diff --git a/pennylane/devices/__init__.py b/pennylane/devices/__init__.py index 87dc6c3f6e6..a542ba7df1d 100644 --- a/pennylane/devices/__init__.py +++ b/pennylane/devices/__init__.py @@ -27,8 +27,6 @@ default_qubit default_qubit_legacy - default_qubit_jax - default_qubit_autograd default_gaussian default_mixed default_qutrit @@ -154,9 +152,6 @@ def execute(self, circuits, execution_config = qml.devices.DefaultExecutionConfi from .default_qubit import DefaultQubit from .legacy_facade import LegacyDeviceFacade -# DefaultQubitTF and DefaultQubitAutograd not imported here since this -# would lead to an automatic import of tensorflow and autograd, which are -# not PennyLane core dependencies. # DefaultTensor is not imported here to avoid warnings # from quimb in case it is installed on the system. from .default_qubit_legacy import DefaultQubitLegacy diff --git a/pennylane/devices/_qubit_device.py b/pennylane/devices/_qubit_device.py index a54c2985290..ef9f41f0378 100644 --- a/pennylane/devices/_qubit_device.py +++ b/pennylane/devices/_qubit_device.py @@ -1110,7 +1110,11 @@ def classical_shadow(self, obs, circuit): n_snapshots = self.shots seed = obs.seed - with qml.workflow.set_shots(self, shots=1): + original_shots = self.shots + original_shot_vector = self._shot_vector + + try: + self.shots = 1 # slow implementation but works for all devices n_qubits = len(wires) mapped_wires = np.array(self.map_wires(wires)) @@ -1139,6 +1143,10 @@ def classical_shadow(self, obs, circuit): ) outcomes[t] = self.generate_samples()[0][mapped_wires] + finally: + self.shots = original_shots + # pylint: disable=attribute-defined-outside-init + self._shot_vector = original_shot_vector return self._cast(self._stack([outcomes, recipes]), dtype=np.int8) diff --git a/pennylane/devices/default_qubit_autograd.py b/pennylane/devices/default_qubit_autograd.py deleted file mode 100644 index abcc6e0452f..00000000000 --- a/pennylane/devices/default_qubit_autograd.py +++ /dev/null @@ -1,141 +0,0 @@ -# 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. -"""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): - r"""Simulator plugin based on ``"default.qubit.legacy"``, written using Autograd. - - **Short name:** ``default.qubit.autograd`` - - This device provides a pure-state qubit simulator written using Autograd. As a result, it - supports classical backpropagation as a means to compute the gradient. This can be faster than - the parameter-shift rule for analytic quantum gradients when the number of parameters to be - optimized is large. - - To use this device, you will need to install Autograd: - - .. code-block:: console - - 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 - (``diff_method="backprop"``) with the Autograd interface. This is the default method of - differentiation when creating a QNode with this device. - - Using this method, the created QNode is a 'white-box', and is - tightly integrated with your Autograd computation: - - >>> dev = qml.device("default.qubit.autograd", wires=1) - >>> @qml.qnode(dev, interface="autograd", diff_method="backprop") - ... def circuit(x): - ... qml.RX(x[1], wires=0) - ... qml.Rot(x[0], x[1], x[2], wires=0) - ... return qml.expval(qml.Z(0)) - >>> weights = np.array([0.2, 0.5, 0.1], requires_grad=True) - >>> grad_fn = qml.grad(circuit) - >>> print(grad_fn(weights)) - array([-2.2526717e-01 -1.0086454e+00 1.3877788e-17]) - - There are a couple of things to keep in mind when using the ``"backprop"`` - differentiation method for QNodes: - - * You must use the ``"autograd"`` interface for classical backpropagation, as Autograd is - used as the device backend. - - * Only exact expectation values, variances, and probabilities are differentiable. - When instantiating the device with ``analytic=False``, differentiating QNode - outputs will result in an error. - - Args: - wires (int): the number of wires to initialize the device with - shots (None, int): How many times the circuit should be evaluated (or sampled) to estimate - the expectation values. Defaults to ``None`` if not specified, which means that the device - returns analytical results. - analytic (bool): Indicates if the device should calculate expectations - and variances analytically. In non-analytic mode, the ``diff_method="backprop"`` - QNode differentiation method is not supported and it is recommended to consider - switching device to ``default.qubit`` and using ``diff_method="parameter-shift"``. - """ - - name = "Default qubit (Autograd) PennyLane plugin" - short_name = "default.qubit.autograd" - - _dot = staticmethod(pnp.dot) - _abs = staticmethod(pnp.abs) - _reduce_sum = staticmethod(lambda array, axes: pnp.sum(array, axis=tuple(axes))) - _reshape = staticmethod(pnp.reshape) - _flatten = staticmethod(lambda array: array.flatten()) - _einsum = staticmethod(pnp.einsum) - _cast = staticmethod(pnp.asarray) - _transpose = staticmethod(pnp.transpose) - _tensordot = staticmethod(pnp.tensordot) - _conj = staticmethod(pnp.conj) - _real = staticmethod(pnp.real) - _imag = staticmethod(pnp.imag) - _roll = staticmethod(pnp.roll) - _stack = staticmethod(pnp.stack) - _size = staticmethod(pnp.size) - _ndim = staticmethod(pnp.ndim) - - @staticmethod - def _asarray(array, dtype=None): - return pnp.asarray(array, dtype=dtype) - - @staticmethod - 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) - - # prevent using special apply methods for these gates due to slowdown in Autograd - # implementation - del self._apply_ops["PauliY"] - del self._apply_ops["Hadamard"] - del self._apply_ops["CZ"] - - @classmethod - def capabilities(cls): - capabilities = super().capabilities().copy() - capabilities.update(passthru_interface="autograd") - return capabilities - - @staticmethod - def _scatter(indices, array, new_dimensions): - new_array = pnp.zeros(new_dimensions, dtype=array.dtype.type) - new_array[indices] = array - return new_array diff --git a/pennylane/devices/default_qubit_jax.py b/pennylane/devices/default_qubit_jax.py deleted file mode 100644 index 0cae4827e41..00000000000 --- a/pennylane/devices/default_qubit_jax.py +++ /dev/null @@ -1,371 +0,0 @@ -# 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. -"""This module contains a jax implementation of the :class:`~.DefaultQubitLegacy` -reference plugin. -""" -import warnings - -# pylint: disable=ungrouped-imports -import numpy as np - -import pennylane as qml -from pennylane.devices import DefaultQubitLegacy -from pennylane.pulse import ParametrizedEvolution -from pennylane.typing import TensorLike - -try: - import jax - import jax.numpy as jnp - from jax.experimental.ode import odeint - - from pennylane.pulse.parametrized_hamiltonian_pytree import ParametrizedHamiltonianPytree - -except ImportError as e: # pragma: no cover - raise ImportError("default.qubit.jax device requires installing jax>0.3.20") from e - - -class DefaultQubitJax(DefaultQubitLegacy): - r"""Simulator plugin based on ``"default.qubit.legacy"``, written using jax. - - **Short name:** ``default.qubit.jax`` - - This device provides a pure-state qubit simulator written using jax. As a result, it - supports classical backpropagation as a means to compute the gradient. This can be faster than - the parameter-shift rule for analytic quantum gradients when the number of parameters to be - optimized is large. - - To use this device, you will need to install jax: - - .. code-block:: console - - 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 - (``diff_method="backprop"``) with the JAX interface. This is the default method of - differentiation when creating a QNode with this device. - - Using this method, the created QNode is a 'white-box', and is - tightly integrated with your JAX computation: - - >>> dev = qml.device("default.qubit.jax", wires=1) - >>> @qml.qnode(dev, interface="jax", diff_method="backprop") - ... def circuit(x): - ... qml.RX(x[1], wires=0) - ... qml.Rot(x[0], x[1], x[2], wires=0) - ... return qml.expval(qml.Z(0)) - >>> weights = jnp.array([0.2, 0.5, 0.1]) - >>> grad_fn = jax.grad(circuit) - >>> print(grad_fn(weights)) - array([-2.2526717e-01 -1.0086454e+00 1.3877788e-17]) - - There are a couple of things to keep in mind when using the ``"backprop"`` - differentiation method for QNodes: - - * You must use the ``"jax"`` interface for classical backpropagation, as JAX is - used as the device backend. - - .. details:: - :title: Usage Details - - JAX does randomness in a special way when compared to NumPy, in that all randomness needs to - be seeded. While we handle this for you automatically in op-by-op mode, when using ``jax.jit``, - the automatically generated seed gets constantant compiled. - - Example: - - .. code-block:: python - - dev = qml.device("default.qubit.jax", wires=1, shots=10) - - @jax.jit - @qml.qnode(dev, interface="jax", diff_method="backprop") - def circuit(): - qml.Hadamard(0) - return qml.sample(qml.Z(0)) - - a = circuit() - b = circuit() # Bad! b will be the exact same samples as a. - - - To fix this, you should wrap your qnode in another function that takes a PRNGKey, and pass - that in during your device construction. - - .. code-block:: python - - @jax.jit - def keyed_circuit(key): - dev = qml.device("default.qubit.jax", prng_key=key, wires=1, shots=10) - @qml.qnode(dev, interface="jax", diff_method="backprop") - def circuit(): - qml.Hadamard(0) - return qml.sample(qml.Z(0)) - return circuit() - - key1 = jax.random.PRNGKey(0) - key2 = jax.random.PRNGKey(1) - a = keyed_circuit(key1) - b = keyed_circuit(key2) # b will be different samples now. - - Check out the `JAX random documentation `__ - for more information. - - Args: - wires (int): The number of wires to initialize the device with. - shots (None, int): How many times the circuit should be evaluated (or sampled) to estimate - the expectation values. Defaults to ``None`` if not specified, which means that the device - returns analytical results. - analytic (bool): Indicates if the device should calculate expectations - and variances analytically. In non-analytic mode, the ``diff_method="backprop"`` - QNode differentiation method is not supported and it is recommended to consider - switching device to ``default.qubit`` and using ``diff_method="parameter-shift"``. - Or keeping ``default.qubit.jax`` but switching to - ``diff_method=qml.gradients.stoch_pulse_grad`` for pulse programming. - prng_key (Optional[jax.random.PRNGKey]): An optional ``jax.random.PRNGKey``. This is the key to the - pseudo random number generator. If None, a random key will be generated. - - """ - - name = "Default qubit (jax) PennyLane plugin" - short_name = "default.qubit.jax" - - _asarray = staticmethod(jnp.array) - _dot = staticmethod(jnp.dot) - _abs = staticmethod(jnp.abs) - _reduce_sum = staticmethod(lambda array, axes: jnp.sum(array, axis=tuple(axes))) - _reshape = staticmethod(jnp.reshape) - _flatten = staticmethod(lambda array: array.ravel()) - _einsum = staticmethod(jnp.einsum) - _cast = staticmethod(jnp.array) - _transpose = staticmethod(jnp.transpose) - _tensordot = staticmethod( - lambda a, b, axes: jnp.tensordot( - a, b, axes if isinstance(axes, int) else list(map(tuple, axes)) - ) - ) - _conj = staticmethod(jnp.conj) - _real = staticmethod(jnp.real) - _imag = staticmethod(jnp.imag) - _roll = staticmethod(jnp.roll) - _stack = staticmethod(jnp.stack) - _const_mul = staticmethod(jnp.multiply) - _size = staticmethod(jnp.size) - _ndim = staticmethod(jnp.ndim) - - 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 - else: - c_dtype = jnp.complex64 - r_dtype = jnp.float32 - super().__init__(wires, r_dtype=r_dtype, c_dtype=c_dtype, shots=shots, analytic=analytic) - - # prevent using special apply methods for these gates due to slowdown in jax - # implementation - del self._apply_ops["PauliY"] - del self._apply_ops["Hadamard"] - del self._apply_ops["CZ"] - self._prng_key = prng_key - - @classmethod - def capabilities(cls): - capabilities = super().capabilities().copy() - capabilities.update(passthru_interface="jax") - return capabilities - - def _apply_parametrized_evolution(self, state: TensorLike, operation: ParametrizedEvolution): - # given that wires is a static value (it is not a tracer), we can use an if statement - if ( - 2 * len(operation.wires) > self.num_wires - and not operation.hyperparameters["complementary"] - ): - # the device state vector contains less values than the operation matrix --> evolve state - return self._evolve_state_vector_under_parametrized_evolution(state, operation) - # the device state vector contains more/equal values than the operation matrix --> evolve matrix - return self._apply_operation(state, operation) - - def _evolve_state_vector_under_parametrized_evolution( - self, state: TensorLike, operation: ParametrizedEvolution - ): - """Uses an odeint solver to compute the evolution of the input ``state`` under the given - ``ParametrizedEvolution`` operation. - - Args: - state (array[complex]): input state - operation (ParametrizedEvolution): operation to apply on the state - - Raises: - ValueError: If the parameters and time windows of the ``ParametrizedEvolution`` are - not defined. - - Returns: - _type_: _description_ - """ - if operation.data is None or operation.t is None: - raise ValueError( - "The parameters and the time window are required to execute a ParametrizedEvolution " - "You can update these values by calling the ParametrizedEvolution class: EV(params, t)." - ) - - state = self._flatten(state) - - with jax.ensure_compile_time_eval(): - H_jax = ParametrizedHamiltonianPytree.from_hamiltonian( - operation.H, dense=operation.dense, wire_order=self.wires - ) - - def fun(y, t): - """dy/dt = -i H(t) y""" - return (-1j * H_jax(operation.data, t=t)) @ y - - result = odeint(fun, state, operation.t, **operation.odeint_kwargs) - out_shape = [2] * self.num_wires - if operation.hyperparameters["return_intermediate"]: - return self._reshape(result, [-1] + out_shape) - return self._reshape(result[-1], out_shape) - - @staticmethod - def _scatter(indices, array, new_dimensions): - new_array = jnp.zeros(new_dimensions, dtype=array.dtype.type) - new_array = new_array.at[indices].set(array) - return new_array - - def sample_basis_states(self, number_of_states, state_probability): - """Sample from the computational basis states based on the state - probability. - - This is an auxiliary method to the generate_samples method. - - Args: - number_of_states (int): the number of basis states to sample from - - Returns: - List[int]: the sampled basis states - """ - if self.shots is None: - raise qml.QuantumFunctionError( - "The number of shots has to be explicitly set on the device " - "when using sample-based measurements." - ) - - shots = self.shots - - if self._prng_key is None: - # Assuming op-by-op, so we'll just make one. - key = jax.random.PRNGKey(np.random.randint(0, 2**31)) - else: - key = self._prng_key - if jnp.ndim(state_probability) == 2: - # Produce separate keys for each of the probabilities along the broadcasted axis - keys = [] - for _ in state_probability: - key, subkey = jax.random.split(key) - keys.append(subkey) - return jnp.array( - [ - jax.random.choice(_key, number_of_states, shape=(shots,), p=prob) - for _key, prob in zip(keys, state_probability) - ] - ) - return jax.random.choice(key, number_of_states, shape=(shots,), p=state_probability) - - @staticmethod - def states_to_binary(samples, num_wires, dtype=jnp.int32): - """Convert basis states from base 10 to binary representation. - - This is an auxiliary method to the generate_samples method. - - Args: - samples (List[int]): samples of basis states in base 10 representation - num_wires (int): the number of qubits - dtype (type): Type of the internal integer array to be used. Can be - important to specify for large systems for memory allocation - purposes. - - Returns: - List[int]: basis states in binary representation - """ - powers_of_two = 1 << jnp.arange(num_wires, dtype=dtype) - states_sampled_base_ten = samples[..., None] & powers_of_two - return (states_sampled_base_ten > 0).astype(dtype)[..., ::-1] - - @staticmethod - def _count_unbinned_samples(indices, batch_size, dim): - """Count the occurences of sampled indices and convert them to relative - counts in order to estimate their occurence probability.""" - - shape = (dim + 1,) if batch_size is None else (batch_size, dim + 1) - prob = qml.math.convert_like(jnp.zeros(shape, dtype=jnp.float64), indices) - if batch_size is None: - basis_states, counts = jnp.unique(indices, return_counts=True, size=dim, fill_value=-1) - for state, count in zip(basis_states, counts): - prob = prob.at[state].set(count / len(indices)) - # resize prob which discards the 'filled values' - return prob[:-1] - - for i, idx in enumerate(indices): - basis_states, counts = jnp.unique(idx, return_counts=True, size=dim, fill_value=-1) - for state, count in zip(basis_states, counts): - prob = prob.at[i, state].set(count / len(idx)) - - # resize prob which discards the 'filled values' - return prob[:, :-1] - - @staticmethod - def _count_binned_samples(indices, batch_size, dim, bin_size, num_bins): - """Count the occurences of bins of sampled indices and convert them to relative - counts in order to estimate their occurence probability per bin.""" - - # extend the probability vectors to store 'filled values' - shape = (dim + 1, num_bins) if batch_size is None else (batch_size, dim + 1, num_bins) - prob = qml.math.convert_like(jnp.zeros(shape, dtype=jnp.float64), indices) - if batch_size is None: - indices = indices.reshape((num_bins, bin_size)) - - # count the basis state occurrences, and construct the probability vector for each bin - for b, idx in enumerate(indices): - idx = qml.math.convert_like(idx, indices) - basis_states, counts = jnp.unique(idx, return_counts=True, size=dim, fill_value=-1) - for state, count in zip(basis_states, counts): - prob = prob.at[state, b].set(count / bin_size) - - # resize prob which discards the 'filled values' - return prob[:-1] - - indices = indices.reshape((batch_size, num_bins, bin_size)) - # count the basis state occurrences, and construct the probability vector - # for each bin and broadcasting index - for i, _indices in enumerate(indices): # First iterate over broadcasting dimension - for b, idx in enumerate(_indices): # Then iterate over bins dimension - idx = qml.math.convert_like(idx, indices) - basis_states, counts = jnp.unique(idx, return_counts=True, size=dim, fill_value=-1) - for state, count in zip(basis_states, counts): - prob = prob.at[i, state, b].set(count / bin_size) - # resize prob which discards the 'filled values' - return prob[:, :-1] diff --git a/pennylane/devices/default_qubit_legacy.py b/pennylane/devices/default_qubit_legacy.py index 8aa98b607ba..868c426d47f 100644 --- a/pennylane/devices/default_qubit_legacy.py +++ b/pennylane/devices/default_qubit_legacy.py @@ -713,10 +713,7 @@ def capabilities(cls): supports_analytic_computation=True, supports_broadcasting=True, returns_state=True, - passthru_devices={ - "autograd": "default.qubit.autograd", - "jax": "default.qubit.jax", - }, + passthru_devices={}, ) return capabilities diff --git a/pennylane/devices/legacy_facade.py b/pennylane/devices/legacy_facade.py index a35aa4b25f5..41c1e0dea2c 100644 --- a/pennylane/devices/legacy_facade.py +++ b/pennylane/devices/legacy_facade.py @@ -15,7 +15,6 @@ Defines a LegacyDeviceFacade class for converting legacy devices to the new interface. """ -import warnings # pylint: disable=not-callable, unused-argument from contextlib import contextmanager @@ -318,79 +317,6 @@ def supports_derivatives(self, execution_config=None, circuit=None) -> bool: return False - # pylint: disable=protected-access - def _create_temp_device(self, batch): - """Create a temporary device for use in a backprop execution.""" - params = [] - for t in batch: - params.extend(t.get_parameters(trainable_only=False)) - interface = qml.math.get_interface(*params) - if interface == "numpy": - return self._device - - mapped_interface = qml.workflow.execution.INTERFACE_MAP.get(interface, interface) - - backprop_interface = self._device.capabilities().get("passthru_interface", None) - if mapped_interface == backprop_interface: - return self._device - - backprop_devices = self._device.capabilities().get("passthru_devices", None) - - if backprop_devices is None: - raise qml.DeviceError(f"Device {self} does not support backpropagation.") - - if backprop_devices[mapped_interface] == self._device.short_name: - return self._device - - if self.target_device.short_name != "default.qubit.legacy": - warnings.warn( - "The switching of devices for backpropagation is now deprecated in v0.38 and " - "will be removed in v0.39, as this behavior was developed purely for the " - "deprecated default.qubit.legacy.", - qml.PennyLaneDeprecationWarning, - ) - - # create new backprop device - expand_fn = self._device.expand_fn - batch_transform = self._device.batch_transform - if hasattr(self._device, "_debugger"): - debugger = self._device._debugger - else: - debugger = "No debugger" - tracker = self._device.tracker - - with warnings.catch_warnings(): - warnings.filterwarnings( - action="ignore", - category=qml.PennyLaneDeprecationWarning, - message=r"use 'default.qubit'", - ) - # we already warned about backprop device switching - new_device = qml.device( - backprop_devices[mapped_interface], - wires=self._device.wires, - shots=self._device.shots, - ).target_device - - new_device.expand_fn = expand_fn - new_device.batch_transform = batch_transform - if debugger != "No debugger": - new_device._debugger = debugger - new_device.tracker = tracker - - return new_device - - # pylint: disable=protected-access - def _update_original_device(self, temp_device): - """After performing an execution with a backprop device, update the state of the original device.""" - # Update for state vector simulators that have the _pre_rotated_state attribute - if hasattr(self._device, "_pre_rotated_state"): - self._device._pre_rotated_state = temp_device._pre_rotated_state - - # Update for state vector simulators that have the _state attribute - if hasattr(self._device, "_state"): - self._device._state = temp_device._state - def _validate_backprop_method(self, tape): if tape.shots: return False @@ -441,11 +367,7 @@ def _validate_device_method(self, _): return self._device.capabilities().get("provides_jacobian", False) def execute(self, circuits, execution_config=DefaultExecutionConfig): - dev = ( - self._create_temp_device(circuits) - if execution_config.gradient_method == "backprop" - else self._device - ) + dev = self.target_device kwargs = {} if dev.capabilities().get("supports_mid_measure", False): @@ -453,16 +375,10 @@ def execute(self, circuits, execution_config=DefaultExecutionConfig): first_shot = circuits[0].shots if all(t.shots == first_shot for t in circuits): - results = _set_shots(dev, first_shot)(dev.batch_execute)(circuits, **kwargs) - else: - results = tuple( - _set_shots(dev, t.shots)(dev.batch_execute)((t,), **kwargs)[0] for t in circuits - ) - - if dev is not self._device: - self._update_original_device(dev) - - return results + return _set_shots(dev, first_shot)(dev.batch_execute)(circuits, **kwargs) + return tuple( + _set_shots(dev, t.shots)(dev.batch_execute)((t,), **kwargs)[0] for t in circuits + ) def execute_and_compute_derivatives(self, circuits, execution_config=DefaultExecutionConfig): first_shot = circuits[0].shots diff --git a/pennylane/devices/qubit/apply_operation.py b/pennylane/devices/qubit/apply_operation.py index 8e082ee9820..dccf49ef11b 100644 --- a/pennylane/devices/qubit/apply_operation.py +++ b/pennylane/devices/qubit/apply_operation.py @@ -71,7 +71,7 @@ def apply_operation_einsum(op: qml.operation.Operator, state, is_state_batched: Returns: array[complex]: output_state """ - mat = op.matrix() + mat = qml.math.cast_like(op.matrix(), 1j) total_indices = len(state.shape) - is_state_batched num_indices = len(op.wires) @@ -114,7 +114,7 @@ def apply_operation_tensordot(op: qml.operation.Operator, state, is_state_batche Returns: array[complex]: output_state """ - mat = op.matrix() + mat = qml.math.cast_like(op.matrix(), 1j) total_indices = len(state.shape) - is_state_batched num_indices = len(op.wires) diff --git a/pennylane/devices/tests/conftest.py b/pennylane/devices/tests/conftest.py index 8e5e9740da1..f462138c4d5 100755 --- a/pennylane/devices/tests/conftest.py +++ b/pennylane/devices/tests/conftest.py @@ -36,7 +36,6 @@ # List of all devices that are included in PennyLane LIST_CORE_DEVICES = { "default.qubit", - "default.qubit.autograd", } @@ -111,12 +110,6 @@ def validate_diff_method(device, diff_method, device_kwargs): if diff_method == "backprop" and device_kwargs.get("shots") is not None: pytest.skip(reason="test should only be run in analytic mode") dev = device(1) - if isinstance(dev, qml.Device): - passthru_devices = dev.capabilities().get("passthru_devices") - if diff_method == "backprop" and passthru_devices is None: - pytest.skip(reason="device does not support backprop") - return - config = qml.devices.ExecutionConfig(gradient_method=diff_method) if not dev.supports_derivatives(execution_config=config): pytest.skip(reason="device does not support diff_method") @@ -142,12 +135,6 @@ def _device(wires): f"plugin and all of its dependencies must be installed." ) - if isinstance(dev, qml.Device): - capabilities = dev.capabilities() - if capabilities.get("model", None) != "qubit": - # exit the tests if device based on cv model (currently not supported) - pytest.exit("The device test suite currently only runs on qubit-based devices.") - return dev return _device diff --git a/pennylane/devices/tests/test_compare_default_qubit.py b/pennylane/devices/tests/test_compare_default_qubit.py index 05b2b6509e8..75e91809e2a 100755 --- a/pennylane/devices/tests/test_compare_default_qubit.py +++ b/pennylane/devices/tests/test_compare_default_qubit.py @@ -38,9 +38,6 @@ def test_hermitian_expectation(self, device, tol, benchmark): if dev.shots: pytest.skip("Device is in non-analytical mode.") - if isinstance(dev, qml.Device) and "Hermitian" not in dev.observables: - pytest.skip("Device does not support the Hermitian observable.") - if dev.name == "default.qubit": pytest.skip("Device is default.qubit.") @@ -107,7 +104,7 @@ def test_projector_expectation(self, device, state, tol, benchmark): if dev.shots: pytest.skip("Device is in non-analytical mode.") - if isinstance(dev, qml.Device) and "Projector" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Projector" not in dev.observables: pytest.skip("Device does not support the Projector observable.") if dev.name == "default.qubit": diff --git a/pennylane/devices/tests/test_gates.py b/pennylane/devices/tests/test_gates.py index b948f6962c4..7226e021994 100644 --- a/pennylane/devices/tests/test_gates.py +++ b/pennylane/devices/tests/test_gates.py @@ -358,7 +358,7 @@ def test_supported_gates_can_be_implemented(self, device_kwargs, operation): device_kwargs["wires"] = 4 # maximum size of current gates dev = qml.device(**device_kwargs) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if operation not in dev.operations: pytest.skip("operation not supported.") else: @@ -395,7 +395,7 @@ def test_basis_state(self, device, basis_state, tol, skip_if): """Test basis state initialization.""" n_wires = 4 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"returns_probs": False}) @qml.qnode(dev) @@ -413,7 +413,7 @@ def test_state_prep(self, device, init_state, tol, skip_if): """Test StatePrep initialisation.""" n_wires = 1 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"returns_probs": False}) rnd_state = init_state(n_wires) @@ -433,7 +433,7 @@ def test_single_qubit_no_parameters(self, device, init_state, op, mat, tol, skip """Test PauliX application.""" n_wires = 1 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"returns_probs": False}) rnd_state = init_state(n_wires) @@ -457,7 +457,7 @@ def test_single_qubit_parameters( """Test single qubit gates taking a single scalar argument.""" n_wires = 1 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"returns_probs": False}) rnd_state = init_state(n_wires) @@ -477,7 +477,7 @@ def test_rotation(self, device, init_state, tol, skip_if, benchmark): """Test three axis rotation gate.""" n_wires = 1 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"returns_probs": False}) rnd_state = init_state(n_wires) @@ -501,7 +501,7 @@ def test_two_qubit_no_parameters(self, device, init_state, op, mat, tol, skip_if """Test two qubit gates.""" n_wires = 2 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"returns_probs": False}) if not dev.supports_operation(op(wires=range(n_wires)).name): pytest.skip("op not supported") @@ -527,7 +527,7 @@ def test_two_qubit_parameters( """Test parametrized two qubit gates taking a single scalar argument.""" n_wires = 2 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"returns_probs": False}) rnd_state = init_state(n_wires) @@ -549,7 +549,7 @@ def test_qubit_unitary(self, device, init_state, mat, tol, skip_if, benchmark): n_wires = int(np.log2(len(mat))) dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if "QubitUnitary" not in dev.operations: pytest.skip("Skipped because device does not support QubitUnitary.") @@ -574,7 +574,7 @@ def test_special_unitary(self, device, init_state, theta_, tol, skip_if, benchma n_wires = int(np.log(len(theta_) + 1) / np.log(4)) dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if "SpecialUnitary" not in dev.operations: pytest.skip("Skipped because device does not support SpecialUnitary.") @@ -603,7 +603,7 @@ def test_three_qubit_no_parameters(self, device, init_state, op, mat, tol, skip_ n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"returns_probs": False}) rnd_state = init_state(n_wires) diff --git a/pennylane/devices/tests/test_gates_with_expval.py b/pennylane/devices/tests/test_gates_with_expval.py index 7239ab2bcd3..f2e7e8951d6 100755 --- a/pennylane/devices/tests/test_gates_with_expval.py +++ b/pennylane/devices/tests/test_gates_with_expval.py @@ -256,7 +256,7 @@ def test_supported_gate_two_wires_no_parameters(self, device, tol, name, expecte dev = device(n_wires) op = getattr(qml.ops, name) - if isinstance(dev, qml.Device) and not dev.supports_operation(op): + if isinstance(dev, qml.devices.LegacyDevice) and not dev.supports_operation(op): pytest.skip("operation not supported") @qml.qnode(dev) diff --git a/pennylane/devices/tests/test_measurements.py b/pennylane/devices/tests/test_measurements.py index 2a81366ba36..cde18014781 100644 --- a/pennylane/devices/tests/test_measurements.py +++ b/pennylane/devices/tests/test_measurements.py @@ -116,7 +116,7 @@ def test_supported_observables_can_be_implemented(self, device_kwargs, observabl if dev.shots and observable == "SparseHamiltonian": pytest.skip("SparseHamiltonian only supported in analytic mode") - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): assert hasattr(dev, "observables") if observable not in dev.observables: pytest.skip("observable not supported") @@ -313,7 +313,7 @@ def test_hermitian_expectation(self, device, tol): n_wires = 2 dev = device(n_wires) - if isinstance(dev, qml.Device) and "Hermitian" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Hermitian" not in dev.observables: pytest.skip("Skipped because device does not support the Hermitian observable.") theta = 0.432 @@ -342,7 +342,7 @@ def test_projector_expectation(self, device, tol): n_wires = 2 dev = device(n_wires) - if isinstance(dev, qml.Device) and "Projector" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Projector" not in dev.observables: pytest.skip("Skipped because device does not support the Projector observable.") theta = 0.732 @@ -380,7 +380,7 @@ def test_multi_mode_hermitian_expectation(self, device, tol): n_wires = 2 dev = device(n_wires) - if isinstance(dev, qml.Device) and "Hermitian" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Hermitian" not in dev.observables: pytest.skip("Skipped because device does not support the Hermitian observable.") theta = 0.432 @@ -426,7 +426,7 @@ def circuit(): def test_op_arithmetic_matches_default_qubit(self, o, device, tol): """Test that devices (which support the observable) match default.qubit results.""" dev = device(2) - if isinstance(dev, qml.Device) and o.name not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and o.name not in dev.observables: pytest.skip(f"Skipped because device does not support the {o.name} observable.") def circuit(): @@ -448,7 +448,7 @@ def test_paulix_pauliy(self, device, tol, skip_if): """Test that a tensor product involving PauliX and PauliY works correctly""" n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) theta = 0.432 @@ -473,7 +473,7 @@ def test_pauliz_hadamard(self, device, tol, skip_if): """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly""" n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) theta = 0.432 @@ -517,7 +517,7 @@ def circ(obs): """ n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) @qml.qnode(dev) @@ -545,7 +545,7 @@ def circ(wire_labels): """ dev = device(wires=3) dev_custom_labels = device(wires=label_map) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) def circ(wire_labels): @@ -567,7 +567,7 @@ def test_hermitian(self, device, tol, skip_if): n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if "Hermitian" not in dev.observables: pytest.skip("Skipped because device does not support the Hermitian observable.") @@ -609,7 +609,7 @@ def test_projector(self, device, tol, skip_if): n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if "Projector" not in dev.observables: pytest.skip("Skipped because device does not support the Projector observable.") @@ -661,7 +661,7 @@ def test_sparse_hamiltonian_expval(self, device, tol): n_wires = 4 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if "SparseHamiltonian" not in dev.observables: pytest.skip( "Skipped because device does not support the SparseHamiltonian observable." @@ -724,7 +724,7 @@ def test_sample_values_hermitian(self, device, tol): if not dev.shots: pytest.skip("Device is in analytic mode, cannot test sampling.") - if isinstance(dev, qml.Device) and "Hermitian" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Hermitian" not in dev.observables: pytest.skip("Skipped because device does not support the Hermitian observable.") A_ = np.array([[1, 2j], [-2j, 0]]) @@ -760,7 +760,7 @@ def test_sample_values_projector(self, device, tol): if not dev.shots: pytest.skip("Device is in analytic mode, cannot test sampling.") - if isinstance(dev, qml.Device) and "Projector" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Projector" not in dev.observables: pytest.skip("Skipped because device does not support the Projector observable.") theta = 0.543 @@ -808,7 +808,7 @@ def test_sample_values_hermitian_multi_qubit(self, device, tol): if not dev.shots: pytest.skip("Device is in analytic mode, cannot test sampling.") - if isinstance(dev, qml.Device) and "Hermitian" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Hermitian" not in dev.observables: pytest.skip("Skipped because device does not support the Hermitian observable.") theta = 0.543 @@ -857,7 +857,7 @@ def test_sample_values_projector_multi_qubit(self, device, tol): if not dev.shots: pytest.skip("Device is in analytic mode, cannot test sampling.") - if isinstance(dev, qml.Device) and "Projector" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Projector" not in dev.observables: pytest.skip("Skipped because device does not support the Projector observable.") theta = 0.543 @@ -915,7 +915,7 @@ def test_paulix_pauliy(self, device, tol, skip_if): if not dev.shots: pytest.skip("Device is in analytic mode, cannot test sampling.") - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) theta = 0.432 @@ -959,7 +959,7 @@ def test_pauliz_hadamard(self, device, tol, skip_if): if not dev.shots: pytest.skip("Device is in analytic mode, cannot test sampling.") - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) theta = 0.432 @@ -1001,7 +1001,7 @@ def test_hermitian(self, device, tol, skip_if): if not dev.shots: pytest.skip("Device is in analytic mode, cannot test sampling.") - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if "Hermitian" not in dev.observables: pytest.skip("Skipped because device does not support the Hermitian observable.") @@ -1095,7 +1095,7 @@ def test_projector(self, device, tol, skip_if): # pylint: disable=too-many-stat if not dev.shots: pytest.skip("Device is in analytic mode, cannot test sampling.") - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if "Projector" not in dev.observables: pytest.skip("Skipped because device does not support the Projector observable.") @@ -1270,7 +1270,7 @@ def test_var_hermitian(self, device, tol): n_wires = 2 dev = device(n_wires) - if isinstance(dev, qml.Device) and "Hermitian" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Hermitian" not in dev.observables: pytest.skip("Skipped because device does not support the Hermitian observable.") phi = 0.543 @@ -1306,7 +1306,7 @@ def test_var_projector(self, device, tol): n_wires = 2 dev = device(n_wires) - if isinstance(dev, qml.Device) and "Projector" not in dev.observables: + if isinstance(dev, qml.devices.LegacyDevice) and "Projector" not in dev.observables: pytest.skip("Skipped because device does not support the Projector observable.") phi = 0.543 @@ -1367,7 +1367,7 @@ def test_paulix_pauliy(self, device, tol, skip_if): """Test that a tensor product involving PauliX and PauliY works correctly""" n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) theta = 0.432 @@ -1399,7 +1399,7 @@ def test_pauliz_hadamard(self, device, tol, skip_if): """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly""" n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) theta = 0.432 @@ -1449,7 +1449,7 @@ def circ(obs): """ n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) @qml.qnode(dev) @@ -1476,7 +1476,7 @@ def circ(wire_labels): """ dev = device(wires=3) dev_custom_labels = device(wires=label_map) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): skip_if(dev, {"supports_tensor_observables": False}) def circ(wire_labels): @@ -1498,7 +1498,7 @@ def test_hermitian(self, device, tol, skip_if): n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if "Hermitian" not in dev.observables: pytest.skip("Skipped because device does not support the Hermitian observable.") @@ -1570,7 +1570,7 @@ def test_projector(self, device, tol, skip_if): n_wires = 3 dev = device(n_wires) - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if "Projector" not in dev.observables: pytest.skip("Skipped because device does not support the Projector observable.") diff --git a/pennylane/devices/tests/test_templates.py b/pennylane/devices/tests/test_templates.py index 07b0b56b39c..86e1101340a 100644 --- a/pennylane/devices/tests/test_templates.py +++ b/pennylane/devices/tests/test_templates.py @@ -33,7 +33,7 @@ def check_op_supported(op, dev): """Skip test if device does not support an operation. Works with both device APIs""" - if isinstance(dev, qml.Device): + if isinstance(dev, qml.devices.LegacyDevice): if op.name not in dev.operations: pytest.skip("operation not supported.") else: diff --git a/pennylane/devices/tests/test_tracker.py b/pennylane/devices/tests/test_tracker.py index 75fc6145ac8..c46a5ad952e 100644 --- a/pennylane/devices/tests/test_tracker.py +++ b/pennylane/devices/tests/test_tracker.py @@ -26,7 +26,9 @@ def test_tracker_initialization(self, device): dev = device(1) - if isinstance(dev, qml.Device) and not dev.capabilities().get("supports_tracker", False): + if isinstance(dev, qml.devices.LegacyDevice) and not dev.capabilities().get( + "supports_tracker", False + ): pytest.skip("Device does not support a tracker") assert isinstance(dev.tracker, qml.Tracker) @@ -36,7 +38,9 @@ def test_tracker_updated_in_execution_mode(self, device): dev = device(1) - if isinstance(dev, qml.Device) and not dev.capabilities().get("supports_tracker", False): + if isinstance(dev, qml.devices.LegacyDevice) and not dev.capabilities().get( + "supports_tracker", False + ): pytest.skip("Device does not support a tracker") @qml.qnode(dev, diff_method="parameter-shift") diff --git a/pennylane/fermi/fermionic.py b/pennylane/fermi/fermionic.py index 939827232f5..193cc0b4be8 100644 --- a/pennylane/fermi/fermionic.py +++ b/pennylane/fermi/fermionic.py @@ -55,6 +55,22 @@ def __init__(self, operator): super().__init__(operator) + def adjoint(self): + r"""Return the adjoint of FermiWord.""" + n = len(self.items()) + adjoint_dict = {} + for key, value in reversed(self.items()): + position = n - key[0] - 1 + orbital = key[1] + fermi = "+" if value == "-" else "-" + adjoint_dict[(position, orbital)] = fermi + + return FermiWord(adjoint_dict) + + def items(self): + """Returns the dictionary items in sorted order.""" + return self.sorted_dic.items() + @property def wires(self): r"""Return wires in a FermiWord.""" @@ -331,6 +347,16 @@ class FermiSentence(dict): def __init__(self, operator): super().__init__(operator) + def adjoint(self): + r"""Return the adjoint of FermiSentence.""" + adjoint_dict = {} + for key, value in self.items(): + word = key.adjoint() + scalar = qml.math.conj(value) + adjoint_dict[word] = scalar + + return FermiSentence(adjoint_dict) + @property def wires(self): r"""Return wires of the FermiSentence.""" @@ -657,9 +683,14 @@ def __init__(self, orbital): raise ValueError( f"FermiC: expected a single, positive integer value for orbital, but received {orbital}" ) + self.orbital = orbital operator = {(0, orbital): "+"} super().__init__(operator) + def adjoint(self): + """Return the adjoint of FermiC.""" + return FermiA(self.orbital) + class FermiA(FermiWord): r"""FermiA(orbital) @@ -694,5 +725,10 @@ def __init__(self, orbital): raise ValueError( f"FermiA: expected a single, positive integer value for orbital, but received {orbital}" ) + self.orbital = orbital operator = {(0, orbital): "-"} super().__init__(operator) + + def adjoint(self): + """Return the adjoint of FermiA.""" + return FermiC(self.orbital) diff --git a/pennylane/gradients/jvp.py b/pennylane/gradients/jvp.py index e3634cce492..67428ab5c68 100644 --- a/pennylane/gradients/jvp.py +++ b/pennylane/gradients/jvp.py @@ -295,11 +295,18 @@ def jvp(tape, tangent, gradient_fn, gradient_kwargs=None): if len(tape.trainable_params) == 0: # The tape has no trainable parameters; the JVP # is simply none. - def zero_vjp(_): - res = tuple(np.zeros(mp.shape(None, tape.shots)) for mp in tape.measurements) + def zero_jvp_for_single_shots(s): + res = tuple( + np.zeros(mp.shape(shots=s), dtype=mp.numeric_type) for mp in tape.measurements + ) return res[0] if len(tape.measurements) == 1 else res - return tuple(), zero_vjp + def zero_jvp(_): + if tape.shots.has_partitioned_shots: + return tuple(zero_jvp_for_single_shots(s) for s in tape.shots) + return zero_jvp_for_single_shots(tape.shots.total_shots) + + return tuple(), zero_jvp multi_m = len(tape.measurements) > 1 diff --git a/pennylane/measurements/classical_shadow.py b/pennylane/measurements/classical_shadow.py index 6c502858cf8..200a48a25a5 100644 --- a/pennylane/measurements/classical_shadow.py +++ b/pennylane/measurements/classical_shadow.py @@ -284,7 +284,11 @@ def process(self, tape, device): n_snapshots = device.shots seed = self.seed - with qml.workflow.set_shots(device, shots=1): + original_shots = device.shots + original_shot_vector = device._shot_vector # pylint: disable=protected-access + + try: + device.shots = 1 # slow implementation but works for all devices n_qubits = len(wires) mapped_wires = np.array(device.map_wires(wires)) @@ -311,6 +315,9 @@ def process(self, tape, device): device.apply(tape.operations, rotations=tape.diagonalizing_gates + rotations) outcomes[t] = device.generate_samples()[0][mapped_wires] + finally: + device.shots = original_shots + device._shot_vector = original_shot_vector # pylint: disable=protected-access return qml.math.cast(qml.math.stack([outcomes, recipes]), dtype=np.int8) diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py index 1a95f9e7ce0..3915ef481f0 100644 --- a/pennylane/measurements/expval.py +++ b/pennylane/measurements/expval.py @@ -151,5 +151,5 @@ def _calculate_expectation(self, probabilities): Args: probabilities (array): the probabilities of collapsing to eigen states """ - eigvals = qml.math.asarray(self.eigvals(), dtype="float64") + eigvals = qml.math.cast_like(self.eigvals(), 1.0) return qml.math.dot(probabilities, eigvals) diff --git a/pennylane/operation.py b/pennylane/operation.py index 0e7e8528941..77f989e6383 100644 --- a/pennylane/operation.py +++ b/pennylane/operation.py @@ -1517,33 +1517,6 @@ def adjoint(self) -> "Operator": # pylint:disable=no-self-use """ raise AdjointUndefinedError - def expand(self) -> "qml.tape.QuantumScript": - """Returns a tape that contains the decomposition of the operator. - - .. warning:: - This function is deprecated and will be removed in version 0.39. - The same behaviour can be achieved simply through 'qml.tape.QuantumScript(self.decomposition())'. - - Returns: - .QuantumTape: quantum tape - """ - warnings.warn( - "'Operator.expand' is deprecated and will be removed in version 0.39. " - "The same behaviour can be achieved simply through 'qml.tape.QuantumScript(self.decomposition())'.", - qml.PennyLaneDeprecationWarning, - ) - - if not self.has_decomposition: - raise DecompositionUndefinedError - - qscript = qml.tape.QuantumScript(self.decomposition()) - - if not self.data: - # original operation has no trainable parameters - qscript.trainable_params = {} - - return qscript - @property def arithmetic_depth(self) -> int: """Arithmetic depth of the operator.""" diff --git a/pennylane/ops/op_math/controlled_decompositions.py b/pennylane/ops/op_math/controlled_decompositions.py index 26de3761cf6..93a6e82b02f 100644 --- a/pennylane/ops/op_math/controlled_decompositions.py +++ b/pennylane/ops/op_math/controlled_decompositions.py @@ -55,7 +55,7 @@ def _convert_to_su2(U, return_global_phase=False): with np.errstate(divide="ignore", invalid="ignore"): dets = math.linalg.det(U) - global_phase = math.cast_like(math.angle(dets), 1j) / 2 + global_phase = math.cast_like(math.angle(dets), 1.0) / 2 U_SU2 = math.cast_like(U, dets) * math.exp(-1j * global_phase) return (U_SU2, global_phase) if return_global_phase else U_SU2 diff --git a/pennylane/ops/qubit/special_unitary.py b/pennylane/ops/qubit/special_unitary.py index 0cff0418f31..14807e1468f 100644 --- a/pennylane/ops/qubit/special_unitary.py +++ b/pennylane/ops/qubit/special_unitary.py @@ -480,8 +480,8 @@ def compute_matrix(theta: TensorLike, num_wires: int) -> TensorLike: [-0.0942679 +0.47133952j, 0.83004499+0.28280371j]]) """ interface = qml.math.get_interface(theta) - if interface == "tensorflow": - theta = qml.math.cast_like(theta, 1j) + theta = qml.math.cast_like(theta, 1j) + if num_wires > 5: matrices = product(_pauli_matrices, repeat=num_wires) # Drop the identity from the generator of matrices diff --git a/pennylane/qcut/cutstrategy.py b/pennylane/qcut/cutstrategy.py index a26a528b338..2e92a8cdb87 100644 --- a/pennylane/qcut/cutstrategy.py +++ b/pennylane/qcut/cutstrategy.py @@ -41,7 +41,7 @@ class CutStrategy: check out the :func:`qml.cut_circuit() ` transform for more details. Args: - devices (Union[qml.Device, Sequence[qml.Device]]): Single, or Sequence of, device(s). + devices (Union[qml.devices.Device, Sequence[qml.devices.Device]]): Single, or Sequence of, device(s). Optional only when ``max_free_wires`` is provided. max_free_wires (int): Number of wires for the largest available device. Optional only when ``devices`` is provided where it defaults to the maximum number of wires among diff --git a/pennylane/resource/specs.py b/pennylane/resource/specs.py index 2123f8653e1..64233ab24d9 100644 --- a/pennylane/resource/specs.py +++ b/pennylane/resource/specs.py @@ -208,18 +208,24 @@ def specs_qnode(*args, **kwargs) -> Union[list[dict], dict]: else qnode.diff_method ) - if isinstance(qnode.gradient_fn, qml.transforms.core.TransformDispatcher): - info["gradient_fn"] = _get_absolute_import_path(qnode.gradient_fn) + gradient_fn = qml.QNode.get_gradient_fn( + qnode.device, + qnode.interface, + qnode.diff_method, + tape=tape, + )[0] + if isinstance(gradient_fn, qml.transforms.core.TransformDispatcher): + info["gradient_fn"] = _get_absolute_import_path(gradient_fn) try: - info["num_gradient_executions"] = len(qnode.gradient_fn(tape)[0]) + info["num_gradient_executions"] = len(gradient_fn(tape)[0]) except Exception as e: # pylint: disable=broad-except # In the case of a broad exception, we don't want the `qml.specs` transform # to fail. Instead, we simply indicate that the number of gradient executions # is not supported for the reason specified. info["num_gradient_executions"] = f"NotSupported: {str(e)}" else: - info["gradient_fn"] = qnode.gradient_fn + info["gradient_fn"] = gradient_fn infos.append(info) diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py index 9758e8fc185..7cc6809ff82 100644 --- a/pennylane/tape/qscript.py +++ b/pennylane/tape/qscript.py @@ -555,7 +555,7 @@ def trainable_params(self) -> list[int]: .. note:: For devices that support native backpropagation (such as - ``default.qubit.tf`` and ``default.qubit.autograd``), this + ``default.qubit`` and ``default.mixed``), this property contains no relevant information when using backpropagation to compute gradients. diff --git a/pennylane/templates/subroutines/fable.py b/pennylane/templates/subroutines/fable.py index 16f1160ddd0..d9637676738 100644 --- a/pennylane/templates/subroutines/fable.py +++ b/pennylane/templates/subroutines/fable.py @@ -166,17 +166,19 @@ def compute_decomposition(input_matrix, wires, tol=0): # pylint:disable=argumen for c_wire in nots: op_list.append(qml.CNOT(wires=[c_wire] + ancilla)) op_list.append(qml.RY(2 * theta, wires=ancilla)) + nots = {} nots[wire_map[control_index]] = 1 + continue + + if qml.math.abs(2 * theta) > tol: + for c_wire in nots: + op_list.append(qml.CNOT(wires=[c_wire] + ancilla)) + op_list.append(qml.RY(2 * theta, wires=ancilla)) + nots = {} + if wire_map[control_index] in nots: + del nots[wire_map[control_index]] else: - if abs(2 * theta) > tol: - for c_wire in nots: - op_list.append(qml.CNOT(wires=[c_wire] + ancilla)) - op_list.append(qml.RY(2 * theta, wires=ancilla)) - nots = {} - if wire_map[control_index] in nots: - del nots[wire_map[control_index]] - else: - nots[wire_map[control_index]] = 1 + nots[wire_map[control_index]] = 1 for c_wire in nots: op_list.append(qml.CNOT([c_wire] + ancilla)) diff --git a/pennylane/workflow/construct_batch.py b/pennylane/workflow/construct_batch.py index bfd788d7b00..1701e459f37 100644 --- a/pennylane/workflow/construct_batch.py +++ b/pennylane/workflow/construct_batch.py @@ -18,7 +18,7 @@ from collections.abc import Callable from contextlib import nullcontext from functools import wraps -from typing import Literal, Optional, Union +from typing import Literal, Union import pennylane as qml from pennylane.tape import QuantumScriptBatch @@ -56,20 +56,24 @@ def wrapped_expand_fn(tape, *args, **kwargs): return qml.transform(wrapped_expand_fn) -def _get_full_transform_program(qnode: QNode) -> "qml.transforms.core.TransformProgram": +def _get_full_transform_program( + qnode: QNode, gradient_fn +) -> "qml.transforms.core.TransformProgram": program = qml.transforms.core.TransformProgram(qnode.transform_program) - if getattr(qnode.gradient_fn, "expand_transform", False): + if getattr(gradient_fn, "expand_transform", False): program.add_transform( - qml.transform(qnode.gradient_fn.expand_transform), + qml.transform(gradient_fn.expand_transform), **qnode.gradient_kwargs, ) - config = _make_execution_config(qnode, qnode.gradient_fn) + config = _make_execution_config(qnode, gradient_fn) return program + qnode.device.preprocess(config)[0] -def get_transform_program(qnode: "QNode", level=None) -> "qml.transforms.core.TransformProgram": +def get_transform_program( + qnode: "QNode", level=None, gradient_fn="unset" +) -> "qml.transforms.core.TransformProgram": """Extract a transform program at a designated level. Args: @@ -81,6 +85,8 @@ def get_transform_program(qnode: "QNode", level=None) -> "qml.transforms.core.Tr * ``int``: How many transforms to include, starting from the front of the program * ``slice``: a slice to select out components of the transform program. + gradient_fn (None, str, TransformDispatcher): The processed gradient fn for the workflow. + Returns: TransformProgram: the transform program corresponding to the requested level. @@ -174,7 +180,10 @@ def circuit(): TransformProgram(validate_device_wires, mid_circuit_measurements, decompose, validate_measurements, validate_observables) """ - full_transform_program = _get_full_transform_program(qnode) + if gradient_fn == "unset": + gradient_fn = QNode.get_gradient_fn(qnode.device, qnode.interface, qnode.diff_method)[0] + + full_transform_program = _get_full_transform_program(qnode, gradient_fn) num_user = len(qnode.transform_program) if qnode.transform_program.has_final_transform: @@ -193,7 +202,7 @@ def circuit(): elif level == "gradient": readd_final_transform = True - level = num_user + 1 if getattr(qnode.gradient_fn, "expand_transform", False) else num_user + level = num_user + 1 if getattr(gradient_fn, "expand_transform", False) else num_user elif isinstance(level, str): raise ValueError( f"level {level} not recognized. Acceptable strings are 'device', 'top', 'user', and 'gradient'." @@ -211,8 +220,8 @@ def circuit(): def construct_batch( - qnode: QNode, - level: Optional[Union[Literal["top", "user", "device", "gradient"], int, slice]] = "user", + qnode: Union[QNode, "qml.qnn.KerasLayer", "qml.qnn.TorchLayer"], + level: Union[Literal["top", "user", "device", "gradient"], int, slice, None] = "user", ) -> Callable: """Construct the batch of tapes and post processing for a designated stage in the transform program. @@ -350,8 +359,10 @@ def batch_constructor(*args, **kwargs) -> tuple[QuantumScriptBatch, Postprocessi params = initial_tape.get_parameters(trainable_only=False) initial_tape.trainable_params = qml.math.get_trainable_indices(params) - qnode._update_gradient_fn(tape=initial_tape) - program = get_transform_program(qnode, level=level) + gradient_fn = QNode.get_gradient_fn( + qnode.device, qnode.interface, qnode.diff_method, tape=initial_tape + )[0] + program = get_transform_program(qnode, level=level, gradient_fn=gradient_fn) return program((initial_tape,)) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index 2465c70e481..7445bcea2b7 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -165,7 +165,7 @@ def _get_ml_boundary_execute( else: from .interfaces.jax import jax_jvp_execute as ml_boundary - except ImportError as e: # pragma: no-cover + except ImportError as e: # pragma: no cover raise qml.QuantumFunctionError( f"{mapped_interface} not found. Please install the latest " f"version of {mapped_interface} to enable the '{mapped_interface}' interface." diff --git a/pennylane/workflow/interfaces/jax_jit.py b/pennylane/workflow/interfaces/jax_jit.py index d52f686148e..cdfade8a979 100644 --- a/pennylane/workflow/interfaces/jax_jit.py +++ b/pennylane/workflow/interfaces/jax_jit.py @@ -96,7 +96,7 @@ def _get_counts_shape(mp: "qml.measurements.CountsMP", num_device_wires=0): return outcome_counts -def _result_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"): +def _result_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.devices.Device"): """Auxiliary function for creating the shape and dtype object structure given a tape.""" @@ -125,7 +125,7 @@ def struct(mp, shots): return tuple(shape) if tape.shots.has_partitioned_shots else shape[0] -def _jac_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.Device"): +def _jac_shape_dtype_struct(tape: "qml.tape.QuantumScript", device: "qml.devices.Device"): """The shape of a jacobian for a single tape given a device. Args: @@ -167,14 +167,9 @@ def pure_callback_wrapper(p): new_tapes = _set_fn(tapes.vals, p) return _to_jax(execute_fn(new_tapes)) - if isinstance(device, qml.Device): - device_supports_vectorization = device.capabilities().get("supports_broadcasting") - else: - # first order way of determining native parameter broadcasting support - # will be inaccurate when inclusion of broadcast_expand depends on ExecutionConfig values (like adjoint) - device_supports_vectorization = ( - qml.transforms.broadcast_expand not in device.preprocess()[0] - ) + # first order way of determining native parameter broadcasting support + # will be inaccurate when inclusion of broadcast_expand depends on ExecutionConfig values (like adjoint) + device_supports_vectorization = qml.transforms.broadcast_expand not in device.preprocess()[0] out = jax.pure_callback( pure_callback_wrapper, shape_dtype_structs, params, vectorized=device_supports_vectorization ) diff --git a/pennylane/workflow/jacobian_products.py b/pennylane/workflow/jacobian_products.py index f8163813366..d3ae50dca65 100644 --- a/pennylane/workflow/jacobian_products.py +++ b/pennylane/workflow/jacobian_products.py @@ -18,7 +18,7 @@ import inspect import logging from collections.abc import Callable, Sequence -from typing import Optional, Union +from typing import Optional import numpy as np from cachetools import LRUCache @@ -46,6 +46,17 @@ def _compute_vjps(jacs, dys, tapes): return tuple(vjps) +def _zero_jvp_single_shots(shots, tape): + jvp = tuple(np.zeros(mp.shape(shots=shots), dtype=mp.numeric_type) for mp in tape.measurements) + return jvp[0] if len(tape.measurements) == 1 else jvp + + +def _zero_jvp(tape): + if tape.shots.has_partitioned_shots: + return tuple(_zero_jvp_single_shots(s, tape) for s in tape.shots) + return _zero_jvp_single_shots(tape.shots.total_shots, tape) + + def _compute_jvps(jacs, tangents, tapes): """Compute the jvps of multiple tapes, directly for a Jacobian and tangents.""" f = {True: qml.gradients.compute_jvp_multi, False: qml.gradients.compute_jvp_single} @@ -54,16 +65,7 @@ def _compute_jvps(jacs, tangents, tapes): for jac, dx, t in zip(jacs, tangents, tapes): multi = len(t.measurements) > 1 if len(t.trainable_params) == 0: - empty_shots = qml.measurements.Shots(None) - zeros_jvp = tuple( - np.zeros(mp.shape(None, empty_shots), dtype=mp.numeric_type) - for mp in t.measurements - ) - zeros_jvp = zeros_jvp[0] if len(t.measurements) == 1 else zeros_jvp - if t.shots.has_partitioned_shots: - jvps.append(tuple(zeros_jvp for _ in range(t.shots.num_copies))) - else: - jvps.append(zeros_jvp) + jvps.append(_zero_jvp(t)) elif t.shots.has_partitioned_shots: jvps.append(tuple(f[multi](dx, j) for j in jac)) else: @@ -332,7 +334,7 @@ def compute_jacobian(self, tapes: QuantumScriptBatch): class DeviceDerivatives(JacobianProductCalculator): - """Calculate jacobian products via a device provided jacobian. This class relies on either ``qml.Device.gradients`` or + """Calculate jacobian products via a device provided jacobian. This class relies on ``qml.devices.Device.compute_derivatives``. Args: @@ -397,7 +399,7 @@ def __repr__(self): def __init__( self, - device: Union["qml.devices.Device", "qml.Device"], + device: "qml.devices.Device", execution_config: Optional["qml.devices.ExecutionConfig"] = None, ): if execution_config is None: diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py index f8f724e58cf..408a0794674 100644 --- a/pennylane/workflow/qnode.py +++ b/pennylane/workflow/qnode.py @@ -524,7 +524,7 @@ def __init__( # input arguments self.func = func self.device = device - self._interface = interface + self._interface = None if diff_method is None else interface self.diff_method = diff_method mcm_config = qml.devices.MCMConfig(mcm_method=mcm_method, postselect_mode=postselect_mode) cache = (max_diff > 1) if cache == "auto" else cache @@ -542,14 +542,48 @@ def __init__( # internal data attributes self._tape = None self._qfunc_output = None - self._user_gradient_kwargs = gradient_kwargs - self.gradient_fn = None - self.gradient_kwargs = {} + self._gradient_fn = None + self.gradient_kwargs = gradient_kwargs self._transform_program = TransformProgram() - self._update_gradient_fn() functools.update_wrapper(self, func) + # validation check. Will raise error if bad diff_method + if diff_method is not None: + QNode.get_gradient_fn(self.device, self.interface, self.diff_method) + + @property + def gradient_fn(self): + """A processed version of ``QNode.diff_method``. + + .. warning:: + + This property is deprecated in v0.39 and will be removed in v0.40. + + Please see ``QNode.diff_method`` instead. + + """ + warnings.warn( + "QNode.gradient_fn is deprecated. Please use QNode.diff_method instead.", + qml.PennyLaneDeprecationWarning, + ) + if self.diff_method is None: + return None + + if ( + self.device.name == "lightning.qubit" + and qml.metric_tensor in self.transform_program + and self.diff_method == "best" + ): + return qml.gradients.param_shift + + if self.tape is None and self.device.shots: + tape = qml.tape.QuantumScript([], [], shots=self.device.shots) + else: + tape = self.tape + + return QNode.get_gradient_fn(self.device, self.interface, self.diff_method, tape=tape)[0] + def __copy__(self) -> "QNode": copied_qnode = QNode.__new__(QNode) for attr, value in vars(self).items(): @@ -590,7 +624,6 @@ def interface(self, value: SupportedInterfaceUserInput): ) self._interface = INTERFACE_MAP[value] - self._update_gradient_fn(shots=self.device.shots) @property def transform_program(self) -> TransformProgram: @@ -605,28 +638,6 @@ def add_transform(self, transform_container: TransformContainer): """ self._transform_program.push_back(transform_container=transform_container) - def _update_gradient_fn(self, shots=None, tape: Optional["qml.tape.QuantumTape"] = None): - if self.diff_method is None: - self._interface = None - self.gradient_fn = None - self.gradient_kwargs = {} - return - if tape is None and shots: - tape = qml.tape.QuantumScript([], [], shots=shots) - - diff_method = self.diff_method - if ( - self.device.name == "lightning.qubit" - and qml.metric_tensor in self.transform_program - and self.diff_method == "best" - ): - diff_method = "parameter-shift" - - self.gradient_fn, self.gradient_kwargs, self.device = QNode.get_gradient_fn( - self.device, self.interface, diff_method, tape=tape - ) - self.gradient_kwargs.update(self._user_gradient_kwargs or {}) - # pylint: disable=too-many-return-statements @staticmethod @debug_logger @@ -652,6 +663,8 @@ def get_gradient_fn( tuple[str or .TransformDispatcher, dict, .device.Device: Tuple containing the ``gradient_fn``, ``gradient_kwargs``, and the device to use when calling the execute function. """ + if diff_method is None: + return None, {}, device config = _make_execution_config(None, diff_method) @@ -859,8 +872,22 @@ def _execution_component(self, args: tuple, kwargs: dict) -> qml.typing.Result: Result """ - + if ( + self.device.name == "lightning.qubit" + and qml.metric_tensor in self.transform_program + and self.diff_method == "best" + ): + gradient_fn = qml.gradients.param_shift + else: + gradient_fn = QNode.get_gradient_fn( + self.device, self.interface, self.diff_method, tape=self.tape + )[0] execute_kwargs = copy.copy(self.execute_kwargs) + + gradient_kwargs = copy.copy(self.gradient_kwargs) + if gradient_fn is qml.gradients.param_shift_cv: + gradient_kwargs["dev"] = self.device + mcm_config = copy.copy(execute_kwargs["mcm_config"]) if not self._tape.shots: mcm_config.postselect_mode = None @@ -875,7 +902,7 @@ def _execution_component(self, args: tuple, kwargs: dict) -> qml.typing.Result: full_transform_program = qml.transforms.core.TransformProgram(self.transform_program) inner_transform_program = qml.transforms.core.TransformProgram() - config = _make_execution_config(self, self.gradient_fn, mcm_config) + config = _make_execution_config(self, gradient_fn, mcm_config) device_transform_program, config = self.device.preprocess(execution_config=config) if config.use_device_gradient: @@ -884,10 +911,10 @@ def _execution_component(self, args: tuple, kwargs: dict) -> qml.typing.Result: inner_transform_program += device_transform_program # Add the gradient expand to the program if necessary - if getattr(self.gradient_fn, "expand_transform", False): + if getattr(gradient_fn, "expand_transform", False): full_transform_program.insert_front_transform( - qml.transform(self.gradient_fn.expand_transform), - **self.gradient_kwargs, + qml.transform(gradient_fn.expand_transform), + **gradient_kwargs, ) # Calculate the classical jacobians if necessary @@ -900,12 +927,12 @@ def _execution_component(self, args: tuple, kwargs: dict) -> qml.typing.Result: res = qml.execute( (self._tape,), device=self.device, - gradient_fn=self.gradient_fn, + gradient_fn=gradient_fn, interface=self.interface, transform_program=full_transform_program, inner_transform=inner_transform_program, config=config, - gradient_kwargs=self.gradient_kwargs, + gradient_kwargs=gradient_kwargs, **execute_kwargs, ) res = res[0] @@ -924,6 +951,9 @@ def _execution_component(self, args: tuple, kwargs: dict) -> qml.typing.Result: def _impl_call(self, *args, **kwargs) -> qml.typing.Result: + # construct the tape + self.construct(args, kwargs) + old_interface = self.interface if old_interface == "auto": interface = ( @@ -933,27 +963,12 @@ def _impl_call(self, *args, **kwargs) -> qml.typing.Result: ) self._interface = INTERFACE_MAP[interface] - if self._qfunc_uses_shots_arg: - override_shots = False - else: - if "shots" not in kwargs: - kwargs["shots"] = self.device.shots - override_shots = kwargs["shots"] - - # construct the tape - self.construct(args, kwargs) - - original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device] - self._update_gradient_fn(shots=override_shots, tape=self._tape) - try: res = self._execution_component(args, kwargs) finally: if old_interface == "auto": self._interface = "auto" - _, self.gradient_kwargs, self.device = original_grad_fn - return res def __call__(self, *args, **kwargs) -> qml.typing.Result: diff --git a/pennylane/workflow/set_shots.py b/pennylane/workflow/set_shots.py index 1bc7dff7f33..aaaed1a1a44 100644 --- a/pennylane/workflow/set_shots.py +++ b/pennylane/workflow/set_shots.py @@ -17,6 +17,7 @@ """ # pylint: disable=protected-access import contextlib +import warnings import pennylane as qml from pennylane.measurements import Shots @@ -26,6 +27,20 @@ def set_shots(device, shots): r"""Context manager to temporarily change the shots of a device. + + .. warning:: + + ``set_shots`` is deprecated and will be removed in PennyLane version v0.40. + + To dynamically update the shots on the workflow, shots can be manually set on a ``QNode`` call: + + >>> circuit(shots=my_new_shots) + + When working with the internal tapes, shots should be set on each tape. + + >>> tape = qml.tape.QuantumScript([], [qml.sample()], shots=50) + + This context manager can be used in two ways. As a standard context manager: @@ -47,6 +62,12 @@ def set_shots(device, shots): "The new device interface is not compatible with `set_shots`. " "Set shots when calling the qnode or put the shots on the QuantumTape." ) + warnings.warn( + "set_shots is deprecated.\n" + "Please dyanmically update shots via keyword argument when calling a QNode " + " or set shots on the tape.", + qml.PennyLaneDeprecationWarning, + ) if isinstance(shots, Shots): shots = shots.shot_vector if shots.has_partitioned_shots else shots.total_shots if shots == device.shots: diff --git a/setup.py b/setup.py index f1f77907b6a..e13673fb1fa 100644 --- a/setup.py +++ b/setup.py @@ -51,8 +51,6 @@ "default.qubit = pennylane.devices:DefaultQubit", "default.qubit.legacy = pennylane.devices:DefaultQubitLegacy", "default.gaussian = pennylane.devices:DefaultGaussian", - "default.qubit.autograd = pennylane.devices.default_qubit_autograd:DefaultQubitAutograd", - "default.qubit.jax = pennylane.devices.default_qubit_jax:DefaultQubitJax", "default.mixed = pennylane.devices.default_mixed:DefaultMixed", "null.qubit = pennylane.devices.null_qubit:NullQubit", "default.qutrit = pennylane.devices.default_qutrit:DefaultQutrit", diff --git a/tests/conftest.py b/tests/conftest.py index 3094a2fd784..d478b6f0998 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -36,20 +36,6 @@ TOL_STOCHASTIC = 0.05 -# pylint: disable=too-few-public-methods -class DummyDevice(DefaultGaussian): - """Dummy device to allow Kerr operations""" - - _operation_map = DefaultGaussian._operation_map.copy() - _operation_map["Kerr"] = lambda *x, **y: np.identity(2) - - -@pytest.fixture(autouse=True) -def set_numpy_seed(): - np.random.seed(9872653) - yield - - @pytest.fixture(scope="function", autouse=True) def capture_legacy_device_deprecation_warnings(): with warnings.catch_warnings(record=True) as recwarn: @@ -67,6 +53,20 @@ def capture_legacy_device_deprecation_warnings(): warnings.warn(message=w.message, category=w.category) +# pylint: disable=too-few-public-methods +class DummyDevice(DefaultGaussian): + """Dummy device to allow Kerr operations""" + + _operation_map = DefaultGaussian._operation_map.copy() + _operation_map["Kerr"] = lambda *x, **y: np.identity(2) + + +@pytest.fixture(autouse=True) +def set_numpy_seed(): + np.random.seed(9872653) + yield + + @pytest.fixture(scope="session") def tol(): """Numerical tolerance for equality tests.""" @@ -181,12 +181,12 @@ def mock_device(monkeypatch): """A mock instance of the abstract Device class""" with monkeypatch.context() as m: - dev = qml.Device + dev = qml.devices.LegacyDevice m.setattr(dev, "__abstractmethods__", frozenset()) m.setattr(dev, "short_name", "mock_device") m.setattr(dev, "capabilities", lambda cls: {"model": "qubit"}) m.setattr(dev, "operations", {"RX", "RY", "RZ", "CNOT", "SWAP"}) - yield qml.Device(wires=2) # pylint:disable=abstract-class-instantiated + yield qml.devices.LegacyDevice(wires=2) # pylint:disable=abstract-class-instantiated # pylint: disable=protected-access diff --git a/tests/devices/default_qubit/test_default_qubit_preprocessing.py b/tests/devices/default_qubit/test_default_qubit_preprocessing.py index 29a1fdca73a..74c30f8661d 100644 --- a/tests/devices/default_qubit/test_default_qubit_preprocessing.py +++ b/tests/devices/default_qubit/test_default_qubit_preprocessing.py @@ -139,7 +139,6 @@ def circuit(x): with dev.tracker: qml.grad(circuit)(qml.numpy.array(0.1)) - assert circuit.gradient_fn == "adjoint" assert dev.tracker.totals["execute_and_derivative_batches"] == 1 diff --git a/tests/devices/test_default_mixed_autograd.py b/tests/devices/test_default_mixed_autograd.py index 6d690054033..c3f98d5955f 100644 --- a/tests/devices/test_default_mixed_autograd.py +++ b/tests/devices/test_default_mixed_autograd.py @@ -85,7 +85,6 @@ def circuit(x): expected = -np.sin(p) - assert circuit.gradient_fn == "backprop" assert np.isclose(circuit(p), expected, atol=tol, rtol=0) def test_correct_state(self, tol): @@ -267,7 +266,6 @@ def circuit(p): qml.RX(p[2] / 2, wires=0) return qml.expval(qml.PauliZ(0)) - assert circuit.gradient_fn == "backprop" res = circuit(weights) expected = np.cos(3 * x) * np.cos(y) * np.cos(z / 2) - np.sin(3 * x) * np.sin(z / 2) @@ -340,9 +338,6 @@ def cost(x): assert np.allclose(res, circuit2(p), atol=tol, rtol=0) - assert circuit1.gradient_fn == "backprop" - assert circuit2.gradient_fn is qml.gradients.param_shift - grad_fn = qml.jacobian(circuit1, 0) res = grad_fn(p) assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0) @@ -549,14 +544,6 @@ def cost(params): expected_cost = (np.sin(lam) * np.sin(phi) - np.cos(theta) * np.cos(lam) * np.cos(phi)) ** 2 assert np.allclose(res, expected_cost, atol=tol, rtol=0) - # Check that the correct differentiation method is being used. - if diff_method == "backprop": - assert circuit.gradient_fn == "backprop" - elif diff_method == "parameter-shift": - assert circuit.gradient_fn is qml.gradients.param_shift - else: - assert circuit.gradient_fn is qml.gradients.finite_diff - res = qml.grad(cost)(params) expected_grad = ( np.array( diff --git a/tests/devices/test_default_mixed_jax.py b/tests/devices/test_default_mixed_jax.py index 3b0f5a2470d..c8ae1ce8ac1 100644 --- a/tests/devices/test_default_mixed_jax.py +++ b/tests/devices/test_default_mixed_jax.py @@ -59,7 +59,6 @@ def circuit(x): expected = -np.sin(p) - assert circuit.gradient_fn == "backprop" assert np.isclose(circuit(p), expected, atol=tol, rtol=0) def test_correct_state(self, tol): @@ -405,7 +404,6 @@ def circuit(p): qml.RX(p[2] / 2, wires=0) return qml.expval(qml.PauliZ(0)) - assert circuit.gradient_fn == "backprop" res = decorator(circuit)(weights) expected = np.cos(3 * x) * np.cos(y) * np.cos(z / 2) - np.sin(3 * x) * np.sin(z / 2) @@ -477,9 +475,6 @@ def circuit(x): res = decorator(circuit1)(p_jax) assert np.allclose(res, circuit2(p), atol=tol, rtol=0) - assert circuit1.gradient_fn == "backprop" - assert circuit2.gradient_fn is qml.gradients.param_shift - res = decorator(jacobian_fn(circuit1, 0))(p_jax) assert np.allclose(res, jax.jacobian(circuit2)(p), atol=tol, rtol=0) @@ -710,14 +705,6 @@ def cost(params): expected_cost = (np.sin(lam) * np.sin(phi) - np.cos(theta) * np.cos(lam) * np.cos(phi)) ** 2 assert np.allclose(res, expected_cost, atol=tol, rtol=0) - # Check that the correct differentiation method is being used. - if diff_method == "backprop": - assert circuit.gradient_fn == "backprop" - elif diff_method == "parameter-shift": - assert circuit.gradient_fn is qml.gradients.param_shift - else: - assert circuit.gradient_fn is qml.gradients.finite_diff - res = jax.grad(cost)(params) expected_grad = ( diff --git a/tests/devices/test_default_mixed_tf.py b/tests/devices/test_default_mixed_tf.py index 6fc555a1283..256f70a4668 100644 --- a/tests/devices/test_default_mixed_tf.py +++ b/tests/devices/test_default_mixed_tf.py @@ -60,7 +60,6 @@ def circuit(x): expected = -np.sin(p) - assert circuit.gradient_fn == "backprop" assert np.isclose(circuit(p), expected, atol=tol, rtol=0) def test_correct_state(self, tol): @@ -300,7 +299,6 @@ def circuit(p): qml.RX(p[2] / 2, wires=0) return qml.expval(qml.PauliZ(0)) - assert circuit.gradient_fn == "backprop" res = circuit(weights) expected = np.cos(3 * x) * np.cos(y) * np.cos(z / 2) - np.sin(3 * x) * np.sin(z / 2) @@ -380,9 +378,6 @@ def cost(x): assert np.allclose(res, circuit2(p), atol=tol, rtol=0) - assert circuit1.gradient_fn == "backprop" - assert circuit2.gradient_fn is qml.gradients.param_shift - res = tape.jacobian(res, p_tf) assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0) @@ -631,14 +626,6 @@ def cost(params): expected_cost = (np.sin(lam) * np.sin(phi) - np.cos(theta) * np.cos(lam) * np.cos(phi)) ** 2 assert np.allclose(res, expected_cost, atol=tol, rtol=0) - # Check that the correct differentiation method is being used. - if diff_method == "backprop": - assert circuit.gradient_fn == "backprop" - elif diff_method == "parameter-shift": - assert circuit.gradient_fn is qml.gradients.param_shift - else: - assert circuit.gradient_fn is qml.gradients.finite_diff - with tf.GradientTape() as tape: out = cost(params) diff --git a/tests/devices/test_default_mixed_torch.py b/tests/devices/test_default_mixed_torch.py index d85242e15fd..b5de0d0867a 100644 --- a/tests/devices/test_default_mixed_torch.py +++ b/tests/devices/test_default_mixed_torch.py @@ -54,7 +54,6 @@ def circuit(x): expected = -np.sin(p) - assert circuit.gradient_fn == "backprop" assert np.isclose(circuit(p), expected, atol=tol, rtol=0) def test_correct_state(self, tol): @@ -308,7 +307,6 @@ def circuit(p): qml.RX(p[2] / 2, wires=0) return qml.expval(qml.PauliZ(0)) - assert circuit.gradient_fn == "backprop" res = circuit(weights) expected = np.cos(3 * x) * np.cos(y) * np.cos(z / 2) - np.sin(3 * x) * np.sin(z / 2) @@ -381,9 +379,6 @@ def circuit(x): res = circuit1(p_torch) assert qml.math.allclose(qml.math.stack(res), circuit2(p), atol=tol, rtol=0) - assert circuit1.gradient_fn == "backprop" - assert circuit2.gradient_fn is qml.gradients.param_shift - grad = torch.autograd.functional.jacobian(circuit1, p_torch) grad_expected = torch.autograd.functional.jacobian(circuit2, p_torch_2) @@ -608,14 +603,6 @@ def cost(params): expected_cost = (np.sin(lam) * np.sin(phi) - np.cos(theta) * np.cos(lam) * np.cos(phi)) ** 2 assert torch.allclose(res, torch.tensor(expected_cost), atol=tol, rtol=0) - # Check that the correct differentiation method is being used. - if diff_method == "backprop": - assert circuit.gradient_fn == "backprop" - elif diff_method == "parameter-shift": - assert circuit.gradient_fn is qml.gradients.param_shift - else: - assert circuit.gradient_fn is qml.gradients.finite_diff - res.backward() res = params.grad diff --git a/tests/devices/test_default_qubit_autograd.py b/tests/devices/test_default_qubit_autograd.py deleted file mode 100644 index 2041e2f402a..00000000000 --- a/tests/devices/test_default_qubit_autograd.py +++ /dev/null @@ -1,803 +0,0 @@ -# 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. -# 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 the ``default.qubit.autograd`` device. -""" -import pytest - -import pennylane as qml -from pennylane import DeviceError -from pennylane import numpy as np -from pennylane.devices.default_qubit_autograd import DefaultQubitAutograd - - -@pytest.mark.autograd -def test_analytic_deprecation(): - """Tests if the kwarg `analytic` is used and displays error message.""" - msg = "The analytic argument has been replaced by shots=None. " - msg += "Please use shots=None instead of analytic=True." - - with pytest.raises( - DeviceError, - match=msg, - ): - qml.device("default.qubit.autograd", wires=1, shots=1, analytic=True) - - -@pytest.mark.autograd -class TestQNodeIntegration: - """Integration tests for default.qubit.autograd. This test ensures it integrates - properly with the PennyLane UI, in particular the new QNode.""" - - def test_defines_correct_capabilities(self): - """Test that the device defines the right capabilities""" - - dev = qml.device("default.qubit.autograd", wires=1) - cap = dev.capabilities() - capabilities = { - "model": "qubit", - "supports_finite_shots": True, - "supports_tensor_observables": True, - "returns_probs": True, - "returns_state": True, - "supports_inverse_operations": True, - "supports_analytic_computation": True, - "passthru_interface": "autograd", - "supports_broadcasting": True, - "passthru_devices": { - "autograd": "default.qubit.autograd", - "jax": "default.qubit.jax", - }, - } - assert cap == capabilities - - def test_load_device(self): - """Test that the plugin device loads correctly""" - dev = qml.device("default.qubit.autograd", wires=2) - assert dev.num_wires == 2 - assert dev.shots == qml.measurements.Shots(None) - assert dev.short_name == "default.qubit.autograd" - assert dev.capabilities()["passthru_interface"] == "autograd" - - def test_qubit_circuit(self, tol): - """Test that the device provides the correct - result for a simple circuit.""" - p = np.array(0.543) - - dev = qml.device("default.qubit.autograd", wires=1) - - @qml.qnode(dev, interface="autograd") - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliY(0)) - - expected = -np.sin(p) - - assert circuit.gradient_fn == "backprop" - assert np.isclose(circuit(p), expected, atol=tol, rtol=0) - - def test_qubit_circuit_broadcasted(self, tol): - """Test that the device provides the correct - result for a simple broadcasted circuit.""" - p = np.array([0.543, 0.21, 1.5]) - - dev = qml.device("default.qubit.autograd", wires=1) - - @qml.qnode(dev, interface="autograd") - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliY(0)) - - expected = -np.sin(p) - - assert circuit.gradient_fn == "backprop" - assert np.allclose(circuit(p), expected, atol=tol, rtol=0) - - def test_correct_state(self, tol): - """Test that the device state is correct after applying a - quantum function on the device""" - - dev = qml.device("default.qubit.autograd", wires=2) - - state = dev.state - expected = np.array([1, 0, 0, 0]) - assert np.allclose(state, expected, atol=tol, rtol=0) - - @qml.qnode(dev, interface="autograd", diff_method="backprop") - def circuit(): - qml.Hadamard(wires=0) - qml.RZ(np.pi / 4, wires=0) - return qml.expval(qml.PauliZ(0)) - - circuit() - state = dev.state - - amplitude = np.exp(-1j * np.pi / 8) / np.sqrt(2) - - expected = np.array([amplitude, 0, np.conj(amplitude), 0]) - assert np.allclose(state, expected, atol=tol, rtol=0) - - def test_correct_state_broadcasted(self, tol): - """Test that the device state is correct after applying a - broadcasted quantum function on the device""" - - dev = qml.device("default.qubit.autograd", wires=2) - - state = dev.state - expected = np.array([1, 0, 0, 0]) - assert np.allclose(state, expected, atol=tol, rtol=0) - - @qml.qnode(dev, interface="autograd", diff_method="backprop") - def circuit(): - qml.Hadamard(wires=0) - qml.RZ(np.array([np.pi / 4, np.pi / 2]), wires=0) - return qml.expval(qml.PauliZ(0)) - - circuit() - state = dev.state - - phase = np.exp(-1j * np.pi / 8) - - expected = np.array( - [ - [phase / np.sqrt(2), 0, np.conj(phase) / np.sqrt(2), 0], - [phase**2 / np.sqrt(2), 0, np.conj(phase) ** 2 / np.sqrt(2), 0], - ] - ) - assert np.allclose(state, expected, atol=tol, rtol=0) - - -@pytest.mark.autograd -class TestDtypePreserved: - """Test that the user-defined dtype of the device is preserved for QNode - evaluation""" - - @pytest.mark.parametrize("r_dtype", [np.float32, np.float64]) - @pytest.mark.parametrize( - "measurement", - [ - qml.expval(qml.PauliY(0)), - qml.var(qml.PauliY(0)), - qml.probs(wires=[1]), - qml.probs(wires=[2, 0]), - ], - ) - def test_real_dtype(self, r_dtype, measurement): - """Test that the default qubit plugin returns the correct - real data type for a simple circuit""" - p = 0.543 - - dev = qml.device("default.qubit.autograd", wires=3) - dev.target_device.R_DTYPE = r_dtype - - @qml.qnode(dev, diff_method="backprop") - def circuit(x): - qml.RX(x, wires=0) - return qml.apply(measurement) - - res = circuit(p) - assert res.dtype == r_dtype - - @pytest.mark.parametrize("r_dtype", [np.float32, np.float64]) - @pytest.mark.parametrize( - "measurement", - [ - qml.expval(qml.PauliY(0)), - qml.var(qml.PauliY(0)), - qml.probs(wires=[1]), - qml.probs(wires=[2, 0]), - ], - ) - def test_real_dtype_broadcasted(self, r_dtype, measurement): - """Test that the default qubit plugin returns the correct - real data type for a simple broadcasted circuit""" - p = np.array([0.543, 0.21, 1.6]) - - dev = qml.device("default.qubit.autograd", wires=3) - dev.target_device.R_DTYPE = r_dtype - - @qml.qnode(dev, diff_method="backprop") - def circuit(x): - qml.RX(x, wires=0) - return qml.apply(measurement) - - res = circuit(p) - assert res.dtype == r_dtype - - @pytest.mark.parametrize("c_dtype_name", ["complex64", "complex128"]) - @pytest.mark.parametrize( - "measurement", - [qml.state(), qml.density_matrix(wires=[1]), qml.density_matrix(wires=[2, 0])], - ) - def test_complex_dtype(self, c_dtype_name, measurement): - """Test that the default qubit plugin returns the correct - complex data type for a simple circuit""" - p = 0.543 - c_dtype = np.dtype(c_dtype_name) - - dev = qml.device("default.qubit.autograd", wires=3) - dev.target_device.C_DTYPE = c_dtype - - @qml.qnode(dev, diff_method="backprop") - def circuit(x): - qml.RX(x, wires=0) - return qml.apply(measurement) - - res = circuit(p) - assert res.dtype == c_dtype - - @pytest.mark.parametrize("c_dtype_name", ["complex64", "complex128"]) - def test_complex_dtype_broadcasted(self, c_dtype_name): - """Test that the default qubit plugin returns the correct - complex data type for a simple broadcasted circuit""" - p = np.array([0.543, 0.21, 1.6]) - c_dtype = np.dtype(c_dtype_name) - - dev = qml.device("default.qubit.autograd", wires=3) - dev.target_device.C_DTYPE = c_dtype - - measurement = qml.state() - - @qml.qnode(dev, diff_method="backprop") - def circuit(x): - qml.RX(x, wires=0) - return qml.apply(measurement) - - res = circuit(p) - assert res.dtype == c_dtype - - -@pytest.mark.autograd -class TestPassthruIntegration: - """Tests for integration with the PassthruQNode""" - - def test_jacobian_variable_multiply(self, tol): - """Test that jacobian of a QNode with an attached default.qubit.autograd device - gives the correct result in the case of parameters multiplied by scalars""" - x = 0.43316321 - y = 0.2162158 - z = 0.75110998 - weights = np.array([x, y, z], requires_grad=True) - - dev = qml.device("default.qubit.autograd", wires=1) - - @qml.qnode(dev, interface="autograd", diff_method="backprop") - def circuit(p): - qml.RX(3 * p[0], wires=0) - qml.RY(p[1], wires=0) - qml.RX(p[2] / 2, wires=0) - return qml.expval(qml.PauliZ(0)) - - assert circuit.gradient_fn == "backprop" - res = circuit(weights) - - expected = np.cos(3 * x) * np.cos(y) * np.cos(z / 2) - np.sin(3 * x) * np.sin(z / 2) - assert np.allclose(res, expected, atol=tol, rtol=0) - - grad_fn = qml.jacobian(circuit, 0) - res = grad_fn(np.array(weights)) - - expected = np.array( - [ - -3 * (np.sin(3 * x) * np.cos(y) * np.cos(z / 2) + np.cos(3 * x) * np.sin(z / 2)), - -np.cos(3 * x) * np.sin(y) * np.cos(z / 2), - -0.5 * (np.sin(3 * x) * np.cos(z / 2) + np.cos(3 * x) * np.cos(y) * np.sin(z / 2)), - ] - ) - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_jacobian_variable_multiply_broadcasted(self, tol): - """Test that jacobian of a QNode with an attached default.qubit.autograd device - gives the correct result in the case of broadcasted parameters multiplied by scalars""" - x = np.array([0.43316321, 92.1, -0.5129]) - y = np.array([0.2162158, 0.241, -0.51]) - z = np.array([0.75110998, 0.12512, 9.12]) - weights = np.array([x, y, z], requires_grad=True) - - dev = qml.device("default.qubit.autograd", wires=1) - - @qml.qnode(dev, interface="autograd", diff_method="backprop") - def circuit(p): - qml.RX(3 * p[0], wires=0) - qml.RY(p[1], wires=0) - qml.RX(p[2] / 2, wires=0) - return qml.expval(qml.PauliZ(0)) - - assert circuit.gradient_fn == "backprop" - res = circuit(weights) - - expected = np.cos(3 * x) * np.cos(y) * np.cos(z / 2) - np.sin(3 * x) * np.sin(z / 2) - assert np.allclose(res, expected, atol=tol, rtol=0) - - grad_fn = qml.jacobian(circuit, 0) - res = grad_fn(np.array(weights)) - - expected = np.array( - [ - -3 * (np.sin(3 * x) * np.cos(y) * np.cos(z / 2) + np.cos(3 * x) * np.sin(z / 2)), - -np.cos(3 * x) * np.sin(y) * np.cos(z / 2), - -0.5 * (np.sin(3 * x) * np.cos(z / 2) + np.cos(3 * x) * np.cos(y) * np.sin(z / 2)), - ] - ) - - assert all(np.allclose(res[i, :, i], expected[:, i], atol=tol, rtol=0) for i in range(3)) - - def test_jacobian_repeated(self, tol): - """Test that jacobian of a QNode with an attached default.qubit.autograd device - gives the correct result in the case of repeated parameters""" - x = 0.43316321 - y = 0.2162158 - z = 0.75110998 - p = np.array([x, y, z], requires_grad=True) - dev = qml.device("default.qubit.autograd", wires=1) - - @qml.qnode(dev, interface="autograd", diff_method="backprop") - def circuit(x): - qml.RX(x[1], wires=0) - qml.Rot(x[0], x[1], x[2], wires=0) - return qml.expval(qml.PauliZ(0)) - - res = circuit(p) - - expected = np.cos(y) ** 2 - np.sin(x) * np.sin(y) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) - - grad_fn = qml.jacobian(circuit, 0) - res = grad_fn(p) - - expected = np.array( - [-np.cos(x) * np.sin(y) ** 2, -2 * (np.sin(x) + 1) * np.sin(y) * np.cos(y), 0] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_jacobian_repeated_broadcasted(self, tol): - """Test that jacobian of a QNode with an attached default.qubit.autograd device - gives the correct result in the case of repeated broadcasted parameters""" - x = np.array([0.43316321, 92.1, -0.5129]) - y = np.array([0.2162158, 0.241, -0.51]) - z = np.array([0.75110998, 0.12512, 9.12]) - p = np.array([x, y, z], requires_grad=True) - dev = qml.device("default.qubit.autograd", wires=1) - - @qml.qnode(dev, interface="autograd", diff_method="backprop") - def circuit(x): - qml.RX(x[1], wires=0) - qml.Rot(x[0], x[1], x[2], wires=0) - return qml.expval(qml.PauliZ(0)) - - res = circuit(p) - - expected = np.cos(y) ** 2 - np.sin(x) * np.sin(y) ** 2 - assert np.allclose(res, expected, atol=tol, rtol=0) - - grad_fn = qml.jacobian(circuit, 0) - res = grad_fn(p) - - expected = np.array( - [ - -np.cos(x) * np.sin(y) ** 2, - -2 * (np.sin(x) + 1) * np.sin(y) * np.cos(y), - np.zeros_like(x), - ] - ) - assert all(np.allclose(res[i, :, i], expected[:, i], atol=tol, rtol=0) for i in range(3)) - - def test_jacobian_agrees_backprop_parameter_shift(self, tol): - """Test that jacobian of a QNode with an attached default.qubit.autograd device - gives the correct result with respect to the parameter-shift method""" - p = np.array([0.43316321, 0.2162158, 0.75110998, 0.94714242], requires_grad=True) - - def circuit(x): - for i in range(0, len(p), 2): - qml.RX(x[i], wires=0) - qml.RY(x[i + 1], wires=1) - for i in range(2): - qml.CNOT(wires=[i, i + 1]) - return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(1)) - - dev1 = qml.device("default.qubit.legacy", wires=3) - dev2 = qml.device("default.qubit.legacy", wires=3) - - def cost(x): - return qml.math.stack(circuit(x)) - - circuit1 = qml.QNode(cost, dev1, diff_method="backprop", interface="autograd") - circuit2 = qml.QNode(cost, dev2, diff_method="parameter-shift") - - res = circuit1(p) - - assert np.allclose(res, circuit2(p), atol=tol, rtol=0) - - assert circuit1.gradient_fn == "backprop" - assert circuit2.gradient_fn is qml.gradients.param_shift - - grad_fn = qml.jacobian(circuit1, 0) - res = grad_fn(p) - assert np.allclose(res, qml.jacobian(circuit2)(p), atol=tol, rtol=0) - - @pytest.mark.parametrize("wires", [[0], ["abc"]]) - def test_state_differentiability(self, wires, tol): - """Test that the device state can be differentiated""" - dev = qml.device("default.qubit.autograd", wires=wires) - - @qml.qnode(dev, diff_method="backprop", interface="autograd") - def circuit(a): - qml.RY(a, wires=wires[0]) - return qml.state() - - a = np.array(0.54, requires_grad=True) - - def cost(a): - """A function of the device quantum state, as a function - of input QNode parameters.""" - res = np.abs(circuit(a)) ** 2 - return res[1] - res[0] - - grad = qml.grad(cost)(a) - expected = np.sin(a) - assert np.allclose(grad, expected, atol=tol, rtol=0) - - def test_state_differentiability_broadcasted(self, tol): - """Test that the broadcasted device state can be differentiated""" - dev = qml.device("default.qubit.autograd", wires=1) - - @qml.qnode(dev, diff_method="backprop", interface="autograd") - def circuit(a): - qml.RY(a, wires=0) - return qml.expval(qml.PauliZ(0)) - - a = np.array([0.54, 0.32, 1.2], requires_grad=True) - - def cost(a): - """A function of the device quantum state, as a function - of input QNode parameters.""" - circuit(a) - res = np.abs(dev.state) ** 2 - return res[:, 1] - res[:, 0] - - grad = qml.jacobian(cost)(a) - expected = np.diag(np.sin(a)) - assert np.allclose(grad, expected, atol=tol, rtol=0) - - def test_prob_differentiability(self, tol): - """Test that the device probability can be differentiated""" - dev = qml.device("default.qubit.autograd", wires=2) - - @qml.qnode(dev, diff_method="backprop", interface="autograd") - def circuit(a, b): - qml.RX(a, wires=0) - qml.RY(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[1]) - - a = np.array(0.54, requires_grad=True) - b = np.array(0.12, requires_grad=True) - - def cost(a, b): - prob_wire_1 = circuit(a, b) - return prob_wire_1[1] - prob_wire_1[0] # pylint:disable=unsubscriptable-object - - res = cost(a, b) - expected = -np.cos(a) * np.cos(b) - assert np.allclose(res, expected, atol=tol, rtol=0) - - grad = qml.grad(cost)(a, b) - expected = [np.sin(a) * np.cos(b), np.cos(a) * np.sin(b)] - assert np.allclose(grad, expected, atol=tol, rtol=0) - - def test_prob_differentiability_broadcasted(self, tol): - """Test that the broadcasted device probability can be differentiated""" - dev = qml.device("default.qubit.autograd", wires=2) - - @qml.qnode(dev, diff_method="backprop", interface="autograd") - def circuit(a, b): - qml.RX(a, wires=0) - qml.RY(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[1]) - - a = np.array([0.54, 0.32, 1.2], requires_grad=True) - b = np.array(0.12, requires_grad=True) - - def cost(a, b): - prob_wire_1 = circuit(a, b) - return prob_wire_1[:, 1] - prob_wire_1[:, 0] # pylint:disable=unsubscriptable-object - - res = cost(a, b) - expected = -np.cos(a) * np.cos(b) - assert np.allclose(res, expected, atol=tol, rtol=0) - - jac = qml.jacobian(cost)(a, b) - expected = np.array([np.sin(a) * np.cos(b), np.cos(a) * np.sin(b)]) - expected = (np.diag(expected[0]), expected[1]) # Only first parameter is broadcasted - assert all(np.allclose(j, e, atol=tol, rtol=0) for j, e in zip(jac, expected)) - - def test_backprop_gradient(self, tol): - """Tests that the gradient of the qnode is correct""" - dev = qml.device("default.qubit.autograd", wires=2) - - @qml.qnode(dev, diff_method="backprop", interface="autograd") - def circuit(a, b): - qml.RX(a, wires=0) - qml.CRX(b, wires=[0, 1]) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - a = np.array(-0.234, requires_grad=True) - b = np.array(0.654, requires_grad=True) - - res = circuit(a, b) - expected_cost = 0.5 * (np.cos(a) * np.cos(b) + np.cos(a) - np.cos(b) + 1) - assert np.allclose(res, expected_cost, atol=tol, rtol=0) - - res = qml.grad(circuit)(a, b) - expected_grad = np.array( - [-0.5 * np.sin(a) * (np.cos(b) + 1), 0.5 * np.sin(b) * (1 - np.cos(a))] - ) - assert np.allclose(res, expected_grad, atol=tol, rtol=0) - - def test_backprop_gradient_broadcasted(self, tol): - """Tests that the gradient of the broadcasted qnode is correct""" - dev = qml.device("default.qubit.autograd", wires=2) - - @qml.qnode(dev, diff_method="backprop", interface="autograd") - def circuit(a, b): - qml.RX(a, wires=0) - qml.CRX(b, wires=[0, 1]) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - a = np.array(0.12, requires_grad=True) - b = np.array([0.54, 0.32, 1.2], requires_grad=True) - - res = circuit(a, b) - expected_cost = 0.5 * (np.cos(a) * np.cos(b) + np.cos(a) - np.cos(b) + 1) - assert np.allclose(res, expected_cost, atol=tol, rtol=0) - - res = qml.jacobian(circuit)(a, b) - expected = np.array([-0.5 * np.sin(a) * (np.cos(b) + 1), 0.5 * np.sin(b) * (1 - np.cos(a))]) - expected = (expected[0], np.diag(expected[1])) - assert all(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) - - @pytest.mark.parametrize( - "x, shift", - [np.array((0.0, 0.0), requires_grad=True), np.array((0.5, -0.5), requires_grad=True)], - ) - def test_hessian_at_zero(self, x, shift): - """Tests that the Hessian at vanishing state vector amplitudes - is correct.""" - dev = qml.device("default.qubit.autograd", wires=1) - - @qml.qnode(dev, interface="autograd", diff_method="backprop") - def circuit(x): - qml.RY(shift, wires=0) - qml.RY(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - assert qml.math.isclose(qml.jacobian(circuit)(x), 0.0) - assert qml.math.isclose(qml.jacobian(qml.jacobian(circuit))(x), -1.0) - assert qml.math.isclose(qml.grad(qml.grad(circuit))(x), -1.0) - - @pytest.mark.parametrize("operation", [qml.U3, qml.U3.compute_decomposition]) - @pytest.mark.parametrize("diff_method", ["backprop", "parameter-shift", "finite-diff"]) - def test_autograd_interface_gradient(self, operation, diff_method, tol): - """Tests that the gradient of an arbitrary U3 gate is correct - using the Autograd interface, using a variety of differentiation methods.""" - dev = qml.device("default.qubit.autograd", wires=1) - state = np.array(1j * np.array([1, -1]) / np.sqrt(2), requires_grad=False) - - @qml.qnode(dev, diff_method=diff_method, interface="autograd") - def circuit(x, weights, w): - """In this example, a mixture of scalar - arguments, array arguments, and keyword arguments are used.""" - qml.StatePrep(state, wires=w) - operation(x, weights[0], weights[1], wires=w) - return qml.expval(qml.PauliX(w)) - - def cost(params): - """Perform some classical processing""" - return circuit(params[0], params[1:], w=0) ** 2 - - theta = 0.543 - phi = -0.234 - lam = 0.654 - - params = np.array([theta, phi, lam], requires_grad=True) - - res = cost(params) - expected_cost = (np.sin(lam) * np.sin(phi) - np.cos(theta) * np.cos(lam) * np.cos(phi)) ** 2 - assert np.allclose(res, expected_cost, atol=tol, rtol=0) - - # Check that the correct differentiation method is being used. - if diff_method == "backprop": - assert circuit.gradient_fn == "backprop" - elif diff_method == "parameter-shift": - assert circuit.gradient_fn is qml.gradients.param_shift - else: - assert circuit.gradient_fn is qml.gradients.finite_diff - - res = qml.grad(cost)(params) - expected_grad = ( - np.array( - [ - np.sin(theta) * np.cos(lam) * np.cos(phi), - np.cos(theta) * np.cos(lam) * np.sin(phi) + np.sin(lam) * np.cos(phi), - np.cos(theta) * np.sin(lam) * np.cos(phi) + np.cos(lam) * np.sin(phi), - ] - ) - * 2 - * (np.sin(lam) * np.sin(phi) - np.cos(theta) * np.cos(lam) * np.cos(phi)) - ) - assert np.allclose(res, expected_grad, atol=tol, rtol=0) - - @pytest.mark.xfail(reason="Not applicable anymore.") - @pytest.mark.parametrize("interface", ["tf", "torch"]) - def test_error_backprop_wrong_interface(self, interface): - """Tests that an error is raised if diff_method='backprop' but not using - the Autograd interface""" - dev = qml.device("default.qubit.autograd", wires=1) - - def circuit(x, w=None): - qml.RZ(x, wires=w) - return qml.expval(qml.PauliX(w)) - - with pytest.raises( - qml.QuantumFunctionError, - match="default.qubit.autograd only supports diff_method='backprop' when using the autograd interface", - ): - qml.qnode(dev, diff_method="backprop", interface=interface)(circuit) - - -@pytest.mark.autograd -class TestHighLevelIntegration: - """Tests for integration with higher level components of PennyLane.""" - - def test_do_not_split_analytic_autograd(self, mocker): - """Tests that the Hamiltonian is not split for shots=None using the autograd device.""" - dev = qml.device("default.qubit.autograd", wires=2) - H = qml.Hamiltonian(np.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)]) - - @qml.qnode(dev, diff_method="backprop", interface="autograd") - def circuit(): - return qml.expval(H) - - spy = mocker.spy(dev.target_device, "expval") - - circuit() - # evaluated one expval altogether - assert spy.call_count == 1 - - def test_do_not_split_analytic_autograd_broadcasted(self, mocker): - """Tests that the Hamiltonian is not split for shots=None - and broadcasting using the autograd device.""" - dev = qml.device("default.qubit.autograd", wires=2) - H = qml.Hamiltonian(np.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)]) - - @qml.qnode(dev, diff_method="backprop", interface="autograd") - def circuit(): - qml.RX(np.zeros(5), 0) - return qml.expval(H) - - spy = mocker.spy(dev.target_device, "expval") - - circuit() - # evaluated one expval altogether - assert spy.call_count == 1 - - def test_template_integration(self): - """Test that a PassthruQNode default.qubit.autograd works with templates.""" - dev = qml.device("default.qubit.autograd", wires=2) - - @qml.qnode(dev, diff_method="backprop") - def circuit(weights): - qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1]) - return qml.expval(qml.PauliZ(0)) - - shape = qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=2) - weights = np.random.random(shape, requires_grad=True) - - grad = qml.grad(circuit)(weights) - assert grad.shape == weights.shape - - -# pylint: disable=protected-access -@pytest.mark.autograd -class TestOps: - """Unit tests for operations supported by the default.qubit.autograd device""" - - def test_multirz_jacobian(self): - """Test that the patched numpy functions are used for the MultiRZ - operation and the jacobian can be computed.""" - wires = 4 - dev = qml.device("default.qubit.autograd", wires=wires) - - @qml.qnode(dev, diff_method="backprop") - def circuit(param): - qml.MultiRZ(param, wires=[0, 1]) - return qml.probs(wires=list(range(wires))) - - param = np.array(0.3, requires_grad=True) - res = qml.jacobian(circuit)(param) - assert np.allclose(res, np.zeros(wires**2)) - - def test_full_subsystem(self, mocker): - """Test applying a state vector to the full subsystem""" - dev = DefaultQubitAutograd(wires=["a", "b", "c"]) - state = np.array([1, 0, 0, 0, 1, 0, 1, 1]) / 2.0 - state_wires = qml.wires.Wires(["a", "b", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - - assert np.all(dev._state.flatten() == state) - spy.assert_not_called() - - def test_partial_subsystem(self, mocker): - """Test applying a state vector to a subset of wires of the full subsystem""" - - dev = DefaultQubitAutograd(wires=["a", "b", "c"]) - state = np.array([1, 0, 1, 0]) / np.sqrt(2.0) - state_wires = qml.wires.Wires(["a", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - res = np.sum(dev._state, axis=(1,)).flatten() - - assert np.all(res == state) - spy.assert_called() - - -@pytest.mark.autograd -class TestOpsBroadcasted: - """Unit tests for broadcasted operations supported by the default.qubit.autograd device""" - - def test_multirz_jacobian_broadcasted(self): - """Test that the patched numpy functions are used for the MultiRZ - operation and the jacobian can be computed.""" - wires = 4 - dev = qml.device("default.qubit.autograd", wires=wires) - - @qml.qnode(dev, diff_method="backprop") - def circuit(param): - qml.MultiRZ(param, wires=[0, 1]) - return qml.probs(wires=list(range(wires))) - - param = np.array([0.3, 0.9, -4.3], requires_grad=True) - res = qml.jacobian(circuit)(param) - assert np.allclose(res, np.zeros((3, wires**2, 3))) - - def test_full_subsystem_broadcasted(self, mocker): - """Test applying a state vector to the full subsystem""" - dev = DefaultQubitAutograd(wires=["a", "b", "c"]) - state = np.array([[1, 0, 0, 0, 1, 0, 1, 1], [0, 0, 0, 1, 1, 1, 1, 0]]) / 2.0 - state_wires = qml.wires.Wires(["a", "b", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - - assert np.all(dev._state.reshape((2, 8)) == state) - spy.assert_not_called() - - def test_partial_subsystem_broadcasted(self, mocker): - """Test applying a state vector to a subset of wires of the full subsystem""" - - dev = DefaultQubitAutograd(wires=["a", "b", "c"]) - state = np.array([[1, 0, 1, 0], [0, 1, 0, 1], [1, 1, 0, 0]]) / np.sqrt(2.0) - state_wires = qml.wires.Wires(["a", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - res = np.sum(dev._state, axis=(2,)).reshape((3, 4)) - - assert np.allclose(res, state) - spy.assert_called() diff --git a/tests/devices/test_default_qubit_jax.py b/tests/devices/test_default_qubit_jax.py deleted file mode 100644 index 4e151e5b986..00000000000 --- a/tests/devices/test_default_qubit_jax.py +++ /dev/null @@ -1,1336 +0,0 @@ -# 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. -""" -Integration tests for the ``default.qubit.jax`` device. -""" -import numpy as np -import pytest - -import pennylane as qml -from pennylane import DeviceError -from pennylane.pulse import ParametrizedHamiltonian - -jax = pytest.importorskip("jax", minversion="0.2") -from pennylane.devices.default_qubit_jax import ( # pylint: disable=wrong-import-position - DefaultQubitJax, -) - -jnp = jax.numpy - - -@pytest.mark.jax -def test_analytic_deprecation(): - """Tests if the kwarg `analytic` is used and displays error message.""" - msg = "The analytic argument has been replaced by shots=None. " - msg += "Please use shots=None instead of analytic=True." - - with pytest.raises( - DeviceError, - match=msg, - ): - qml.device("default.qubit.jax", wires=1, shots=1, analytic=True) - - -# pylint: disable=too-many-public-methods -@pytest.mark.jax -class TestQNodeIntegration: - """Integration tests for default.qubit.jax. This test ensures it integrates - properly with the PennyLane UI, in particular the new QNode.""" - - def test_defines_correct_capabilities(self): - """Test that the device defines the right capabilities""" - - dev = qml.device("default.qubit.jax", wires=1) - cap = dev.capabilities() - capabilities = { - "model": "qubit", - "supports_finite_shots": True, - "supports_tensor_observables": True, - "returns_probs": True, - "returns_state": True, - "supports_inverse_operations": True, - "supports_analytic_computation": True, - "supports_broadcasting": True, - "passthru_interface": "jax", - "passthru_devices": { - "autograd": "default.qubit.autograd", - "jax": "default.qubit.jax", - }, - } - assert cap == capabilities - - def test_defines_correct_capabilities_directly_from_class(self): - """Test that the device defines the right capabilities""" - - dev = DefaultQubitJax(wires=1) - cap = dev.capabilities() - assert cap["passthru_interface"] == "jax" - - def test_load_device(self): - """Test that the plugin device loads correctly""" - dev = qml.device("default.qubit.jax", wires=2) - assert dev.num_wires == 2 - assert dev.shots == qml.measurements.Shots(None) - assert dev.short_name == "default.qubit.jax" - assert dev.capabilities()["passthru_interface"] == "jax" - - @pytest.mark.parametrize( - "jax_enable_x64, c_dtype, r_dtype", - ([True, np.complex128, np.float64], [False, np.complex64, np.float32]), - ) - def test_float_precision(self, jax_enable_x64, c_dtype, r_dtype): - """Test that the plugin device uses the same float precision as the jax config.""" - jax.config.update("jax_enable_x64", jax_enable_x64) - dev = qml.device("default.qubit.jax", wires=2) - assert dev.state.dtype == c_dtype - assert dev.state.real.dtype == r_dtype - jax.config.update("jax_enable_x64", True) - - def test_qubit_circuit(self, tol): - """Test that the device provides the correct - result for a simple circuit.""" - p = jnp.array(0.543) - - dev = qml.device("default.qubit.jax", wires=1) - - @qml.qnode(dev, interface="jax") - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliY(0)) - - expected = -jnp.sin(p) - assert jnp.isclose(circuit(p), expected, atol=tol, rtol=0) - - def test_qubit_circuit_with_jit(self, tol, benchmark): - """Test that the device provides the correct - result for a simple circuit under a jax.jit.""" - p = jnp.array(0.543) - - dev = qml.device("default.qubit.jax", wires=1) - - @jax.jit - @qml.qnode(dev, interface="jax") - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliY(0)) - - res = benchmark(circuit, p) - expected = -jnp.sin(p) - # Do not test isinstance here since the @jax.jit changes the function - # type. Just test that it works and spits our the right value. - assert jnp.isclose(res, expected, atol=tol, rtol=0) - - # Test with broadcasted parameters - p = jnp.array([0.543, 0.21, 1.5]) - expected = -jnp.sin(p) - assert jnp.allclose(circuit(p), expected, atol=tol, rtol=0) - - def test_qubit_circuit_broadcasted(self, tol): - """Test that the device provides the correct - result for a simple broadcasted circuit.""" - p = jnp.array([0.543, 0.21, 1.5]) - - dev = qml.device("default.qubit.jax", wires=1) - - @qml.qnode(dev, interface="jax") - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliY(0)) - - expected = -jnp.sin(p) - - assert jnp.allclose(circuit(p), expected, atol=tol, rtol=0) - - def test_correct_state(self, tol): - """Test that the device state is correct after applying a - quantum function on the device""" - - dev = qml.device("default.qubit.jax", wires=2) - - state = dev.state - expected = jnp.array([1, 0, 0, 0]) - assert jnp.allclose(state, expected, atol=tol, rtol=0) - - @qml.qnode(dev, interface="jax", diff_method="backprop") - def circuit(): - qml.Hadamard(wires=0) - qml.RZ(jnp.pi / 4, wires=0) - return qml.expval(qml.PauliZ(0)) - - circuit() - state = dev.state - - amplitude = jnp.exp(-1j * jnp.pi / 8) / jnp.sqrt(2) - - expected = jnp.array([amplitude, 0, jnp.conj(amplitude), 0]) - assert jnp.allclose(state, expected, atol=tol, rtol=0) - - def test_correct_state_broadcasted(self, tol): - """Test that the device state is correct after applying a - broadcasted quantum function on the device""" - - dev = qml.device("default.qubit.jax", wires=2) - - state = dev.state - expected = jnp.array([1, 0, 0, 0]) - assert jnp.allclose(state, expected, atol=tol, rtol=0) - - @qml.qnode(dev, interface="jax", diff_method="backprop") - def circuit(): - qml.Hadamard(wires=0) - qml.RZ(jnp.array([np.pi / 4, np.pi / 2]), wires=0) - return qml.expval(qml.PauliZ(0)) - - circuit() - state = dev.state - - phase = jnp.exp(-1j * jnp.pi / 8) - - expected = np.array( - [ - [phase / jnp.sqrt(2), 0, jnp.conj(phase) / jnp.sqrt(2), 0], - [phase**2 / jnp.sqrt(2), 0, jnp.conj(phase) ** 2 / jnp.sqrt(2), 0], - ] - ) - assert jnp.allclose(state, expected, atol=tol, rtol=0) - - def test_correct_state_returned(self, tol): - """Test that the device state is correct after applying a - quantum function on the device""" - dev = qml.device("default.qubit.jax", wires=2) - - @qml.qnode(dev, interface="jax", diff_method="backprop") - def circuit(): - qml.Hadamard(wires=0) - qml.RZ(jnp.pi / 4, wires=0) - return qml.state() - - state = circuit() - - amplitude = jnp.exp(-1j * jnp.pi / 8) / jnp.sqrt(2) - - expected = jnp.array([amplitude, 0, jnp.conj(amplitude), 0]) - assert jnp.allclose(state, expected, atol=tol, rtol=0) - - def test_correct_state_returned_broadcasted(self, tol): - """Test that the device state is correct after applying a - broadcasted quantum function on the device""" - dev = qml.device("default.qubit.jax", wires=2) - - @qml.qnode(dev, interface="jax", diff_method="backprop") - def circuit(): - qml.Hadamard(wires=0) - qml.RZ(jnp.array([np.pi / 4, np.pi / 2]), wires=0) - return qml.state() - - state = circuit() - - phase = jnp.exp(-1j * jnp.pi / 8) - - expected = np.array( - [ - [phase / jnp.sqrt(2), 0, jnp.conj(phase) / jnp.sqrt(2), 0], - [phase**2 / jnp.sqrt(2), 0, jnp.conj(phase) ** 2 / jnp.sqrt(2), 0], - ] - ) - assert jnp.allclose(state, expected, atol=tol, rtol=0) - - def test_probs_jax(self, tol, benchmark): - """Test that returning probs works with jax""" - dev = qml.device("default.qubit.jax", wires=1, shots=100) - expected = jnp.array([0.0, 1.0]) - - @qml.qnode(dev, interface="jax", diff_method=None) - def circuit(): - qml.PauliX(wires=0) - return qml.probs(wires=0) - - result = benchmark(circuit) - assert jnp.allclose(result, expected, atol=tol) - - def test_probs_jax_broadcasted(self, tol): - """Test that returning probs works with jax""" - dev = qml.device("default.qubit.jax", wires=1, shots=100) - expected = jnp.array([[0.0, 1.0]] * 3) - - @qml.qnode(dev, interface="jax", diff_method=None) - def circuit(): - qml.RX(jnp.zeros(3), 0) - qml.PauliX(wires=0) - return qml.probs(wires=0) - - result = circuit() - assert jnp.allclose(result, expected, atol=tol) - - def test_probs_jax_jit(self, tol): - """Test that returning probs works with jax and jit""" - dev = qml.device("default.qubit.jax", wires=1, shots=100) - expected = jnp.array([0.0, 1.0]) - - @qml.qnode(dev, interface="jax", diff_method=None) - def circuit(z): - qml.RX(z, wires=0) - qml.PauliX(wires=0) - return qml.probs(wires=0) - - result = circuit(0.0) - assert jnp.allclose(result, expected, atol=tol) - - # Test with broadcasting - result = circuit(jnp.zeros(3)) - expected = jnp.array([[0.0, 1.0]] * 3) - assert jnp.allclose(result, expected, atol=tol) - - def test_custom_shots_probs_jax_jit(self, tol): - """Test that returning probs works with jax and jit when using custom shot vector""" - dev = qml.device("default.qubit.jax", wires=1, shots=(3, 2)) - expected = jnp.array([[0.0, 1.0], [0.0, 1.0]]) - - @jax.jit - @qml.qnode(dev, diff_method=None, interface="jax") - def circuit(): - qml.PauliX(wires=0) - return qml.probs(wires=0) - - result = circuit() - assert jnp.allclose(qml.math.hstack(result[0]), expected[0], atol=tol) - assert jnp.allclose(qml.math.hstack(result[1]), expected[1], atol=tol) - - @pytest.mark.skip("Shot lists are not supported with broadcasting yet") - def test_custom_shots_probs_jax_jit_broadcasted(self, tol): - """Test that returning probs works with jax and jit when - using a custom shot vector and broadcasting""" - dev = qml.device("default.qubit.jax", wires=1, shots=(2, 2)) - expected = jnp.array([[[0.0, 1.0], [0.0, 1.0]]] * 5) - - @jax.jit - @qml.qnode(dev, diff_method=None, interface="jax") - def circuit(): - qml.RX(jnp.zeros(5), 0) - qml.PauliX(wires=0) - return qml.probs(wires=0) - - result = circuit() - assert jnp.allclose(result, expected, atol=tol) - - def test_sampling_with_jit(self): - """Test that sampling works with a jax.jit""" - - @jax.jit - def circuit(x, key): - dev = qml.device("default.qubit.jax", wires=1, shots=1000, prng_key=key) - - @qml.qnode(dev, interface="jax", diff_method=None) - def inner_circuit(): - qml.RX(x, wires=0) - qml.Hadamard(0) - return qml.sample(qml.PauliZ(wires=0)) - - return inner_circuit() - - a = circuit(0.0, jax.random.PRNGKey(0)) - b = circuit(0.0, jax.random.PRNGKey(0)) - c = circuit(0.0, jax.random.PRNGKey(1)) - np.testing.assert_array_equal(a, b) - assert not np.all(a == c) - - # Test with broadcasting - d = circuit(jnp.zeros(5), jax.random.PRNGKey(9)) - assert qml.math.shape(d) == (5, 1000) - - @pytest.mark.parametrize( - "state_vector", - [np.array([0.5 + 0.5j, 0.5 + 0.5j, 0, 0]), jnp.array([0.5 + 0.5j, 0.5 + 0.5j, 0, 0])], - ) - def test_qubit_state_vector_arg_jax_jit(self, state_vector, tol): - """Test that Qubit state vector as argument works with a jax.jit""" - dev = qml.device("default.qubit.jax", wires=list(range(2))) - - @jax.jit - @qml.qnode(dev, interface="jax") - def circuit(x): - wires = list(range(2)) - qml.StatePrep(x, wires=wires) - return [qml.expval(qml.PauliX(wires=i)) for i in wires] - - res = circuit(state_vector) - assert jnp.allclose(jnp.array(res), jnp.array([0, 1]), atol=tol, rtol=0) - - @pytest.mark.parametrize( - "state_vector", - [np.array([0.5 + 0.5j, 0.5 + 0.5j, 0, 0]), jnp.array([0.5 + 0.5j, 0.5 + 0.5j, 0, 0])], - ) - def test_qubit_state_vector_arg_jax(self, state_vector, tol): - """Test that Qubit state vector as argument works with jax""" - dev = qml.device("default.qubit.jax", wires=list(range(2))) - - @qml.qnode(dev, interface="jax") - def circuit(x): - wires = list(range(2)) - qml.StatePrep(x, wires=wires) - return [qml.expval(qml.PauliX(wires=i)) for i in wires] - - res = circuit(state_vector) - assert jnp.allclose(jnp.array(res), jnp.array([0, 1]), atol=tol, rtol=0) - - @pytest.mark.parametrize( - "state_vector", - [np.array([0.5 + 0.5j, 0.5 + 0.5j, 0, 0]), jnp.array([0.5 + 0.5j, 0.5 + 0.5j, 0, 0])], - ) - def test_qubit_state_vector_jax_jit(self, state_vector, tol): - """Test that Qubit state vector works with a jax.jit""" - dev = qml.device("default.qubit.jax", wires=list(range(2))) - - @jax.jit - @qml.qnode(dev, interface="jax") - def circuit(x): - qml.StatePrep(state_vector, wires=dev.wires) - for w in dev.wires: - qml.RZ(x, wires=w, id="x") - return qml.expval(qml.PauliZ(wires=0)) - - res = circuit(0.1) - assert jnp.allclose(jnp.array(res), 1, atol=tol, rtol=0) - - @pytest.mark.parametrize( - "state_vector", - [np.array([0.5 + 0.5j, 0.5 + 0.5j, 0, 0]), jnp.array([0.5 + 0.5j, 0.5 + 0.5j, 0, 0])], - ) - def test_qubit_state_vector_jax(self, state_vector, tol): - """Test that Qubit state vector works with a jax""" - dev = qml.device("default.qubit.jax", wires=list(range(2))) - - @qml.qnode(dev, interface="jax") - def circuit(x): - qml.StatePrep(state_vector, wires=dev.wires) - for w in dev.wires: - qml.RZ(x, wires=w, id="x") - return qml.expval(qml.PauliZ(wires=0)) - - res = circuit(0.1) - assert jnp.allclose(jnp.array(res), 1, atol=tol, rtol=0) - - @pytest.mark.parametrize( - "state_vector", - [np.array([0.1 + 0.1j, 0.2 + 0.2j, 0, 0]), jnp.array([0.1 + 0.1j, 0.2 + 0.2j, 0, 0])], - ) - def test_qubit_state_vector_jax_not_normed(self, state_vector): - """Test that an error is raised when Qubit state vector is not normed works with a jax""" - dev = qml.device("default.qubit.jax", wires=list(range(2))) - - @qml.qnode(dev, interface="jax") - def circuit(x): - qml.StatePrep(state_vector, wires=dev.wires) - for w in dev.wires: - qml.RZ(x, wires=w, id="x") - return qml.expval(qml.PauliZ(wires=0)) - - with pytest.raises(ValueError, match="The state must be a vector of norm 1.0"): - circuit(0.1) - - def test_sampling_op_by_op(self): - """Test that op-by-op sampling works as a new user would expect""" - dev = qml.device("default.qubit.jax", wires=1, shots=1000) - - @qml.qnode(dev, interface="jax", diff_method=None) - def circuit(): - qml.Hadamard(0) - return qml.sample(qml.PauliZ(wires=0)) - - a = circuit() - b = circuit() - assert not np.all(a == b) - - def test_sampling_analytic_mode(self): - """Test that when sampling with shots=None an error is raised.""" - dev = qml.device("default.qubit.jax", wires=1, shots=None) - - @qml.qnode(dev, interface="jax", diff_method=None) - def circuit(): - return qml.sample(qml.PauliZ(wires=0)) - - with pytest.raises( - qml.QuantumFunctionError, - match="The number of shots has to be explicitly set on the device " - "when using sample-based measurements.", - ): - circuit() - - def test_sampling_analytic_mode_with_counts(self): - """Test that when sampling with counts and shots=None an error is raised.""" - dev = qml.device("default.qubit.jax", wires=1, shots=None) - - @qml.qnode(dev, interface="jax", diff_method=None) - def circuit(): - return qml.counts(qml.PauliZ(wires=0)) - - with pytest.raises( - qml.QuantumFunctionError, - match="The number of shots has to be explicitly set on the device " - "when using sample-based measurements.", - ): - circuit() - - def test_gates_dont_crash(self): - """Test for gates that weren't covered by other tests.""" - dev = qml.device("default.qubit.jax", wires=2, shots=1000) - - @qml.qnode(dev, interface="jax", diff_method=None) - def circuit(): - qml.CRZ(0.0, wires=[0, 1]) - qml.CRX(0.0, wires=[0, 1]) - qml.PhaseShift(0.0, wires=0) - qml.ControlledPhaseShift(0.0, wires=[1, 0]) - qml.CRot(1.0, 0.0, 0.0, wires=[0, 1]) - qml.CRY(0.0, wires=[0, 1]) - return qml.sample(qml.PauliZ(wires=0)) - - circuit() # Just don't crash. - - def test_diagonal_doesnt_crash(self): - """Test that diagonal gates can be used.""" - dev = qml.device("default.qubit.jax", wires=1, shots=1000) - - @qml.qnode(dev, interface="jax", diff_method=None) - def circuit(): - qml.DiagonalQubitUnitary(np.array([1.0, 1.0]), wires=0) - return qml.sample(qml.PauliZ(wires=0)) - - circuit() # Just don't crash. - - def test_broadcasted_diagonal_doesnt_crash(self): - """Test that diagonal gates can be used.""" - dev = qml.device("default.qubit.jax", wires=1, shots=1000) - - @qml.qnode(dev, interface="jax", diff_method=None) - def circuit(): - qml.DiagonalQubitUnitary(np.array([[-1, -1], [1j, -1], [1.0, 1.0]]), wires=0) - return qml.sample(qml.PauliZ(wires=0)) - - circuit() # Just don't crash. - - @pytest.mark.parametrize("phi", np.pi * np.array([1e-8, 1 / 8, 1 / 4, 1 / 2, 1])) - def test_parametrized_evolution_state_vector(self, phi, mocker): - """Test that when executing a ParametrizedEvolution with ``num_wires >= device.num_wires/2`` - the `_evolve_state_vector_under_parametrized_evolution` method is used.""" - dev = qml.device("default.qubit.jax", wires=1) - H = ParametrizedHamiltonian([1], [qml.PauliX(0)]) - spy = mocker.spy(dev.target_device, "_evolve_state_vector_under_parametrized_evolution") - - @jax.jit - @qml.qnode(dev, interface="jax") - def circuit(): - qml.evolve(H)(params=[], t=phi / 2) - return qml.expval(qml.PauliZ(0)) - - @qml.qnode(dev) - def true_circuit(): - qml.RX(phi, 0) - return qml.expval(qml.PauliZ(0)) - - res = circuit() - spy.assert_called_once() - assert qml.math.allclose(res, true_circuit(), atol=1e-6) - - @pytest.mark.parametrize("phi", np.pi * np.array([1e-8, 1 / 8, 1 / 4, 1 / 2, 1])) - def test_parametrized_evolution_matrix(self, phi, mocker): - """Test that when executing a ParametrizedEvolution with ``num_wires < device.num_wires/2`` - the `_apply_operation` method is used.""" - dev = qml.device("default.qubit.jax", wires=3) - H = ParametrizedHamiltonian([1], [qml.PauliX(0)]) - spy = mocker.spy(dev.target_device, "_evolve_state_vector_under_parametrized_evolution") - spy2 = mocker.spy(dev.target_device, "_apply_operation") - - @jax.jit - @qml.qnode(dev, interface="jax") - def circuit(): - qml.evolve(H)(params=[], t=phi / 2) # corresponds to a PauliX gate - return qml.expval(qml.PauliZ(0)) - - @qml.qnode(dev) - def true_circuit(): - qml.RX(phi, 0) - return qml.expval(qml.PauliZ(0)) - - res = circuit() - spy.assert_not_called() - spy2.assert_called_once() - assert qml.math.allclose(res, true_circuit(), atol=1e-6) - - def test_parametrized_evolution_state_vector_return_intermediate(self, mocker): - """Test that when executing a ParametrizedEvolution with ``num_wires >= device.num_wires/2`` - and ``return_intermediate=True``, the ``_evolve_state_vector_under_parametrized_evolution`` - method is used.""" - dev = qml.device("default.qubit.jax", wires=1) - H = ParametrizedHamiltonian([1], [qml.PauliX(0)]) - spy = mocker.spy(dev.target_device, "_evolve_state_vector_under_parametrized_evolution") - spy2 = mocker.spy(dev.target_device, "_apply_operation") - - phi = jnp.linspace(0.3, 0.7, 7) - phi_for_RX = phi - phi[0] - - @jax.jit - @qml.qnode(dev, interface="jax") - def circuit(): - qml.evolve(H, return_intermediate=True)(params=[], t=phi / 2) - return qml.expval(qml.PauliZ(0)) - - @qml.qnode(dev) - def true_circuit(): - qml.RX(phi_for_RX, 0) - return qml.expval(qml.PauliZ(0)) - - res = circuit() - spy.assert_called_once() - spy2.assert_not_called() - assert qml.math.allclose(res, true_circuit(), atol=1e-6) - - def test_parametrized_evolution_matrix_complementary(self, mocker): - """Test that when executing a ParametrizedEvolution with ``num_wires >= device.num_wires/2`` - but with ``complementary=True``, the `_apply_operation` method is used.""" - dev = qml.device("default.qubit.jax", wires=1) - H = ParametrizedHamiltonian([1], [qml.PauliX(0)]) - spy = mocker.spy(dev.target_device, "_evolve_state_vector_under_parametrized_evolution") - spy2 = mocker.spy(dev.target_device, "_apply_operation") - - phi = jnp.linspace(0.3, 0.7, 7) - phi_for_RX = phi[-1] - phi - - @jax.jit - @qml.qnode(dev, interface="jax") - def circuit(): - qml.evolve(H, return_intermediate=True, complementary=True)(params=[], t=phi / 2) - return qml.expval(qml.PauliZ(0)) - - @qml.qnode(dev) - def true_circuit(): - qml.RX(phi_for_RX, 0) - return qml.expval(qml.PauliZ(0)) - - res = circuit() - spy.assert_not_called() - spy2.assert_called_once() - assert qml.math.allclose(res, true_circuit(), atol=1e-6) - - -@pytest.mark.jax -class TestPassthruIntegration: - """Tests for integration with the PassthruQNode""" - - @pytest.mark.parametrize("jacobian_transform", [jax.jacfwd, jax.jacrev]) - def test_jacobian_variable_multiply(self, tol, jacobian_transform, benchmark): - """Test that jacobian of a QNode with an attached default.qubit.jax device - gives the correct result in the case of parameters multiplied by scalars""" - x = 0.43316321 - y = 0.2162158 - z = 0.75110998 - weights = jnp.array([x, y, z]) - - dev = qml.device("default.qubit.jax", wires=1) - - @qml.qnode(dev, interface="jax") - def circuit(p): - qml.RX(3 * p[0], wires=0) - qml.RY(p[1], wires=0) - qml.RX(p[2] / 2, wires=0) - return qml.expval(qml.PauliZ(0)) - - def workload(): - return circuit(weights), jacobian_transform(circuit, 0)(jnp.array(weights)) - - res, grad = benchmark(workload) - - expected = jnp.cos(3 * x) * jnp.cos(y) * jnp.cos(z / 2) - jnp.sin(3 * x) * jnp.sin(z / 2) - assert jnp.allclose(res, expected, atol=tol, rtol=0) - - expected = jnp.array( - [ - -3 - * (jnp.sin(3 * x) * jnp.cos(y) * jnp.cos(z / 2) + jnp.cos(3 * x) * jnp.sin(z / 2)), - -jnp.cos(3 * x) * jnp.sin(y) * jnp.cos(z / 2), - -0.5 - * (jnp.sin(3 * x) * jnp.cos(z / 2) + jnp.cos(3 * x) * jnp.cos(y) * jnp.sin(z / 2)), - ] - ) - - assert jnp.allclose(grad, expected, atol=tol, rtol=0) - - def test_jacobian_variable_multiply_broadcasted(self, tol): - """Test that jacobian of a QNode with an attached default.qubit.jax device - gives the correct result in the case of broadcasted parameters multiplied by scalars""" - x = jnp.array([0.43316321, 92.1, -0.5129]) - y = jnp.array([0.2162158, 0.241, -0.51]) - z = jnp.array([0.75110998, 0.12512, 9.12]) - weights = jnp.array([x, y, z]) - - dev = qml.device("default.qubit.jax", wires=1) - - @qml.qnode(dev, interface="jax", diff_method="backprop") - def circuit(p): - qml.RX(3 * p[0], wires=0) - qml.RY(p[1], wires=0) - qml.RX(p[2] / 2, wires=0) - return qml.expval(qml.PauliZ(0)) - - assert circuit.gradient_fn == "backprop" - res = circuit(weights) - - expected = jnp.cos(3 * x) * jnp.cos(y) * jnp.cos(z / 2) - jnp.sin(3 * x) * jnp.sin(z / 2) - assert jnp.allclose(res, expected, atol=tol, rtol=0) - - grad_fn = jax.jacobian(circuit, 0) - res = grad_fn(jnp.array(weights)) - - expected = jnp.array( - [ - -3 - * (jnp.sin(3 * x) * jnp.cos(y) * jnp.cos(z / 2) + jnp.cos(3 * x) * jnp.sin(z / 2)), - -jnp.cos(3 * x) * jnp.sin(y) * jnp.cos(z / 2), - -0.5 - * (jnp.sin(3 * x) * jnp.cos(z / 2) + jnp.cos(3 * x) * jnp.cos(y) * jnp.sin(z / 2)), - ] - ) - - assert all(jnp.allclose(res[i, :, i], expected[:, i], atol=tol, rtol=0) for i in range(3)) - - @pytest.mark.parametrize("jacobian_transform", [jax.jacfwd, jax.jacrev]) - def test_jacobian_repeated(self, tol, jacobian_transform): - """Test that jacobian of a QNode with an attached default.qubit.jax device - gives the correct result in the case of repeated parameters""" - x = 0.43316321 - y = 0.2162158 - z = 0.75110998 - p = jnp.array([x, y, z]) - dev = qml.device("default.qubit.jax", wires=1) - - @qml.qnode(dev, interface="jax") - def circuit(x): - qml.RX(x[1], wires=0) - qml.Rot(x[0], x[1], x[2], wires=0) - return qml.expval(qml.PauliZ(0)) - - res = circuit(p) - - expected = jnp.cos(y) ** 2 - jnp.sin(x) * jnp.sin(y) ** 2 - assert jnp.allclose(res, expected, atol=tol, rtol=0) - - grad_fn = jacobian_transform(circuit, 0) - res = grad_fn(p) - - expected = jnp.array( - [-jnp.cos(x) * jnp.sin(y) ** 2, -2 * (jnp.sin(x) + 1) * jnp.sin(y) * jnp.cos(y), 0] - ) - assert jnp.allclose(res, expected, atol=tol, rtol=0) - - def test_jacobian_repeated_broadcasted(self, tol): - """Test that jacobian of a QNode with an attached default.qubit.jax device - gives the correct result in the case of repeated broadcasted parameters""" - p = jnp.array([[0.433, 92.1, -0.512], [0.218, 0.241, -0.51], [0.71, 0.152, 9.12]]) - dev = qml.device("default.qubit.jax", wires=1) - - @qml.qnode(dev, interface="jax", diff_method="backprop") - def circuit(x): - qml.RX(x[1], wires=0) - qml.Rot(x[0], x[1], x[2], wires=0) - return qml.expval(qml.PauliZ(0)) - - res = circuit(p) - - x, y, _ = p - expected = jnp.cos(y) ** 2 - jnp.sin(x) * jnp.sin(y) ** 2 - assert jnp.allclose(res, expected, atol=tol, rtol=0) - - grad_fn = jax.jacobian(circuit) - res = grad_fn(p) - - expected = jnp.array( - [ - -jnp.cos(x) * jnp.sin(y) ** 2, - -2 * (jnp.sin(x) + 1) * jnp.sin(y) * jnp.cos(y), - jnp.zeros_like(x), - ] - ) - assert all(jnp.allclose(res[i, :, i], expected[:, i], atol=tol, rtol=0) for i in range(3)) - - @pytest.mark.parametrize("wires", [[0], ["abc"]]) - def test_state_differentiability(self, wires, tol, benchmark): - """Test that the device state can be differentiated""" - dev = qml.device("default.qubit.jax", wires=wires) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(a): - qml.RY(a, wires=wires[0]) - return qml.state() - - a = jnp.array(0.54) - - def cost(a): - """A function of the device quantum state, as a function - of input QNode parameters.""" - res = jnp.abs(circuit(a)) ** 2 - return res[1] - res[0] - - grad = benchmark(jax.grad(cost), a) - expected = jnp.sin(a) - assert jnp.allclose(grad, expected, atol=tol, rtol=0) - - def test_state_differentiability_broadcasted(self, tol): - """Test that the broadcasted device state can be differentiated""" - dev = qml.device("default.qubit.jax", wires=1) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(a): - qml.RY(a, wires=0) - return qml.expval(qml.PauliZ(0)) - - a = jnp.array([0.54, 0.32, 1.2]) - - def cost(a): - """A function of the device quantum state, as a function - of input QNode parameters.""" - circuit(a) - res = jnp.abs(dev.state) ** 2 - return res[:, 1] - res[:, 0] - - jac = jax.jacobian(cost)(a) - expected = jnp.diag(jnp.sin(a)) - assert jnp.allclose(jac, expected, atol=tol, rtol=0) - - @pytest.mark.parametrize("theta", np.linspace(-2 * np.pi, np.pi, 7)) - def test_CRot_gradient(self, theta, tol, benchmark): - """Tests that the automatic gradient of a arbitrary controlled Euler-angle-parameterized - gate is correct.""" - dev = qml.device("default.qubit.jax", wires=2) - a, b, c = np.array([theta, theta**3, np.sqrt(2) * theta]) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(a, b, c): - qml.StatePrep(np.array([1.0, -1.0]) / np.sqrt(2), wires=0) - qml.CRot(a, b, c, wires=[0, 1]) - return qml.expval(qml.PauliX(0)) - - def workload(): - return circuit(a, b, c), jax.grad(circuit, argnums=(0, 1, 2))(a, b, c) - - res, grad = benchmark(workload) - - expected = -np.cos(b / 2) * np.cos(0.5 * (a + c)) - assert np.allclose(res, expected, atol=tol, rtol=0) - - expected = np.array( - [ - [ - 0.5 * np.cos(b / 2) * np.sin(0.5 * (a + c)), - 0.5 * np.sin(b / 2) * np.cos(0.5 * (a + c)), - 0.5 * np.cos(b / 2) * np.sin(0.5 * (a + c)), - ] - ] - ) - assert np.allclose(grad, expected, atol=tol, rtol=0) - - def test_prob_differentiability(self, tol, benchmark): - """Test that the device probability can be differentiated""" - dev = qml.device("default.qubit.jax", wires=2) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(a, b): - qml.RX(a, wires=0) - qml.RY(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[1]) - - a = jnp.array(0.54) - b = jnp.array(0.12) - - def cost(a, b): - prob_wire_1 = circuit(a, b).squeeze() - return prob_wire_1[1] - prob_wire_1[0] - - def workload(): - return cost(a, b), jax.jit(jax.grad(cost, argnums=(0, 1)))(a, b) - - res, grad = benchmark(workload) - - expected = -jnp.cos(a) * jnp.cos(b) - assert jnp.allclose(res, expected, atol=tol, rtol=0) - - expected = [jnp.sin(a) * jnp.cos(b), jnp.cos(a) * jnp.sin(b)] - assert jnp.allclose(jnp.array(grad), jnp.array(expected), atol=tol, rtol=0) - - def test_prob_differentiability_broadcasted(self, tol): - """Test that the broadcasted device probability can be differentiated""" - dev = qml.device("default.qubit.jax", wires=2) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(a, b): - qml.RX(a, wires=0) - qml.RY(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[1]) - - a = jnp.array([0.54, 0.32, 1.2]) - b = jnp.array(0.12) - - def cost(a, b): - prob_wire_1 = circuit(a, b) - return prob_wire_1[:, 1] - prob_wire_1[:, 0] - - res = cost(a, b) - expected = -jnp.cos(a) * jnp.cos(b) - assert jnp.allclose(res, expected, atol=tol, rtol=0) - - jac = jax.jacobian(cost, argnums=[0, 1])(a, b) - expected = jnp.array([jnp.sin(a) * jnp.cos(b), jnp.cos(a) * jnp.sin(b)]) - expected = (jnp.diag(expected[0]), expected[1]) # Only first parameter is broadcasted - assert all(jnp.allclose(j, e, atol=tol, rtol=0) for j, e in zip(jac, expected)) - - def test_backprop_gradient(self, tol): - """Tests that the gradient of the qnode is correct""" - dev = qml.device("default.qubit.jax", wires=2) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(a, b): - qml.RX(a, wires=0) - qml.CRX(b, wires=[0, 1]) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - a = jnp.array(-0.234) - b = jnp.array(0.654) - - res = circuit(a, b) - expected_cost = 0.5 * (jnp.cos(a) * jnp.cos(b) + jnp.cos(a) - jnp.cos(b) + 1) - assert jnp.allclose(res, expected_cost, atol=tol, rtol=0) - res = jax.grad(circuit, argnums=(0, 1))(a, b) - expected_grad = jnp.array( - [-0.5 * jnp.sin(a) * (jnp.cos(b) + 1), 0.5 * jnp.sin(b) * (1 - jnp.cos(a))] - ) - - assert jnp.allclose(jnp.array(res), jnp.array(expected_grad), atol=tol, rtol=0) - - def test_backprop_gradient_broadcasted(self, tol): - """Tests that the gradient of the broadcasted qnode is correct""" - dev = qml.device("default.qubit.jax", wires=2) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(a, b): - qml.RX(a, wires=0) - qml.CRX(b, wires=[0, 1]) - return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1)) - - a = jnp.array(0.12) - b = jnp.array([0.54, 0.32, 1.2]) - - res = circuit(a, b) - expected_cost = 0.5 * (jnp.cos(a) * jnp.cos(b) + jnp.cos(a) - jnp.cos(b) + 1) - assert jnp.allclose(res, expected_cost, atol=tol, rtol=0) - - res = jax.jacobian(circuit, argnums=[0, 1])(a, b) - expected = jnp.array( - [-0.5 * jnp.sin(a) * (jnp.cos(b) + 1), 0.5 * jnp.sin(b) * (1 - jnp.cos(a))] - ) - expected = (expected[0], jnp.diag(expected[1])) - assert all(jnp.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected)) - - @pytest.mark.parametrize("x, shift", [(0.0, 0.0), (0.5, -0.5)]) - def test_hessian_at_zero(self, x, shift): - """Tests that the Hessian at vanishing state vector amplitudes - is correct.""" - dev = qml.device("default.qubit.jax", wires=1) - - @qml.qnode(dev, interface="jax", diff_method="backprop") - def circuit(x): - qml.RY(shift, wires=0) - qml.RY(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - assert qml.math.isclose(jax.grad(circuit)(x), 0.0) - assert qml.math.isclose(jax.jacobian(jax.jacobian(circuit))(x), -1.0) - assert qml.math.isclose(jax.grad(jax.grad(circuit))(x), -1.0) - - @pytest.mark.parametrize("operation", [qml.U3, qml.U3.compute_decomposition]) - @pytest.mark.parametrize("diff_method", ["backprop"]) - def test_jax_interface_gradient(self, operation, diff_method, tol, benchmark): - """Tests that the gradient of an arbitrary U3 gate is correct - using the Jax interface, using a variety of differentiation methods.""" - dev = qml.device("default.qubit.jax", wires=1) - - @qml.qnode(dev, diff_method=diff_method, interface="jax") - def circuit(x, weights, w=None): - """In this example, a mixture of scalar - arguments, array arguments, and keyword arguments are used.""" - qml.StatePrep(1j * jnp.array([1, -1]) / jnp.sqrt(2), wires=w) - operation(x, weights[0], weights[1], wires=w) - return qml.expval(qml.PauliX(w)) - - def cost(params): - """Perform some classical processing""" - return (circuit(params[0], params[1:], w=0) ** 2).reshape(()) - - theta = 0.543 - phi = -0.234 - lam = 0.654 - - params = jnp.array([theta, phi, lam]) - - def workload(): - return cost(params), jax.grad(cost)(params) - - res, grad = benchmark(workload) - - expected_cost = ( - jnp.sin(lam) * jnp.sin(phi) - jnp.cos(theta) * jnp.cos(lam) * jnp.cos(phi) - ) ** 2 - assert jnp.allclose(res, expected_cost, atol=tol, rtol=0) - - expected_grad = ( - jnp.array( - [ - jnp.sin(theta) * jnp.cos(lam) * jnp.cos(phi), - jnp.cos(theta) * jnp.cos(lam) * jnp.sin(phi) + jnp.sin(lam) * jnp.cos(phi), - jnp.cos(theta) * jnp.sin(lam) * jnp.cos(phi) + jnp.cos(lam) * jnp.sin(phi), - ] - ) - * 2 - * (jnp.sin(lam) * jnp.sin(phi) - jnp.cos(theta) * jnp.cos(lam) * jnp.cos(phi)) - ) - assert jnp.allclose(grad, expected_grad, atol=tol, rtol=0) - - @pytest.mark.xfail(reason="Not applicable anymore.") - @pytest.mark.parametrize("interface", ["autograd", "tf", "torch"]) - def test_error_backprop_wrong_interface(self, interface): - """Tests that an error is raised if diff_method='backprop' but not using - the Jax interface""" - dev = qml.device("default.qubit.jax", wires=1) - - def circuit(x, w=None): - qml.RZ(x, wires=w) - return qml.expval(qml.PauliX(w)) - - error_type = qml.QuantumFunctionError - with pytest.raises( - error_type, - match="default.qubit.jax only supports diff_method='backprop' when using the jax interface", - ): - qml.qnode(dev, diff_method="backprop", interface=interface)(circuit) - - def test_no_jax_interface_applied(self): - """Tests that the JAX interface is not applied and no error is raised if qml.probs is used with the Jax - interface when diff_method='backprop' - - When the JAX interface is applied, we can only get the expectation value and the variance of a QNode. - """ - dev = qml.device("default.qubit.jax", wires=1, shots=None) - - def circuit(): - return qml.probs(wires=0) - - qnode = qml.qnode(dev, diff_method="backprop", interface="jax")(circuit) - assert jnp.allclose(qnode(), jnp.array([1, 0])) - - -@pytest.mark.jax -class TestHighLevelIntegration: - """Tests for integration with higher level components of PennyLane.""" - - def test_do_not_split_analytic_jax(self, mocker): - """Tests that the Hamiltonian is not split for shots=None using the jax device.""" - dev = qml.device("default.qubit.jax", wires=2) - H = qml.Hamiltonian(jnp.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)]) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(): - return qml.expval(H) - - spy = mocker.spy(dev.target_device, "expval") - - circuit() - # evaluated one expval altogether - assert spy.call_count == 1 - - def test_direct_eval_linear_combination_broadcasted_jax(self): - """Tests that the correct result is returned when attempting to evaluate a Hamiltonian with - broadcasting and shots=None directly via its sparse representation with Jax.""" - dev = qml.device("default.qubit.jax", wires=2) - H = qml.ops.LinearCombination(jnp.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)]) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(): - qml.RX(jnp.zeros(5), 0) - return qml.expval(H) - - res = circuit() - assert qml.math.allclose(res, 0.2) - - @pytest.mark.usefixtures("use_legacy_opmath") - def test_direct_eval_hamiltonian_broadcasted_error_jax_legacy_opmath(self): - """Tests that an error is raised when attempting to evaluate a Hamiltonian with - broadcasting and shots=None directly via its sparse representation with Jax.""" - dev = qml.device("default.qubit.jax", wires=2) - H = qml.Hamiltonian(jnp.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)]) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(): - qml.RX(jnp.zeros(5), 0) - return qml.expval(H) - - with pytest.raises(NotImplementedError, match="Hamiltonians for interface!=None"): - circuit() - - def test_template_integration(self): - """Test that a PassthruQNode using default.qubit.jax works with templates.""" - dev = qml.device("default.qubit.jax", wires=2) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(weights): - qml.templates.StronglyEntanglingLayers(weights, wires=[0, 1]) - return qml.expval(qml.PauliZ(0)) - - weights = jnp.array( - np.random.random(qml.templates.StronglyEntanglingLayers.shape(n_layers=2, n_wires=2)) - ) - - grad = jax.grad(circuit)(weights) - assert grad.shape == weights.shape - - -# pylint: disable=protected-access -@pytest.mark.jax -class TestOps: - """Unit tests for operations supported by the default.qubit.jax device""" - - @pytest.mark.parametrize("jacobian_transform", [jax.jacfwd, jax.jacrev]) - def test_multirz_jacobian(self, jacobian_transform, benchmark): - """Test that the patched numpy functions are used for the MultiRZ - operation and the jacobian can be computed.""" - wires = 4 - dev = qml.device("default.qubit.jax", wires=wires) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(param): - qml.MultiRZ(param, wires=[0, 1]) - return qml.probs(wires=list(range(wires))) - - param = 0.3 - res = benchmark(jacobian_transform(circuit), param) - assert jnp.allclose(res, jnp.zeros(wires**2)) - - def test_full_subsystem(self, mocker): - """Test applying a state vector to the full subsystem""" - dev = DefaultQubitJax(wires=["a", "b", "c"]) - state = jnp.array([1, 0, 0, 0, 1, 0, 1, 1]) / 2.0 - state_wires = qml.wires.Wires(["a", "b", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - - assert jnp.all(dev._state.flatten() == state) - spy.assert_not_called() - - def test_partial_subsystem(self, mocker): - """Test applying a state vector to a subset of wires of the full subsystem""" - - dev = DefaultQubitJax(wires=["a", "b", "c"]) - state = jnp.array([1, 0, 1, 0]) / jnp.sqrt(2.0) - state_wires = qml.wires.Wires(["a", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - res = jnp.sum(dev._state, axis=(1,)).flatten() - - assert jnp.all(res == state) - spy.assert_called() - - def test_parametrized_evolution(self): - """Test applying a ParametrizedEvolution to a subset of wires of the full subsystem""" - dev = DefaultQubitJax(wires=["a", "b", "c"]) - state = jnp.array([[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]], dtype=complex) - expected_res = jnp.array( - [[[0.0, 0.0], [0.0, 0.0]], [[1.0, 0.0], [0.0, 0.0]]], dtype=complex - ) - # ev corresponds to a PauliX gate - ev = qml.evolve(ParametrizedHamiltonian([1], [qml.PauliX("a")]))(params=[], t=np.pi / 2) - res = qml.math.abs(dev._apply_parametrized_evolution(state=state, operation=ev)) - - assert qml.math.allclose(res, expected_res, atol=1e-5) - - def test_parametrized_evolution_raises_error(self): - """Test applying a ParametrizedEvolution without params or t specified raises an error.""" - dev = DefaultQubitJax(wires=["a"]) - state = jnp.array([[[1.0, 0.0], [0.0, 0.0]], [[0.0, 0.0], [0.0, 0.0]]], dtype=complex) - ev = qml.evolve(ParametrizedHamiltonian([1], [qml.PauliX("a")])) - with pytest.raises( - ValueError, - match="The parameters and the time window are required to execute a ParametrizedEvolution", - ): - dev._apply_parametrized_evolution(state=state, operation=ev) - - -@pytest.mark.jax -class TestOpsBroadcasted: - """Unit tests for broadcasted operations supported by the default.qubit.jax device""" - - @pytest.mark.parametrize("jacobian_transform", [jax.jacfwd, jax.jacrev]) - def test_multirz_jacobian_broadcasted(self, jacobian_transform): - """Test that the patched numpy functions are used for the MultiRZ - operation and the jacobian can be computed.""" - wires = 4 - dev = qml.device("default.qubit.jax", wires=wires) - - @qml.qnode(dev, diff_method="backprop", interface="jax") - def circuit(param): - qml.MultiRZ(param, wires=[0, 1]) - return qml.probs(wires=list(range(wires))) - - param = jnp.array([0.3, 0.9, -4.3]) - res = jacobian_transform(circuit)(param) - assert jnp.allclose(res, jnp.zeros((3, wires**2, 3))) - - def test_full_subsystem_broadcasted(self, mocker): - """Test applying a state vector to the full subsystem""" - dev = DefaultQubitJax(wires=["a", "b", "c"]) - state = jnp.array([[1, 0, 0, 0, 1, 0, 1, 1], [0, 0, 0, 1, 1, 1, 1, 0]]) / 2.0 - state_wires = qml.wires.Wires(["a", "b", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - - assert jnp.all(dev._state.reshape((2, 8)) == state) - spy.assert_not_called() - - def test_partial_subsystem_broadcasted(self, mocker): - """Test applying a state vector to a subset of wires of the full subsystem""" - - dev = DefaultQubitJax(wires=["a", "b", "c"]) - state = jnp.array([[1, 0, 1, 0], [0, 1, 0, 1], [1, 1, 0, 0]]) / jnp.sqrt(2.0) - state_wires = qml.wires.Wires(["a", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - res = jnp.sum(dev._state, axis=(2,)).reshape((3, 4)) - - assert jnp.allclose(res, state) - spy.assert_called() - - -@pytest.mark.jax -class TestEstimateProb: - """Test the estimate_probability method""" - - @pytest.mark.parametrize( - "wires, expected", [([0], [0.5, 0.5]), (None, [0.5, 0, 0, 0.5]), ([0, 1], [0.5, 0, 0, 0.5])] - ) - def test_estimate_probability(self, wires, expected, monkeypatch): - """Tests the estimate_probability method""" - dev = qml.device("default.qubit.jax", wires=2) - samples = jnp.array([[0, 0], [1, 1], [1, 1], [0, 0]]) - - with monkeypatch.context() as m: - m.setattr(dev.target_device, "_samples", samples) - res = dev.estimate_probability(wires=wires) - - assert np.allclose(res, expected) - - @pytest.mark.parametrize( - "wires, expected", - [ - ([0], [[0.0, 0.5], [1.0, 0.5]]), - (None, [[0.0, 0.5], [0, 0], [0, 0.5], [1.0, 0]]), - ([0, 1], [[0.0, 0.5], [0, 0], [0, 0.5], [1.0, 0]]), - ], - ) - def test_estimate_probability_with_binsize(self, wires, expected, monkeypatch): - """Tests the estimate_probability method with a bin size""" - dev = qml.device("default.qubit.jax", wires=2) - samples = jnp.array([[1, 1], [1, 1], [1, 0], [0, 0]]) - bin_size = 2 - - with monkeypatch.context() as m: - m.setattr(dev.target_device, "_samples", samples) - res = dev.estimate_probability(wires=wires, bin_size=bin_size) - - assert np.allclose(res, expected) - - @pytest.mark.parametrize( - "wires, expected", - [ - ([0], [[0.0, 1.0], [0.5, 0.5], [0.25, 0.75]]), - (None, [[0, 0, 0.25, 0.75], [0.5, 0, 0, 0.5], [0.25, 0, 0.25, 0.5]]), - ([0, 1], [[0, 0, 0.25, 0.75], [0.5, 0, 0, 0.5], [0.25, 0, 0.25, 0.5]]), - ], - ) - def test_estimate_probability_with_broadcasting(self, wires, expected, monkeypatch): - """Tests the estimate_probability method with parameter broadcasting""" - dev = qml.device("default.qubit.jax", wires=2) - samples = jnp.array( - [ - [[1, 0], [1, 1], [1, 1], [1, 1]], - [[0, 0], [1, 1], [1, 1], [0, 0]], - [[1, 0], [1, 1], [1, 1], [0, 0]], - ] - ) - - with monkeypatch.context() as m: - m.setattr(dev.target_device, "_samples", samples) - res = dev.estimate_probability(wires=wires) - - assert np.allclose(res, expected) - - @pytest.mark.parametrize( - "wires, expected", - [ - ( - [0], - [ - [[0, 0, 0.5], [1, 1, 0.5]], - [[0.5, 0.5, 0], [0.5, 0.5, 1]], - [[0, 0.5, 1], [1, 0.5, 0]], - ], - ), - ( - None, - [ - [[0, 0, 0], [0, 0, 0.5], [0.5, 0, 0], [0.5, 1, 0.5]], - [[0.5, 0.5, 0], [0, 0, 0], [0, 0, 0], [0.5, 0.5, 1]], - [[0, 0.5, 0.5], [0, 0, 0.5], [0.5, 0, 0], [0.5, 0.5, 0]], - ], - ), - ( - [0, 1], - [ - [[0, 0, 0], [0, 0, 0.5], [0.5, 0, 0], [0.5, 1, 0.5]], - [[0.5, 0.5, 0], [0, 0, 0], [0, 0, 0], [0.5, 0.5, 1]], - [[0, 0.5, 0.5], [0, 0, 0.5], [0.5, 0, 0], [0.5, 0.5, 0]], - ], - ), - ], - ) - def test_estimate_probability_with_binsize_with_broadcasting( - self, wires, expected, monkeypatch - ): - """Tests the estimate_probability method with a bin size and parameter broadcasting""" - dev = qml.device("default.qubit.jax", wires=2) - bin_size = 2 - samples = jnp.array( - [ - [[1, 0], [1, 1], [1, 1], [1, 1], [1, 1], [0, 1]], - [[0, 0], [1, 1], [1, 1], [0, 0], [1, 1], [1, 1]], - [[1, 0], [1, 1], [1, 1], [0, 0], [0, 1], [0, 0]], - ] - ) - - with monkeypatch.context() as m: - m.setattr(dev.target_device, "_samples", samples) - res = dev.estimate_probability(wires=wires, bin_size=bin_size) - - assert np.allclose(res, expected) diff --git a/tests/devices/test_default_qubit_legacy.py b/tests/devices/test_default_qubit_legacy.py index 2bcaaef11aa..11ca082441c 100644 --- a/tests/devices/test_default_qubit_legacy.py +++ b/tests/devices/test_default_qubit_legacy.py @@ -24,7 +24,6 @@ import pennylane as qml from pennylane import numpy as np from pennylane.devices.default_qubit_legacy import DefaultQubitLegacy, _get_slice -from pennylane.pulse import ParametrizedHamiltonian from pennylane.wires import WireError, Wires U = np.array( @@ -1007,10 +1006,7 @@ def test_defines_correct_capabilities(self): "supports_inverse_operations": True, "supports_analytic_computation": True, "supports_broadcasting": True, - "passthru_devices": { - "autograd": "default.qubit.autograd", - "jax": "default.qubit.jax", - }, + "passthru_devices": {}, } assert cap == capabilities @@ -2085,34 +2081,6 @@ def test_apply_three_qubit_op_controls_split(self, op, method): state_out_einsum = np.einsum("abcdef,kdfe->kacb", matrix, self.state) assert np.allclose(state_out, state_out_einsum) - @pytest.mark.jax - def test_apply_parametrized_evolution_raises_error(self): - """Test that applying a ParametrizedEvolution raises an error.""" - param_ev = qml.evolve(ParametrizedHamiltonian([1], [qml.PauliX(0)])) - with pytest.raises( - NotImplementedError, - match="The device default.qubit.legacy cannot execute a ParametrizedEvolution operation", - ): - self.dev._apply_parametrized_evolution(state=self.state, operation=param_ev) - - @qml.qnode(self.dev) - def circuit(): - qml.apply(param_ev) - return qml.expval(qml.PauliZ(0)) - - with pytest.raises( - qml.DeviceError, - match="Gate ParametrizedEvolution not supported on device default.qubit.", - ): - circuit() - - self.dev.operations.add("ParametrizedEvolution") - with pytest.raises( - NotImplementedError, - match="The device default.qubit.legacy cannot execute a ParametrizedEvolution operation", - ): - circuit() - class TestStateVector: """Unit tests for the _apply_state_vector method""" @@ -2376,7 +2344,7 @@ def test_Hamiltonian_filtered_from_rotations(self, mocker): dev = qml.device("default.qubit.legacy", wires=2, shots=10) H = qml.Hamiltonian([0.1, 0.2], [qml.PauliX(0), qml.PauliZ(1)]) - spy = mocker.spy(qml.QubitDevice, "_get_diagonalizing_gates") + spy = mocker.spy(qml.devices.QubitDevice, "_get_diagonalizing_gates") qs = qml.tape.QuantumScript([qml.RX(1, 0)], [qml.expval(qml.PauliX(0)), qml.expval(H)]) rotations = dev._get_diagonalizing_gates(qs) @@ -2414,35 +2382,11 @@ def circuit(y, z, is_state_batched): def test_super_expval_not_called(self, is_state_batched, mocker): """Tests basic expval result, and ensures QubitDevice.expval is not called.""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.QubitDevice, "expval") + spy = mocker.spy(qml.devices.QubitDevice, "expval") obs = qml.sum(qml.s_prod(0.1, qml.PauliX(0)), qml.s_prod(0.2, qml.PauliZ(0))) assert np.isclose(dev.expval(obs), 0.2) spy.assert_not_called() - @pytest.mark.autograd - def test_trainable_autograd(self, is_state_batched): - """Tests that coeffs passed to a sum are trainable with autograd.""" - if is_state_batched: - pytest.skip( - reason="Broadcasting, qml.jacobian and new return types do not work together" - ) - dev = qml.device("default.qubit.legacy", wires=1) - qnode = qml.QNode(self.circuit, dev, interface="autograd") - y, z = np.array([1.1, 2.2]) - actual = qml.grad(qnode, argnum=[0, 1])(y, z, is_state_batched) - assert np.allclose(actual, self.expected_grad(is_state_batched)) - - @pytest.mark.jax - def test_trainable_jax(self, is_state_batched): - """Tests that coeffs passed to a sum are trainable with jax.""" - import jax - - dev = qml.device("default.qubit.legacy", wires=1) - qnode = qml.QNode(self.circuit, dev, interface="jax") - y, z = jax.numpy.array([1.1, 2.2]) - actual = jax.jacobian(qnode, argnums=[0, 1])(y, z, is_state_batched) - assert np.allclose(actual, self.expected_grad(is_state_batched)) - class TestGetBatchSize: """Tests for the helper method ``_get_batch_size`` of ``QubitDevice``.""" diff --git a/tests/devices/test_default_qubit_legacy_broadcasting.py b/tests/devices/test_default_qubit_legacy_broadcasting.py deleted file mode 100644 index 50ef31d0294..00000000000 --- a/tests/devices/test_default_qubit_legacy_broadcasting.py +++ /dev/null @@ -1,2024 +0,0 @@ -# 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. -# 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 :class:`pennylane.devices.DefaultQubitLegacy` device when using broadcasting. -""" -# pylint: disable=protected-access,cell-var-from-loop,too-many-arguments -import math -from itertools import product - -import pytest -from gate_data import ( - CNOT, - CSWAP, - CZ, - ISWAP, - SISWAP, - SWAP, - CRot3, - CRotx, - CRoty, - CRotz, - H, - I, - IsingXX, - IsingYY, - IsingZZ, - MultiRZ1, - MultiRZ2, - Rot3, - Rotx, - Roty, - Rotz, - Rphi, - S, - T, - Toffoli, - X, - Y, - Z, -) - -import pennylane as qml -from pennylane import DeviceError -from pennylane import numpy as np -from pennylane.devices.default_qubit_legacy import DefaultQubitLegacy - -THETA = np.linspace(0.11, 1, 3) -PHI = np.linspace(0.32, 1, 3) -VARPHI = np.linspace(0.02, 1, 3) - -INVSQ2 = 1 / math.sqrt(2) -T_PHASE = np.exp(1j * np.pi / 4) -T_PHASE_C = np.exp(-1j * np.pi / 4) - -# Variant of diag that does not take the diagonal of a 2d array, but broadcasts diag. -diag = lambda x: np.array([np.diag(_x) for _x in x]) if np.ndim(x) == 2 else np.diag(x) - - -def mat_vec(mat, vec, par=None, inv=False): - if par is not None: - scalar = [np.isscalar(p) for p in par] - if not all(scalar): - batch_size = len(par[scalar.index(False)]) - par = [tuple(p if s else p[i] for p, s in zip(par, scalar)) for i in range(batch_size)] - mat = np.array([mat(*_par) for _par in par]) - else: - mat = mat(*par) - - if inv: - mat = np.moveaxis(mat.conj(), -2, -1) - - return np.einsum("...ij,...j->...i", mat, vec) - - -class TestApplyBroadcasted: - """Tests that operations and inverses of certain operations are applied to a broadcasted - state/with broadcasted parameters (or both) correctly, or that the proper errors are raised. - """ - - triple_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]]) - test_data_no_parameters = [ - (qml.PauliX, triple_state, mat_vec(X, triple_state)), - (qml.PauliY, triple_state, mat_vec(Y, triple_state)), - (qml.PauliZ, triple_state, mat_vec(Z, triple_state)), - (qml.S, triple_state, mat_vec(S, triple_state)), - (qml.T, triple_state, mat_vec(T, triple_state)), - (qml.Hadamard, triple_state, mat_vec(H, triple_state)), - (qml.Identity, triple_state, triple_state), - ] - - @pytest.mark.parametrize("operation,input,expected_output", test_data_no_parameters) - def test_apply_operation_single_wire_no_parameters_broadcasted( - self, qubit_device_1_wire, tol, operation, input, expected_output - ): - """Tests that applying an operation yields the expected output state for single wire - operations that have no parameters.""" - - qubit_device_1_wire.target_device._state = np.array( - input, dtype=qubit_device_1_wire.C_DTYPE - ) - qubit_device_1_wire.apply([operation(wires=[0])]) - - assert np.allclose(qubit_device_1_wire._state, np.array(expected_output), atol=tol, rtol=0) - assert qubit_device_1_wire._state.dtype == qubit_device_1_wire.C_DTYPE - - single_state = np.array([[0, 0.6, 0, 0.8]]) - triple_state = np.array([[1, 0, 0, 0], [0, 0, INVSQ2, -INVSQ2], [0, 0.6, 0, 0.8]]) - - test_data_two_wires_no_param = [ - (qml_op, state, mat_vec(mat_op, state)) - for (qml_op, mat_op), state in product( - zip( - [qml.CNOT, qml.SWAP, qml.CZ, qml.ISWAP, qml.SISWAP, qml.SQISW], - [CNOT, SWAP, CZ, ISWAP, SISWAP, SISWAP], - ), - [single_state, triple_state], - ) - ] - - @pytest.mark.parametrize("operation,input,expected_output", test_data_two_wires_no_param) - def test_apply_operation_two_wires_no_parameters_broadcasted( - self, qubit_device_2_wires, tol, operation, input, expected_output - ): - """Tests that applying an operation yields the expected output state for two wire - operations that have no parameters.""" - - qubit_device_2_wires.target_device._state = np.array( - input, dtype=qubit_device_2_wires.C_DTYPE - ).reshape((-1, 2, 2)) - qubit_device_2_wires.apply([operation(wires=[0, 1])]) - - assert np.allclose( - qubit_device_2_wires._state.reshape((-1, 4)), - np.array(expected_output), - atol=tol, - rtol=0, - ) - assert qubit_device_2_wires._state.dtype == qubit_device_2_wires.C_DTYPE - - quad_state = np.array( - [ - [0.6, 0, 0, 0, 0, 0, 0.8, 0], - [-INVSQ2, INVSQ2, 0, 0, 0, 0, 0, 0], - [0, 0, 0.5, 0.5, 0.5, -0.5, 0, 0], - [0, 0, 0.5, 0, 0.5, -0.5, 0, 0.5], - ] - ) - test_data_three_wires_no_parameters = [(qml.CSWAP, quad_state, mat_vec(CSWAP, quad_state))] - - @pytest.mark.parametrize("operation,input,expected_output", test_data_three_wires_no_parameters) - def test_apply_operation_three_wires_no_parameters_broadcasted( - self, qubit_device_3_wires, tol, operation, input, expected_output - ): - """Tests that applying an operation yields the expected output state for three wire - operations that have no parameters.""" - - qubit_device_3_wires.target_device._state = np.array( - input, dtype=qubit_device_3_wires.C_DTYPE - ).reshape((-1, 2, 2, 2)) - qubit_device_3_wires.apply([operation(wires=[0, 1, 2])]) - - assert np.allclose( - qubit_device_3_wires._state.reshape((-1, 8)), - np.array(expected_output), - atol=tol, - rtol=0, - ) - assert qubit_device_3_wires._state.dtype == qubit_device_3_wires.C_DTYPE - - single_state = np.array([[0, 0, 1, 0]]) - triple_state = np.array( - [ - [0, 0, 1, 0], - [1 / math.sqrt(3), 0, 1 / math.sqrt(3), 1 / math.sqrt(3)], - [0.5, -0.5, 0.5j, -0.5j], - ] - ) - - # TODO[dwierichs]: add tests with qml.BaisState once `_apply_basis_state` supports broadcasting - @pytest.mark.parametrize( - "operation,expected_output,par", - [(qml.StatePrep, s, s) for s in [single_state, triple_state]], - ) - def test_apply_operation_state_preparation_broadcasted( - self, qubit_device_2_wires, tol, operation, expected_output, par - ): - """Tests that applying an operation yields the expected output state for single wire - operations that have no parameters.""" - - par = np.array(par) - qubit_device_2_wires.reset() - qubit_device_2_wires.apply([operation(par, wires=[0, 1])]) - - assert np.allclose( - qubit_device_2_wires._state.reshape((-1, 4)), - np.array(expected_output), - atol=tol, - rtol=0, - ) - - # Collect test cases for single-scalar-parameter single-wire operations and their inverses - # For each operation, we choose broadcasted state, broadcasted params, or both - state_1 = np.array([0.6, 0.8j]) - state_5 = np.array([[INVSQ2, INVSQ2], [0.6, 0.8], [0, 1j], [-1, 0], [-INVSQ2, INVSQ2]]) - scalar_par_1 = [np.pi / 2] - scalar_par_5 = [[np.pi / 3, np.pi, 0.5, -1.2, -3 * np.pi / 2]] - test_data_single_wire_with_parameters = [ - (qml_op, state, mat_vec(mat_op, state, par=par), par) - for (qml_op, mat_op), (state, par) in product( - zip( - [qml.PhaseShift, qml.RX, qml.RY, qml.RZ, qml.MultiRZ], - [Rphi, Rotx, Roty, Rotz, MultiRZ1], - ), - [(state_1, scalar_par_5), (state_5, scalar_par_1), (state_5, scalar_par_5)], - ) - ] - - # Add qml.QubitUnitary test cases - matrix_1_par_1 = [np.array([[1, 1j], [1j, 1]]) * INVSQ2] - matrix_1_par_5 = [ - np.array( - [ - np.array([[1, -1j], [-1j, 1]]) * INVSQ2, - np.array([[1, -1], [1, 1]]) * INVSQ2, - np.array([[T_PHASE_C, 0], [0, T_PHASE]]), - np.array([[1, 0], [0, -1]]), - T, - ] - ) - ] - test_data_single_wire_with_parameters += [ - (qml.QubitUnitary, s, mat_vec(par[0], s), par) - for s, par in [ - (state_1, matrix_1_par_5), - (state_5, matrix_1_par_1), - (state_5, matrix_1_par_5), - ] - ] - - # Add qml.DiagonalQubitUnitary test cases - diag_par_1 = [[np.exp(1j * 0.1), np.exp(1j * np.pi)]] - diag_par_5 = [ - np.array( - [ - [1, -1j], - [np.exp(1j * 0.4), np.exp(1j * -0.4)], - [np.exp(1j * 0.1), np.exp(1j * np.pi)], - [1.0, np.exp(1j * np.pi / 2)], - [1, 1], - ] - ) - ] - test_data_single_wire_with_parameters += [ - (qml.DiagonalQubitUnitary, s, mat_vec(diag(par[0]), s), par) - for s, par in [(state_1, diag_par_5), (state_5, diag_par_1), (state_5, diag_par_5)] - ] - - # Add qml.SpecialUnitary test cases - theta_1_par_1 = [np.array([np.pi / 2, 0, 0])] - theta_1_par_5 = [ - np.array( - [[np.pi / 2, 0, 0], [0, np.pi / 2, 0], [0, 0, np.pi / 2], [0.3, 0, 0], [0.4, 0.2, 1.2]] - ) - ] - test_data_single_wire_with_parameters += [ - (qml.SpecialUnitary, s, mat_vec(qml.SpecialUnitary.compute_matrix(par[0], 1), s), par) - for s, par in [(state_1, theta_1_par_5), (state_5, theta_1_par_1), (state_5, theta_1_par_5)] - ] - - # Add qml.Rot test cases - multi_par_1 = { - "rz_0": [0.632, 0, 0], - "ry": [0, 0.632, 0], - "rz_1": [0, 0, 0.632], - "mixed": [0.12, -2.468, 0.812], - } - multi_par_5 = { - "rz_0": [[np.pi / 2 * i for i in range(5)], 0, 0], - "ry": [0, [np.pi / 2 * i for i in range(5)], 0], - "rz_1": [0, 0, [np.pi / 2 * i for i in range(5)]], - "mixed": [[np.pi / 2 * i for i in range(5)], [np.pi / 2 * i for i in range(5)], np.pi], - } - for like in ["rz_0", "ry", "rz_1", "mixed"]: - states_and_pars = [ - (state_1, multi_par_5[like]), - (state_5, multi_par_1[like]), - (state_5, multi_par_5[like]), - ] - test_data_single_wire_with_parameters += [ - (qml.Rot, s, mat_vec(Rot3, s, par=par), par) for s, par in states_and_pars - ] - - @pytest.mark.parametrize( - "operation,input,expected_output,par", test_data_single_wire_with_parameters - ) - def test_apply_operation_single_wire_with_parameters_broadcasted( - self, qubit_device_1_wire, tol, operation, input, expected_output, par - ): - """Tests that applying an operation yields the expected output state for single wire - operations that have parameters.""" - - qubit_device_1_wire.target_device._state = np.array( - input, dtype=qubit_device_1_wire.C_DTYPE - ) - - par = tuple(np.array(p) for p in par) - qubit_device_1_wire.apply([operation(*par, wires=[0])]) - - assert np.allclose(qubit_device_1_wire._state, np.array(expected_output), atol=tol, rtol=0) - assert qubit_device_1_wire._state.dtype == qubit_device_1_wire.C_DTYPE - - # Collect test cases for single-scalar-parameter two-wires operations and their inverses - # For each operation, we choose broadcasted state, broadcasted params, or both - state_1 = np.array([0.6, 0.8j, -0.6, -0.8j]) * INVSQ2 - state_5 = np.array( - [ - [0, 1, 0, 0], - [0, 0, 0, 1], - [0, INVSQ2, INVSQ2, 0], - [0.5, 0.5j, -0.5, 0.5], - [0.6, 0, -0.8j, 0], - ] - ) - scalar_par_1 = [np.pi / 2] - scalar_par_5 = [[np.pi / 3, np.pi, 0.5, -1.2, -3 * np.pi / 2]] - two_wires_scalar_par_ops = [ - qml.CRX, - qml.CRY, - qml.CRZ, - qml.MultiRZ, - qml.IsingXX, - qml.IsingYY, - qml.IsingZZ, - ] - two_wires_scalar_par_mats = [CRotx, CRoty, CRotz, MultiRZ2, IsingXX, IsingYY, IsingZZ] - test_data_two_wires_with_parameters = [ - (qml_op, state, mat_vec(mat_op, state, par=par), par) - for (qml_op, mat_op), (state, par) in product( - zip(two_wires_scalar_par_ops, two_wires_scalar_par_mats), - [(state_1, scalar_par_5), (state_5, scalar_par_1), (state_5, scalar_par_5)], - ) - ] - - # Add qml.CRot test cases - multi_par_1 = { - "rz_0": [0.632, 0, 0], - "ry": [0, 0.632, 0], - "rz_1": [0, 0, 0.632], - "mixed": [0.12, -2.468, 0.812], - } - multi_par_5 = { - "rz_0": [[np.pi / 2 * i for i in range(5)], 0, 0], - "ry": [0, [np.pi / 2 * i for i in range(5)], 0], - "rz_1": [0, 0, [np.pi / 2 * i for i in range(5)]], - "mixed": [[np.pi / 2 * i for i in range(5)], [np.pi / 2 * i for i in range(5)], np.pi], - } - for like in ["rz_0", "ry", "rz_1", "mixed"]: - states_and_pars = [ - (state_1, multi_par_5[like]), - (state_5, multi_par_1[like]), - (state_5, multi_par_5[like]), - ] - test_data_two_wires_with_parameters += [ - (qml.CRot, s, mat_vec(CRot3, s, par=par), par) for s, par in states_and_pars - ] - - # Add qml.QubitUnitary test cases - matrix_2_par_1 = [SISWAP] - matrix_2_par_5 = [ - np.array( - [ - np.eye(4), - np.array([[1, -1j, 0, 0], [-1j, 1, 0, 0], [0, 0, 1, -1j], [0, 0, -1j, 1]]) * INVSQ2, - np.array([[1, -1, 0, 0], [1, 1, 0, 0], [0, 0, 1, -1j], [0, 0, 1j, -1]]) * INVSQ2, - np.array([[T_PHASE_C, 0, 0, 0], [0, T_PHASE, 0, 0], [0, 0, 1, 0], [0, 0, 0, 1j]]), - SISWAP, - ] - ) - ] - test_data_two_wires_with_parameters += [ - (qml.QubitUnitary, s, mat_vec(par[0], s), par) - for s, par in [ - (state_1, matrix_2_par_5), - (state_5, matrix_2_par_1), - (state_5, matrix_2_par_5), - ] - ] - - # Add qml.DiagonalQubitUnitary test cases - diag_par_1 = [np.exp(1j * np.array([0.1, np.pi, 0.2, -2.4]))] - diag_par_5 = [ - np.array( - [ - np.ones(4), - [1, -1j, 1, 1j], - [np.exp(1j * 0.4), np.exp(1j * -0.4), 1j, 1], - [np.exp(1j * 0.1), np.exp(1j * np.pi), INVSQ2 * (1 + 1j), INVSQ2 * (1 - 1j)], - [1.0, np.exp(1j * np.pi / 2), 1, 1], - ] - ) - ] - test_data_two_wires_with_parameters += [ - (qml.DiagonalQubitUnitary, s, mat_vec(diag(par[0]), s), par) - for s, par in [(state_1, diag_par_5), (state_5, diag_par_1), (state_5, diag_par_5)] - ] - - # Add qml.SpecialUnitary test cases - theta_2_par_1 = [np.linspace(0.1, 3, 15)] - theta_2_par_5 = [np.array([0.4, -0.2, 1.2, -0.5, 2.2])[:, np.newaxis] * np.eye(15)[2::3]] - test_data_two_wires_with_parameters += [ - (qml.SpecialUnitary, s, mat_vec(qml.SpecialUnitary.compute_matrix(par[0], 2), s), par) - for s, par in [(state_1, theta_2_par_5), (state_5, theta_2_par_1), (state_5, theta_2_par_5)] - ] - - @pytest.mark.parametrize( - "operation,input,expected_output,par", test_data_two_wires_with_parameters - ) - def test_apply_operation_two_wires_with_parameters_broadcasted( - self, qubit_device_2_wires, tol, operation, input, expected_output, par - ): - """Tests that applying an operation yields the expected output state for two wire - operations that have parameters.""" - - shape = (5, 2, 2) if np.array(input).size == 20 else (2, 2) - dtype = qubit_device_2_wires.C_DTYPE - qubit_device_2_wires.target_device._state = np.array(input, dtype=dtype).reshape(shape) - par = tuple(np.array(p) for p in par) - qubit_device_2_wires.apply([operation(*par, wires=[0, 1])]) - - assert np.allclose( - qubit_device_2_wires._state.reshape((5, 4)), expected_output, atol=tol, rtol=0 - ) - assert qubit_device_2_wires._state.dtype == qubit_device_2_wires.C_DTYPE - - def test_apply_errors_qubit_state_vector_broadcasted(self, qubit_device_2_wires): - """Test that apply fails for incorrect state preparation, and > 2 qubit gates""" - with pytest.raises(ValueError, match="The state must be a vector of norm 1.0"): - qubit_device_2_wires.apply([qml.StatePrep(np.array([[1, -1], [0, 2]]), wires=[0])]) - - # Also test that the sum-check is *not* performed along the broadcasting dimension - qubit_device_2_wires.apply([qml.StatePrep(np.array([[0.6, 0.8], [0.6, 0.8]]), wires=[0])]) - - with pytest.raises(ValueError, match=r"State must be of length 4"): - # Second dimension does not match 2**num_wires - p = np.array([[1, 0, 1, 1, 0], [0, 1, 1, 0, 1]]) / np.sqrt(3) - qubit_device_2_wires.apply([qml.StatePrep(p, wires=[0, 1])]) - - with pytest.raises(ValueError, match=r"State must be of length 4"): - # Broadcasting dimension is not first dimension - p = np.array([[1, 1, 0], [0, 1, 1], [1, 0, 1], [0, 1, 1]]) / np.sqrt(2) - qubit_device_2_wires.apply([qml.StatePrep(p, wires=[0, 1])]) - - qubit_device_2_wires.reset() - vec = qml.StatePrep(np.array([[0, 1, 0, 0], [0, 0, 1, 0]]), wires=[0, 1]) - with pytest.raises( - DeviceError, - match="Operation StatePrep cannot be used after other Operations have already been applied " - "on a default.qubit.legacy device.", - ): - qubit_device_2_wires.apply([qml.RZ(0.5, wires=[0]), vec]) - - @pytest.mark.skip("Applying a BasisState does not support broadcasting yet") - def test_apply_errors_basis_state_broadcasted(self, qubit_device_2_wires): - """Test that applying the BasisState operation raises the correct errors.""" - with pytest.raises( - ValueError, match="BasisState parameter must consist of 0 or 1 integers." - ): - op = qml.BasisState(np.array([[-0.2, 4.2], [0.5, 1.2]]), wires=[0, 1]) - qubit_device_2_wires.apply([op]) - - with pytest.raises( - ValueError, match="BasisState parameter and wires must be of equal length." - ): - # Test that the error is raised - qubit_device_2_wires.apply( - [qml.BasisState(np.array([[0, 1], [1, 1], [1, 0]]), wires=[0])] - ) - # Test that the broadcasting dimension is allowed to mismatch the length of the wires - qubit_device_2_wires.apply([qml.BasisState(np.array([[0], [1], [0]]), wires=[0])]) - - qubit_device_2_wires.reset() - qubit_device_2_wires.apply([qml.RZ(0.5, wires=[0])]) - vec = qml.BasisState(np.array([[0, 0], [1, 0], [1, 1]]), wires=[0, 1]) - with pytest.raises( - DeviceError, - match="Operation BasisState cannot be used after other Operations have already been applied " - "on a default.qubit.legacy device.", - ): - qubit_device_2_wires.apply([vec]) - - -zero = [1, 0] -one = [0, 1] -plus = [INVSQ2, INVSQ2] -minus = [INVSQ2, -INVSQ2] -y_plus = [INVSQ2, 1j * INVSQ2] -y_minus = [INVSQ2, -1j * INVSQ2] - - -class TestExpvalBroadcasted: - """Tests that expectation values are properly calculated for broadcasted states - or that the proper errors are raised.""" - - @pytest.mark.parametrize( - "operation,input,expected_output", - [ - (qml.PauliX, np.array([plus, zero, minus]), [1, 0, -1]), - (qml.PauliY, np.array([y_plus, zero, y_minus]), [1, 0, -1]), - (qml.PauliZ, np.array([plus, zero, one]), [0, 1, -1]), - (qml.Hadamard, np.array([plus, zero, one]), [INVSQ2, INVSQ2, -INVSQ2]), - (qml.Identity, np.array([minus, zero, one]), [1, 1, 1]), - ], - ) - def test_expval_single_wire_no_parameters_broadcasted( - self, qubit_device_1_wire, tol, operation, input, expected_output - ): - """Tests that expectation values are properly calculated for single-wire observables without parameters.""" - - obs = operation(wires=[0]) - - qubit_device_1_wire.reset() - qubit_device_1_wire.apply( - [qml.StatePrep(np.array(input), wires=[0])], obs.diagonalizing_gates() - ) - res = qubit_device_1_wire.expval(obs) - - assert np.allclose(res, expected_output, atol=tol, rtol=0) - - @pytest.mark.parametrize( - "operation,input,expected_output,par", - [(qml.Hermitian, [zero, one, minus, y_plus], [1, 1, 1, 0], I - Y)], - ) - def test_expval_single_wire_with_parameters_broadcasted( - self, qubit_device_1_wire, tol, operation, input, expected_output, par - ): - """Tests that expectation values are properly calculated for single-wire observables with parameters.""" - - obs = operation(np.array(par), wires=[0]) - - qubit_device_1_wire.reset() - qubit_device_1_wire.apply( - [qml.StatePrep(np.array(input), wires=[0])], obs.diagonalizing_gates() - ) - res = qubit_device_1_wire.expval(obs) - - assert np.allclose(res, expected_output, atol=tol, rtol=0) - - @pytest.mark.parametrize( - "operation,input,expected_output,par", - [ - ( - qml.Hermitian, - [ - [1 / math.sqrt(3), 0, 1 / math.sqrt(3), 1 / math.sqrt(3)], - [0, 0, 0, 1], - [1 / math.sqrt(2), 0, -1 / math.sqrt(2), 0], - ], - [4 / 3, 0, 1], - [[1, 1j, 0, 1], [-1j, 1, 0, 0], [0, 0, 1, -1j], [1, 0, 1j, 0]], - ), - ( - qml.Hermitian, - [[INVSQ2, 0, 0, INVSQ2], [0, INVSQ2, -INVSQ2, 0]], - [1, -1], - [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]], - ), - ], - ) - def test_expval_two_wires_with_parameters_broadcasted( - self, qubit_device_2_wires, tol, operation, input, expected_output, par - ): - """Tests that expectation values are properly calculated for two-wire observables with parameters.""" - - obs = operation(np.array(par), wires=[0, 1]) - - qubit_device_2_wires.reset() - qubit_device_2_wires.apply( - [qml.StatePrep(np.array(input), wires=[0, 1])], obs.diagonalizing_gates() - ) - res = qubit_device_2_wires.expval(obs) - - assert np.allclose(res, expected_output, atol=tol, rtol=0) - - def test_expval_estimate_broadcasted(self): - """Test that the expectation value is not analytically calculated""" - - dev = qml.device("default.qubit.legacy", wires=1, shots=3) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(): - qml.RX(np.zeros(5), wires=0) # Broadcast the tape without applying an op - return qml.expval(qml.PauliX(0)) - - expval = circuit() - - # With 3 samples we are guaranteed to see a difference between - # an estimated variance an an analytically calculated one - assert np.all(expval != 0.0) - - -class TestVarBroadcasted: - """Tests that variances are properly calculated for broadcasted states.""" - - @pytest.mark.parametrize( - "operation,input,expected_output", - [ - (qml.PauliX, [plus, zero, minus], [0, 1, 0]), - (qml.PauliY, [y_plus, zero, y_minus], [0, 1, 0]), - (qml.PauliZ, [plus, zero, one], [1, 0, 0]), - (qml.Hadamard, [plus, zero, one], [0.5, 0.5, 0.5]), - (qml.Identity, [minus, zero, one], [0, 0, 0]), - ], - ) - def test_var_single_wire_no_parameters_broadcasted( - self, qubit_device_1_wire, tol, operation, input, expected_output - ): - """Tests that variances are properly calculated for single-wire observables without parameters.""" - - obs = operation(wires=[0]) - - qubit_device_1_wire.reset() - qubit_device_1_wire.apply( - [qml.StatePrep(np.array(input), wires=[0])], obs.diagonalizing_gates() - ) - res = qubit_device_1_wire.var(obs) - - assert np.allclose(res, expected_output, atol=tol, rtol=0) - - @pytest.mark.parametrize( - "operation,input,expected_output,par", - [(qml.Hermitian, [zero, one, minus, y_plus], [1, 1, 1, 0], I - Y)], - ) - def test_var_single_wire_with_parameters_broadcasted( - self, qubit_device_1_wire, tol, operation, input, expected_output, par - ): - """Tests that variances are properly calculated for single-wire observables with parameters.""" - - obs = operation(np.array(par), wires=[0]) - - qubit_device_1_wire.reset() - qubit_device_1_wire.apply( - [qml.StatePrep(np.array(input), wires=[0])], obs.diagonalizing_gates() - ) - res = qubit_device_1_wire.var(obs) - - assert np.allclose(res, expected_output, atol=tol, rtol=0) - - @pytest.mark.parametrize( - "operation,input,expected_output,par", - [ - ( - qml.Hermitian, - [ - [1 / math.sqrt(3), 0, 1 / math.sqrt(3), 1 / math.sqrt(3)], - [0, 0, 0, 1], - [1 / math.sqrt(2), 0, -1 / math.sqrt(2), 0], - ], - [11 / 9, 2, 3 / 2], - [[1, 1j, 0, 1], [-1j, 1, 0, 0], [0, 0, 1, -1j], [1, 0, 1j, 1]], - ), - ( - qml.Hermitian, - [[INVSQ2, 0, 0, INVSQ2], [0, INVSQ2, -INVSQ2, 0]], - [0, 0], - [[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]], - ), - ], - ) - def test_var_two_wires_with_parameters_broadcasted( - self, qubit_device_2_wires, tol, operation, input, expected_output, par - ): - """Tests that variances are properly calculated for two-wire observables with parameters.""" - - obs = operation(np.array(par), wires=[0, 1]) - - qubit_device_2_wires.reset() - qubit_device_2_wires.apply( - [qml.StatePrep(np.array(input), wires=[0, 1])], obs.diagonalizing_gates() - ) - res = qubit_device_2_wires.var(obs) - - assert np.allclose(res, expected_output, atol=tol, rtol=0) - - def test_var_estimate_broadcasted(self): - """Test that the variance is not analytically calculated""" - - dev = qml.device("default.qubit.legacy", wires=1, shots=3) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(): - qml.RX(np.zeros(5), wires=0) # Broadcast the tape without applying an op - return qml.var(qml.PauliX(0)) - - var = circuit() - - # With 3 samples we are guaranteed to see a difference between - # an estimated variance and an analytically calculated one - assert np.all(var != 1.0) - - -class TestSampleBroadcasted: - """Tests that samples are properly calculated for broadcasted states.""" - - def test_sample_dimensions_broadcasted(self): - """Tests if the samples returned by the sample function have - the correct dimensions - """ - - # Explicitly resetting is necessary as the internal - # state is set to None in __init__ and only properly - # initialized during reset - dev = qml.device("default.qubit.legacy", wires=2, shots=1000) - - dev.apply([qml.RX(np.array([np.pi / 2, 0.0]), 0), qml.RX(np.array([np.pi / 2, 0.0]), 1)]) - - dev.target_device.shots = 10 - dev.target_device._wires_measured = {0} - dev.target_device._samples = dev.generate_samples() - s1 = dev.sample(qml.PauliZ(0)) - assert s1.shape == ( - 2, - 10, - ) - - dev.reset() - dev.target_device.shots = 12 - dev.target_device._wires_measured = {1} - dev.target_device._samples = dev.generate_samples() - s2 = dev.sample(qml.PauliZ(wires=[1])) - assert s2.shape == (12,) - - dev.reset() - dev.apply([qml.RX(np.ones(5), 0), qml.RX(np.ones(5), 1)]) - dev.target_device.shots = 17 - dev.target_device._wires_measured = {0, 1} - dev.target_device._samples = dev.generate_samples() - s3 = dev.sample(qml.PauliX(0) @ qml.PauliZ(1)) - assert s3.shape == (5, 17) - - def test_sample_values_broadcasted(self, tol): - """Tests if the samples returned by sample have - the correct values - """ - - # Explicitly resetting is necessary as the internal - # state is set to None in __init__ and only properly - # initialized during reset - dev = qml.device("default.qubit.legacy", wires=2, shots=1000) - - dev.apply([qml.RX(np.ones(3), wires=[0])]) - dev.target_device._wires_measured = {0} - dev.target_device._samples = dev.generate_samples() - - s1 = dev.sample(qml.PauliZ(0)) - - # s1 should only contain 1 and -1, which is guaranteed if - # they square to 1 - assert np.allclose(s1**2, 1, atol=tol, rtol=0) - - -class TestDefaultQubitIntegrationBroadcasted: - """Integration tests for default.qubit.legacy. This test ensures it integrates - properly with the PennyLane interface, in particular QNode.""" - - @pytest.mark.parametrize("r_dtype", [np.float32, np.float64]) - def test_qubit_circuit_broadcasted(self, qubit_device_1_wire, r_dtype, tol): - """Test that the default qubit plugin provides correct result for a simple circuit""" - - p = np.array([0.543, np.pi / 2, 0.0, 1.0]) - - dev = qubit_device_1_wire - dev.target_device.R_DTYPE = r_dtype - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliY(0)) - - expected = -np.sin(p) - - res = circuit(p) - assert np.allclose(res, expected, atol=tol, rtol=0) - assert res.dtype == r_dtype # pylint:disable=no-member - - def test_qubit_identity_broadcasted(self, qubit_device_1_wire, tol): - """Test that the default qubit plugin provides correct result for the Identity expectation""" - - p = np.array([0.543, np.pi / 2, 0.0, 1.0]) - - @qml.qnode(qubit_device_1_wire) - def circuit(x): - """Test quantum function""" - qml.RX(x, wires=0) - return qml.expval(qml.Identity(0)) - - assert np.allclose(circuit(p), 1, atol=tol, rtol=0) - - def test_nonzero_shots_broadcasted(self, tol): - """Test that the default qubit plugin provides correct result for high shot number""" - - shots = 10**5 - dev = qml.device("default.qubit.legacy", wires=1, shots=shots) - - p = np.array([0.543, np.pi / 2, 0.0, 1.0]) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(x): - """Test quantum function""" - qml.RX(x, wires=0) - return qml.expval(qml.PauliY(0)) - - runs = [] - for _ in range(100): - runs.append(circuit(p)) - - assert np.allclose(np.mean(runs, axis=0), -np.sin(p), atol=tol, rtol=0) - - @pytest.mark.parametrize( - "name,state,expected_output", - [ - ("PauliX", [plus, minus, zero], [1, -1, 0]), - ("PauliY", [y_plus, y_minus, zero], [1, -1, 0]), - ("PauliZ", [plus, one, zero], [0, -1, 1]), - ("Hadamard", [plus, one, zero], [INVSQ2, -INVSQ2, INVSQ2]), - ], - ) - def test_supported_observable_single_wire_no_parameters_broadcasted( - self, qubit_device_1_wire, tol, name, state, expected_output - ): - """Tests supported observables on single wires without parameters.""" - - obs = getattr(qml.ops, name) - - assert qubit_device_1_wire.supports_observable(name) - - @qml.qnode(qubit_device_1_wire) - def circuit(): - qml.StatePrep(np.array(state), wires=[0]) - return qml.expval(obs(wires=[0])) - - assert np.allclose(circuit(), expected_output, atol=tol, rtol=0) - - @pytest.mark.parametrize( - "name,state,expected_output,par", - [ - ("Identity", [zero, one, plus], [1, 1, 1], []), - ("Hermitian", [zero, one, minus], [1, 1, 1], [I - Y]), - ], - ) - def test_supported_observable_single_wire_with_parameters_broadcasted( - self, qubit_device_1_wire, tol, name, state, expected_output, par - ): - """Tests supported observables on single wires with parameters.""" - - obs = getattr(qml.ops, name) - - assert qubit_device_1_wire.supports_observable(name) - - @qml.qnode(qubit_device_1_wire) - def circuit(): - qml.StatePrep(np.array(state), wires=[0]) - return qml.expval(obs(*par, wires=[0])) - - assert np.allclose(circuit(), expected_output, atol=tol, rtol=0) - - @pytest.mark.parametrize( - "name,state,expected_output,par", - [ - ( - "Hermitian", - [ - [1 / math.sqrt(3), 0, 1 / math.sqrt(3), 1 / math.sqrt(3)], - [0, 0, 0, 1], - [1 / math.sqrt(2), 0, -1 / math.sqrt(2), 0], - ], - [4 / 3, 0, 1], - ([[1, 1j, 0, 1], [-1j, 1, 0, 0], [0, 0, 1, -1j], [1, 0, 1j, 0]],), - ), - ( - "Hermitian", - [[INVSQ2, 0, 0, INVSQ2], [0, INVSQ2, -INVSQ2, 0]], - [1, -1], - ([[1, 0, 0, 0], [0, -1, 0, 0], [0, 0, -1, 0], [0, 0, 0, 1]],), - ), - ], - ) - def test_supported_observable_two_wires_with_parameters_broadcasted( - self, qubit_device_2_wires, tol, name, state, expected_output, par - ): - """Tests supported observables on two wires with parameters.""" - - obs = getattr(qml.ops, name) - - assert qubit_device_2_wires.supports_observable(name) - - @qml.qnode(qubit_device_2_wires) - def circuit(): - qml.StatePrep(np.array(state), wires=[0, 1]) - return qml.expval(obs(*par, wires=[0, 1])) - - assert np.allclose(circuit(), expected_output, atol=tol, rtol=0) - - def test_multi_samples_return_correlated_results_broadcasted(self): - """Tests if the samples returned by the sample function are correlated - correctly for correlated observables. - """ - - dev = qml.device("default.qubit.legacy", wires=2, shots=1000) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(): - qml.Hadamard(0) - qml.RX(np.zeros(5), 0) - qml.CNOT(wires=[0, 1]) - return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)) - - outcomes = circuit() - - assert np.array_equal(outcomes[0], outcomes[1]) - - @pytest.mark.parametrize("num_wires", [3, 4, 5, 6, 7, 8]) - def test_multi_samples_correlated_results_more_wires_than_observable_broadcasted( - self, num_wires - ): - """Tests if the samples returned by the sample function are correlated - correctly for correlated observables on larger devices than the observables - """ - - dev = qml.device("default.qubit.legacy", wires=num_wires, shots=1000) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(): - qml.Hadamard(0) - qml.RX(np.zeros(5), 0) - qml.CNOT(wires=[0, 1]) - return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)) - - outcomes = circuit() - - assert np.array_equal(outcomes[0], outcomes[1]) - - -# pylint: disable=unused-argument -@pytest.mark.parametrize( - "theta,phi,varphi", [(THETA, PHI, VARPHI), (THETA, PHI[0], VARPHI), (THETA[0], PHI, VARPHI[1])] -) -class TestTensorExpvalBroadcasted: - """Test tensor expectation values for broadcasted states""" - - def test_paulix_pauliy_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving PauliX and PauliY works correctly""" - dev = qml.device("default.qubit.legacy", wires=3) - dev.reset() - - obs = qml.PauliX(0) @ qml.PauliY(2) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.expval(obs) - - expected = np.sin(theta) * np.sin(phi) * np.sin(varphi) - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_pauliz_identity_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving PauliZ and Identity works correctly""" - dev = qml.device("default.qubit.legacy", wires=3) - dev.reset() - - obs = qml.PauliZ(0) @ qml.Identity(1) @ qml.PauliZ(2) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.expval(obs) - - expected = np.cos(varphi) * np.cos(phi) - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_pauliz_hadamard_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly""" - dev = qml.device("default.qubit.legacy", wires=3) - obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2) - - dev.reset() - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.expval(obs) - - expected = -(np.cos(varphi) * np.sin(phi) + np.sin(varphi) * np.cos(theta)) / np.sqrt(2) - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_hermitian_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving qml.Hermitian works correctly""" - dev = qml.device("default.qubit.legacy", wires=3) - dev.reset() - - A = np.array( - [ - [-6, 2 + 1j, -3, -5 + 2j], - [2 - 1j, 0, 2 - 1j, -5 + 4j], - [-3, 2 + 1j, 0, -4 + 3j], - [-5 - 2j, -5 - 4j, -4 - 3j, -6], - ] - ) - - obs = qml.PauliZ(0) @ qml.Hermitian(A, wires=[1, 2]) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.expval(obs) - - expected = 0.5 * ( - -6 * np.cos(theta) * (np.cos(varphi) + 1) - - 2 * np.sin(varphi) * (np.cos(theta) + np.sin(phi) - 2 * np.cos(phi)) - + 3 * np.cos(varphi) * np.sin(phi) - + np.sin(phi) - ) - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_hermitian_hermitian_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving two Hermitian matrices works correctly""" - dev = qml.device("default.qubit.legacy", wires=3) - - A1 = np.array([[1, 2], [2, 4]]) - - A2 = np.array( - [ - [-6, 2 + 1j, -3, -5 + 2j], - [2 - 1j, 0, 2 - 1j, -5 + 4j], - [-3, 2 + 1j, 0, -4 + 3j], - [-5 - 2j, -5 - 4j, -4 - 3j, -6], - ] - ) - - obs = qml.Hermitian(A1, wires=[0]) @ qml.Hermitian(A2, wires=[1, 2]) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.expval(obs) - - expected = 0.25 * ( - -30 - + 4 * np.cos(phi) * np.sin(theta) - + 3 * np.cos(varphi) * (-10 + 4 * np.cos(phi) * np.sin(theta) - 3 * np.sin(phi)) - - 3 * np.sin(phi) - - 2 - * (5 + np.cos(phi) * (6 + 4 * np.sin(theta)) + (-3 + 8 * np.sin(theta)) * np.sin(phi)) - * np.sin(varphi) - + np.cos(theta) - * ( - 18 - + 5 * np.sin(phi) - + 3 * np.cos(varphi) * (6 + 5 * np.sin(phi)) - + 2 * (3 + 10 * np.cos(phi) - 5 * np.sin(phi)) * np.sin(varphi) - ) - ) - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_hermitian_identity_expectation_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving an Hermitian matrix and the identity works correctly""" - dev = qml.device("default.qubit.legacy", wires=2) - - A = np.array( - [[1.02789352, 1.61296440 - 0.3498192j], [1.61296440 + 0.3498192j, 1.23920938 + 0j]] - ) - - obs = qml.Hermitian(A, wires=[0]) @ qml.Identity(wires=[1]) - - dev.apply( - [qml.RY(theta, wires=[0]), qml.RY(phi, wires=[1]), qml.CNOT(wires=[0, 1])], - obs.diagonalizing_gates(), - ) - - res = dev.expval(obs) - - a = A[0, 0] - re_b = A[0, 1].real - d = A[1, 1] - expected = ((a - d) * np.cos(theta) + 2 * re_b * np.sin(theta) * np.sin(phi) + a + d) / 2 - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_hermitian_two_wires_identity_expectation_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving an Hermitian matrix for two wires and the identity works correctly""" - dev = qml.device("default.qubit.legacy", wires=3) - - A = np.array( - [[1.02789352, 1.61296440 - 0.3498192j], [1.61296440 + 0.3498192j, 1.23920938 + 0j]] - ) - Identity = np.array([[1, 0], [0, 1]]) - Ham = np.kron(np.kron(Identity, Identity), A) - obs = qml.Hermitian(Ham, wires=[2, 1, 0]) - - dev.apply( - [qml.RY(theta, wires=[0]), qml.RY(phi, wires=[1]), qml.CNOT(wires=[0, 1])], - obs.diagonalizing_gates(), - ) - res = dev.expval(obs) - - a = A[0, 0] - re_b = A[0, 1].real - d = A[1, 1] - - expected = ((a - d) * np.cos(theta) + 2 * re_b * np.sin(theta) * np.sin(phi) + a + d) / 2 - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize( - "theta,phi,varphi", [(THETA, PHI, VARPHI), (THETA, PHI[0], VARPHI), (THETA[0], PHI, VARPHI[1])] -) -class TestTensorVarBroadcasted: - """Tests for variance of tensor observables for broadcasted states""" - - def test_paulix_pauliy_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving PauliX and PauliY works correctly""" - dev = qml.device("default.qubit.legacy", wires=3) - - obs = qml.PauliX(0) @ qml.PauliY(2) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.var(obs) - - expected = ( - 8 * np.sin(theta) ** 2 * np.cos(2 * varphi) * np.sin(phi) ** 2 - - np.cos(2 * (theta - phi)) - - np.cos(2 * (theta + phi)) - + 2 * np.cos(2 * theta) - + 2 * np.cos(2 * phi) - + 14 - ) / 16 - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_pauliz_hadamard_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly""" - dev = qml.device("default.qubit.legacy", wires=3) - obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2) - - dev.reset() - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.var(obs) - - expected = ( - 3 - + np.cos(2 * phi) * np.cos(varphi) ** 2 - - np.cos(2 * theta) * np.sin(varphi) ** 2 - - 2 * np.cos(theta) * np.sin(phi) * np.sin(2 * varphi) - ) / 4 - - assert np.allclose(res, expected, atol=tol, rtol=0) - - def test_hermitian_broadcasted(self, theta, phi, varphi, tol): - """Test that a tensor product involving qml.Hermitian works correctly""" - dev = qml.device("default.qubit.legacy", wires=3) - - A = np.array( - [ - [-6, 2 + 1j, -3, -5 + 2j], - [2 - 1j, 0, 2 - 1j, -5 + 4j], - [-3, 2 + 1j, 0, -4 + 3j], - [-5 - 2j, -5 - 4j, -4 - 3j, -6], - ] - ) - - obs = qml.PauliZ(0) @ qml.Hermitian(A, wires=[1, 2]) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - res = dev.var(obs) - - expected = ( - 1057 - - np.cos(2 * phi) - + 12 * (27 + np.cos(2 * phi)) * np.cos(varphi) - - 2 * np.cos(2 * varphi) * np.sin(phi) * (16 * np.cos(phi) + 21 * np.sin(phi)) - + 16 * np.sin(2 * phi) - - 8 * (-17 + np.cos(2 * phi) + 2 * np.sin(2 * phi)) * np.sin(varphi) - - 8 * np.cos(2 * theta) * (3 + 3 * np.cos(varphi) + np.sin(varphi)) ** 2 - - 24 * np.cos(phi) * (np.cos(phi) + 2 * np.sin(phi)) * np.sin(2 * varphi) - - 8 - * np.cos(theta) - * ( - 4 - * np.cos(phi) - * ( - 4 - + 8 * np.cos(varphi) - + np.cos(2 * varphi) - - (1 + 6 * np.cos(varphi)) * np.sin(varphi) - ) - + np.sin(phi) - * ( - 15 - + 8 * np.cos(varphi) - - 11 * np.cos(2 * varphi) - + 42 * np.sin(varphi) - + 3 * np.sin(2 * varphi) - ) - ) - ) / 16 - - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@pytest.mark.parametrize( - "theta,phi,varphi", [(THETA, PHI, VARPHI), (THETA, PHI[0], VARPHI), (THETA[0], PHI, VARPHI[1])] -) -class TestTensorSampleBroadcasted: - """Test tensor sampling for broadcated states""" - - def test_paulix_pauliy_broadcasted(self, theta, phi, varphi, tol_stochastic): - """Test that a tensor product involving PauliX and PauliY works correctly""" - dev = qml.device("default.qubit.legacy", wires=3, shots=int(1e6)) - - obs = qml.PauliX(0) @ qml.PauliY(2) - - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - dev.target_device._wires_measured = {0, 1, 2} - dev.target_device._samples = dev.generate_samples() - dev.sample(obs) - - s1 = obs.eigvals() - p = dev.probability(wires=dev.map_wires(obs.wires)) - - # s1 should only contain 1 and -1 - assert np.allclose(s1**2, 1, atol=tol_stochastic, rtol=0) - - mean = p @ s1 - expected = np.sin(theta) * np.sin(phi) * np.sin(varphi) - assert np.allclose(mean, expected, atol=tol_stochastic, rtol=0) - - var = p @ (s1**2) - (p @ s1).real ** 2 - expected = ( - 8 * np.sin(theta) ** 2 * np.cos(2 * varphi) * np.sin(phi) ** 2 - - np.cos(2 * (theta - phi)) - - np.cos(2 * (theta + phi)) - + 2 * np.cos(2 * theta) - + 2 * np.cos(2 * phi) - + 14 - ) / 16 - assert np.allclose(var, expected, atol=tol_stochastic, rtol=0) - - def test_pauliz_hadamard_broadcasted(self, theta, phi, varphi, tol_stochastic): - """Test that a tensor product involving PauliZ and PauliY and hadamard works correctly""" - dev = qml.device("default.qubit.legacy", wires=3, shots=int(1e6)) - obs = qml.PauliZ(0) @ qml.Hadamard(1) @ qml.PauliY(2) - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - dev.target_device._wires_measured = {0, 1, 2} - dev.target_device._samples = dev.generate_samples() - dev.sample(obs) - - s1 = obs.eigvals() - p = dev.marginal_prob(dev.probability(), wires=obs.wires) - - # s1 should only contain 1 and -1 - assert np.allclose(s1**2, 1, atol=tol_stochastic, rtol=0) - - mean = p @ s1 - expected = -(np.cos(varphi) * np.sin(phi) + np.sin(varphi) * np.cos(theta)) / np.sqrt(2) - assert np.allclose(mean, expected, atol=tol_stochastic, rtol=0) - - var = p @ (s1**2) - (p @ s1).real ** 2 - expected = ( - 3 - + np.cos(2 * phi) * np.cos(varphi) ** 2 - - np.cos(2 * theta) * np.sin(varphi) ** 2 - - 2 * np.cos(theta) * np.sin(phi) * np.sin(2 * varphi) - ) / 4 - assert np.allclose(var, expected, atol=tol_stochastic, rtol=0) - - def test_hermitian_broadcasted(self, theta, phi, varphi, tol_stochastic): - """Test that a tensor product involving qml.Hermitian works correctly""" - dev = qml.device("default.qubit.legacy", wires=3, shots=int(1e6)) - - A = 0.1 * np.array( - [ - [-6, 2 + 1j, -3, -5 + 2j], - [2 - 1j, 0, 2 - 1j, -5 + 4j], - [-3, 2 + 1j, 0, -4 + 3j], - [-5 - 2j, -5 - 4j, -4 - 3j, -6], - ] - ) - - obs = qml.PauliZ(0) @ qml.Hermitian(A, wires=[1, 2]) - dev.apply( - [ - qml.RX(theta, wires=[0]), - qml.RX(phi, wires=[1]), - qml.RX(varphi, wires=[2]), - qml.CNOT(wires=[0, 1]), - qml.CNOT(wires=[1, 2]), - ], - obs.diagonalizing_gates(), - ) - - dev.target_device._wires_measured = {0, 1, 2} - dev.target_device._samples = dev.generate_samples() - dev.sample(obs) - - s1 = obs.eigvals() - p = dev.marginal_prob(dev.probability(), wires=obs.wires) - - # s1 should only contain the eigenvalues of - # the hermitian matrix tensor product Z - z = np.diag([1, -1]) - eigvals = np.linalg.eigvalsh(np.kron(z, A)) - assert set(np.round(s1, 8).tolist()).issubset(set(np.round(eigvals, 8).tolist())) - - mean = p @ s1 - expected = ( - 0.1 - * 0.5 - * ( - -6 * np.cos(theta) * (np.cos(varphi) + 1) - - 2 * np.sin(varphi) * (np.cos(theta) + np.sin(phi) - 2 * np.cos(phi)) - + 3 * np.cos(varphi) * np.sin(phi) - + np.sin(phi) - ) - ) - assert np.allclose(mean, expected, atol=tol_stochastic, rtol=0) - - var = p @ (s1**2) - (p @ s1).real ** 2 - expected = ( - 0.01 - * ( - 1057 - - np.cos(2 * phi) - + 12 * (27 + np.cos(2 * phi)) * np.cos(varphi) - - 2 * np.cos(2 * varphi) * np.sin(phi) * (16 * np.cos(phi) + 21 * np.sin(phi)) - + 16 * np.sin(2 * phi) - - 8 * (-17 + np.cos(2 * phi) + 2 * np.sin(2 * phi)) * np.sin(varphi) - - 8 * np.cos(2 * theta) * (3 + 3 * np.cos(varphi) + np.sin(varphi)) ** 2 - - 24 * np.cos(phi) * (np.cos(phi) + 2 * np.sin(phi)) * np.sin(2 * varphi) - - 8 - * np.cos(theta) - * ( - 4 - * np.cos(phi) - * ( - 4 - + 8 * np.cos(varphi) - + np.cos(2 * varphi) - - (1 + 6 * np.cos(varphi)) * np.sin(varphi) - ) - + np.sin(phi) - * ( - 15 - + 8 * np.cos(varphi) - - 11 * np.cos(2 * varphi) - + 42 * np.sin(varphi) - + 3 * np.sin(2 * varphi) - ) - ) - ) - / 16 - ) - assert np.allclose(var, expected, atol=tol_stochastic, rtol=0) - - -@pytest.mark.parametrize( - "r_dtype,c_dtype", [(np.float32, np.complex64), (np.float64, np.complex128)] -) -class TestDtypePreservedBroadcasted: - """Test that the user-defined dtype of the device is preserved for QNode - evaluation""" - - @pytest.mark.parametrize( - "op", - [ - qml.SingleExcitation, - qml.SingleExcitationPlus, - qml.SingleExcitationMinus, - qml.DoubleExcitation, - qml.DoubleExcitationPlus, - qml.DoubleExcitationMinus, - qml.OrbitalRotation, - qml.FermionicSWAP, - qml.QubitSum, - qml.QubitCarry, - ], - ) - def test_state_dtype_after_op_broadcasted(self, r_dtype, c_dtype, op): - """Test that the default qubit plugin preserves data types of states when an operation is - applied. As TestApply class check most of operators, we here only check some subtle - examples. - """ - - dev = qml.device("default.qubit.legacy", wires=4, r_dtype=r_dtype, c_dtype=c_dtype) - - n_wires = op.num_wires - n_params = op.num_params - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(): - x = np.array([0.543, 0.622, 1.3]) - if n_params == 0: - op(wires=range(n_wires)) - elif n_params == 1: - op(x, wires=range(n_wires)) - else: - op([x] * n_params, wires=range(n_wires)) - return qml.state() - - res = circuit() - assert res.dtype == c_dtype # pylint:disable=no-member - - @pytest.mark.parametrize( - "measurement", - [ - qml.expval(qml.PauliY(0)), - qml.var(qml.PauliY(0)), - qml.probs(wires=[1]), - qml.probs(wires=[2, 0]), - ], - ) - def test_measurement_real_dtype_broadcasted(self, r_dtype, c_dtype, measurement): - """Test that the default qubit plugin provides correct result for a simple circuit""" - p = np.array([0.543, 0.622, 1.3]) - - dev = qml.device("default.qubit.legacy", wires=3, r_dtype=r_dtype, c_dtype=c_dtype) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(x): - qml.RX(x, wires=0) - return qml.apply(measurement) - - res = circuit(p) - assert res.dtype == r_dtype - - def test_measurement_complex_dtype_broadcasted(self, r_dtype, c_dtype): - """Test that the default qubit plugin provides correct result for a simple circuit""" - p = np.array([0.543, 0.622, 1.3]) - m = qml.state() - - dev = qml.device("default.qubit.legacy", wires=3, r_dtype=r_dtype, c_dtype=c_dtype) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(x): - qml.RX(x, wires=0) - return qml.apply(m) - - res = circuit(p) - assert res.dtype == c_dtype - - -class TestProbabilityIntegrationBroadcasted: - """Test probability method for when analytic is True/False""" - - # pylint: disable=no-member, unused-argument - def mock_analytic_counter(self, wires=None): - self.analytic_counter += 1 - return np.array([1, 0, 0, 0], dtype=float) - - def test_probability_broadcasted(self, tol): - """Test that the probability function works for finite and infinite shots""" - dev = qml.device("default.qubit.legacy", wires=2, shots=1000) - dev_analytic = qml.device("default.qubit.legacy", wires=2, shots=None) - - x = np.array([[0.2, 0.5, 0.4], [0.9, 0.8, 0.3]]) - - def circuit(x): - qml.RX(x[0], wires=0) - qml.RY(x[1], wires=0) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[0, 1]) - - prob = qml.QNode(circuit, dev) - prob_analytic = qml.QNode(circuit, dev_analytic) - - assert np.allclose(prob(x).sum(axis=-1), 1, atol=tol, rtol=0) - assert np.allclose(prob_analytic(x), prob(x), atol=0.1, rtol=0) - assert not np.array_equal(prob_analytic(x), prob(x)) - - -class TestWiresIntegrationBroadcasted: - """Test that the device integrates with PennyLane's wire management.""" - - def make_circuit_probs(self, wires): - """Factory for a qnode returning probabilities using arbitrary wire labels.""" - dev = qml.device("default.qubit.legacy", wires=wires) - n_wires = len(wires) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(): - qml.RX(np.array([0.5, 1.2, -0.6]), wires=wires[0 % n_wires]) - qml.RY(np.array([2.0, 0.4, 1.2]), wires=wires[1 % n_wires]) - if n_wires > 1: - qml.CNOT(wires=[wires[0], wires[1]]) - return qml.probs(wires=wires) - - return circuit - - @pytest.mark.parametrize( - "wires1, wires2", - [ - (["a", "c", "d"], [2, 3, 0]), - ([-1, -2, -3], ["q1", "ancilla", 2]), - (["a", "c"], [3, 0]), - ([-1, -2], ["ancilla", 2]), - (["a"], ["nothing"]), - ], - ) - def test_wires_probs_broadcasted(self, wires1, wires2, tol): - """Test that the probability vector of a circuit is independent from the wire labels used.""" - - circuit1 = self.make_circuit_probs(wires1) - circuit2 = self.make_circuit_probs(wires2) - - assert np.allclose(circuit1(), circuit2(), tol) - - -class TestApplyOpsBroadcasted: - """Tests for special methods listed in _apply_ops that use array manipulation tricks to apply - gates in DefaultQubitLegacy.""" - - broadcasted_state = np.arange(2**4 * 3, dtype=np.complex128).reshape((3, 2, 2, 2, 2)) - with pytest.warns(qml.PennyLaneDeprecationWarning): - dev = qml.device("default.qubit.legacy", wires=4) - - single_qubit_ops = [ - (qml.PauliX, dev._apply_x), - (qml.PauliY, dev._apply_y), - (qml.PauliZ, dev._apply_z), - (qml.Hadamard, dev._apply_hadamard), - (qml.S, dev._apply_s), - (qml.T, dev._apply_t), - (qml.SX, dev._apply_sx), - ] - two_qubit_ops = [ - (qml.CNOT, dev._apply_cnot), - (qml.SWAP, dev._apply_swap), - (qml.CZ, dev._apply_cz), - ] - three_qubit_ops = [ - (qml.Toffoli, dev._apply_toffoli), - ] - - @pytest.mark.parametrize("op, method", single_qubit_ops) - def test_apply_single_qubit_op_broadcasted_state(self, op, method): - """Test if the application of single qubit operations to a - broadcasted state is correct.""" - state_out = method(self.broadcasted_state, axes=[2]) - op = op(wires=[1]) - matrix = op.matrix() - state_out_einsum = np.einsum("ab,mibjk->miajk", matrix, self.broadcasted_state) - assert np.allclose(state_out, state_out_einsum) - - @pytest.mark.parametrize("op, method", two_qubit_ops) - def test_apply_two_qubit_op_broadcasted_state(self, op, method): - """Test if the application of two qubit operations to a - broadcasted state is correct.""" - state_out = method(self.broadcasted_state, axes=[1, 2]) - op = op(wires=[0, 1]) - matrix = op.matrix() - matrix = matrix.reshape((2, 2, 2, 2)) - state_out_einsum = np.einsum("abcd,mcdjk->mabjk", matrix, self.broadcasted_state) - assert np.allclose(state_out, state_out_einsum) - - @pytest.mark.parametrize("op, method", two_qubit_ops) - def test_apply_two_qubit_op_reverse_broadcasted_state(self, op, method): - """Test if the application of two qubit operations to a - broadcasted state is correct when the applied wires are reversed.""" - state_out = method(self.broadcasted_state, axes=[3, 2]) - op = op(wires=[2, 1]) - matrix = op.matrix() - matrix = matrix.reshape((2, 2, 2, 2)) - state_out_einsum = np.einsum("abcd,midck->mibak", matrix, self.broadcasted_state) - assert np.allclose(state_out, state_out_einsum) - - @pytest.mark.parametrize("op, method", three_qubit_ops) - def test_apply_three_qubit_op_controls_smaller_broadcasted_state(self, op, method): - """Test if the application of three qubit operations to a broadcasted - state is correct when both control wires are smaller than the target wire.""" - state_out = method(self.broadcasted_state, axes=[1, 3, 4]) - op = op(wires=[0, 2, 3]) - matrix = op.matrix() - matrix = matrix.reshape((2, 2) * 3) - state_out_einsum = np.einsum("abcdef,mdkef->makbc", matrix, self.broadcasted_state) - assert np.allclose(state_out, state_out_einsum) - - @pytest.mark.parametrize("op, method", three_qubit_ops) - def test_apply_three_qubit_op_controls_greater_broadcasted_state(self, op, method): - """Test if the application of three qubit operations to a broadcasted - state is correct when both control wires are greater than the target wire.""" - state_out = method(self.broadcasted_state, axes=[3, 2, 1]) - op = op(wires=[2, 1, 0]) - matrix = op.matrix() - matrix = matrix.reshape((2, 2) * 3) - state_out_einsum = np.einsum("abcdef,mfedk->mcbak", matrix, self.broadcasted_state) - assert np.allclose(state_out, state_out_einsum) - - @pytest.mark.parametrize("op, method", three_qubit_ops) - def test_apply_three_qubit_op_controls_split_broadcasted_state(self, op, method): - """Test if the application of three qubit operations to a broadcasted state is correct - when one control wire is smaller and one control wire is greater than the target wire.""" - state_out = method(self.broadcasted_state, axes=[4, 2, 3]) - op = op(wires=[3, 1, 2]) - matrix = op.matrix() - matrix = matrix.reshape((2, 2) * 3) - state_out_einsum = np.einsum("abcdef,mkdfe->mkacb", matrix, self.broadcasted_state) - assert np.allclose(state_out, state_out_einsum) - - -class TestStateVectorBroadcasted: - """Unit tests for the _apply_state_vector method with broadcasting""" - - def test_full_subsystem_broadcasted(self, mocker): - """Test applying a state vector to the full subsystem""" - dev = DefaultQubitLegacy(wires=["a", "b", "c"]) - state = np.array([[0, 1, 1, 0, 1, 1, 0, 0], [1, 0, 0, 0, 1, 0, 1, 1]]) / 2.0 - state_wires = qml.wires.Wires(["a", "b", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - - assert np.all(dev._state.reshape((2, 8)) == state) - spy.assert_not_called() - - def test_partial_subsystem_broadcasted(self, mocker): - """Test applying a state vector to a subset of wires of the full subsystem""" - - dev = DefaultQubitLegacy(wires=["a", "b", "c"]) - state = np.array([[0, 1, 1, 0], [1, 0, 1, 0], [1, 1, 0, 0]]) / np.sqrt(2.0) - state_wires = qml.wires.Wires(["a", "c"]) - - spy = mocker.spy(dev, "_scatter") - dev._apply_state_vector(state=state, device_wires=state_wires) - # Axes are (broadcasting, wire "a", wire "b", wire "c"), so we sum over axis=2 - res = np.sum(dev._state, axis=(2,)).reshape((3, 4)) - - assert np.all(res == state) - spy.assert_called() - - -class TestApplyOperationBroadcasted: - """Unit tests for the internal _apply_operation method.""" - - def test_internal_apply_ops_case_broadcasted(self, monkeypatch): - """Tests that if we provide an operation that has an internal - implementation, then we use that specific implementation. - - This test provides a new internal function that `default.qubit.legacy` uses to - apply `PauliX` (rather than redefining the gate itself). - """ - dev = qml.device("default.qubit.legacy", wires=1) - - test_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]]) - # Create a dummy operation - expected_test_output = np.ones(1) - supported_gate_application = lambda *args, **kwargs: expected_test_output - - with monkeypatch.context() as m: - # Set the internal ops implementations dict - m.setattr(dev.target_device, "_apply_ops", {"PauliX": supported_gate_application}) - - op = qml.PauliX(0) - - res = dev._apply_operation(test_state, op) - assert np.allclose(res, expected_test_output) - - def test_diagonal_operation_case_broadcasted(self, monkeypatch): - """Tests the case when the operation to be applied is - diagonal in the computational basis and the _apply_diagonal_unitary method is used.""" - dev = qml.device("default.qubit.legacy", wires=1) - par = 0.3 - - test_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]]) - wires = 0 - op = qml.PhaseShift(par, wires=wires) - assert op.name not in dev._apply_ops - - # Set the internal _apply_diagonal_unitary - history = [] - mock_apply_diag = lambda state, matrix, wires: history.append((state, matrix, wires)) - with monkeypatch.context() as m: - m.setattr(dev.target_device, "_apply_diagonal_unitary", mock_apply_diag) - assert dev._apply_diagonal_unitary == mock_apply_diag - - dev._apply_operation(test_state, op) - - res_state, res_mat, res_wires = history[0] - - assert np.allclose(res_state, test_state) - assert np.allclose(res_mat, np.diag(op.matrix())) - assert np.allclose(res_wires, wires) - - def test_apply_einsum_case_broadcasted(self, monkeypatch): - """Tests the case when np.einsum is used to apply an operation in - default.qubit.""" - dev = qml.device("default.qubit.legacy", wires=1) - - test_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]]) - wires = 0 - - # Redefine the S gate so that it is an example for a one-qubit gate - # that is not registered in the diagonal_in_z_basis attribute - # pylint: disable=too-few-public-methods - class TestSGate(qml.operation.Operation): - num_wires = 1 - - # pylint: disable=unused-argument - @staticmethod - def compute_matrix(*params, **hyperparams): - return np.array([[1, 0], [0, 1j]]) - - dev.operations.add("TestSGate") - op = TestSGate(wires=wires) - - assert op.name in dev.operations - assert op.name not in dev._apply_ops - - # Set the internal _apply_unitary_einsum - history = [] - mock_apply_einsum = lambda state, matrix, wires: history.append((state, matrix, wires)) - with monkeypatch.context() as m: - m.setattr(dev.target_device, "_apply_unitary_einsum", mock_apply_einsum) - - dev._apply_operation(test_state, op) - - res_state, res_mat, res_wires = history[0] - - assert np.allclose(res_state, test_state) - assert np.allclose(res_mat, op.matrix()) - assert np.allclose(res_wires, wires) - - def test_apply_tensordot_case_broadcasted(self, monkeypatch): - """Tests the case when np.tensordot is used to apply an operation in - default.qubit.""" - dev = qml.device("default.qubit.legacy", wires=3) - - test_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]]) - wires = [0, 1, 2] - - # Redefine the Toffoli gate so that it is an example for a gate with - # more than two wires - # pylint: disable=too-few-public-methods - class TestToffoli(qml.operation.Operation): - num_wires = 3 - - # pylint: disable=unused-argument - @staticmethod - def compute_matrix(*params, **hyperparams): - return Toffoli - - dev.operations.add("TestToffoli") - op = TestToffoli(wires=wires) - - assert op.name in dev.operations - assert op.name not in dev._apply_ops - - # Set the internal _apply_unitary_tensordot - history = [] - mock_apply_tensordot = lambda state, matrix, wires: history.append((state, matrix, wires)) - - with monkeypatch.context() as m: - m.setattr(dev.target_device, "_apply_unitary", mock_apply_tensordot) - - dev._apply_operation(test_state, op) - - res_state, res_mat, res_wires = history[0] - - assert np.allclose(res_state, test_state) - assert np.allclose(res_mat, op.matrix()) - assert np.allclose(res_wires, wires) - - def test_identity_skipped_broadcasted(self, mocker): - """Test that applying the identity operation does not perform any additional computations.""" - dev = qml.device("default.qubit.legacy", wires=1) - - starting_state = np.array([[1, 0], [INVSQ2, INVSQ2], [0, 1]]) - op = qml.Identity(0) - - spy_diagonal = mocker.spy(dev.target_device, "_apply_diagonal_unitary") - spy_einsum = mocker.spy(dev.target_device, "_apply_unitary_einsum") - spy_unitary = mocker.spy(dev.target_device, "_apply_unitary") - - res = dev._apply_operation(starting_state, op) - assert res is starting_state - - spy_diagonal.assert_not_called() - spy_einsum.assert_not_called() - spy_unitary.assert_not_called() - - -class TestHamiltonianSupportBroadcasted: - """Tests the devices' native support for Hamiltonian observables.""" - - def test_do_not_split_analytic_broadcasted(self, mocker): - """Tests that the Hamiltonian is not split for shots=None.""" - dev = qml.device("default.qubit.legacy", wires=2) - Ham = qml.Hamiltonian(np.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)]) - - @qml.qnode(dev, diff_method="parameter-shift", interface=None) - def circuit(): - qml.RX(np.zeros(5), 0) # Broadcast the state by applying a broadcasted identity - return qml.expval(Ham) - - spy = mocker.spy(dev.target_device, "expval") - - circuit() - # evaluated one expval altogether - assert spy.call_count == 1 - - def test_split_finite_shots_broadcasted(self, mocker): - """Tests that the Hamiltonian is split for finite shots.""" - dev = qml.device("default.qubit.legacy", wires=2, shots=10) - spy = mocker.spy(dev.target_device, "expval") - - ham = qml.Hamiltonian(np.array([0.1, 0.2]), [qml.PauliX(0), qml.PauliZ(1)]) - - @qml.qnode(dev) - def circuit(): - qml.RX(np.zeros(5), 0) # Broadcast the state by applying a broadcasted identity - return qml.expval(ham) - - circuit() - - # evaluated one expval per Pauli observable - assert spy.call_count == 2 - - -original_capabilities = qml.devices.DefaultQubitLegacy.capabilities() - - -@pytest.fixture(scope="function", name="mock_default_qubit") -def mock_default_qubit_fixture(monkeypatch): - """A function to create a mock device that mocks the broadcasting support flag - to be False, so that default support via broadcast_expand transform can be tested""" - - # pylint: disable=unused-argument - def overwrite_support(*cls): - capabilities = original_capabilities.copy() - capabilities.update(supports_broadcasting=False) - return capabilities - - with monkeypatch.context() as m: - m.setattr(qml.devices.DefaultQubitLegacy, "capabilities", overwrite_support) - - def get_default_qubit(wires=1, shots=None): - dev = qml.devices.DefaultQubitLegacy(wires=wires, shots=shots) - return dev - - yield get_default_qubit - - -@pytest.mark.parametrize("shots", [None, 100000]) -class TestBroadcastingSupportViaExpansion: - """Tests that the device correctly makes use of ``broadcast_expand`` to - execute broadcasted tapes if its capability to execute broadcasted tapes - is artificially deactivated.""" - - @pytest.mark.parametrize("x", [0.2, np.array([0.1, 0.6, 0.3]), np.array([0.1])]) - def test_with_single_broadcasted_par(self, x, shots, mock_default_qubit): - """Test that broadcasting on a circuit with a - single parametrized operation works.""" - dev = mock_default_qubit(wires=2, shots=shots) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - circuit.construct((np.array(x),), {}) - out = circuit(np.array(x)) - - assert circuit.device.num_executions == (1 if isinstance(x, float) else len(x)) - tol = 1e-10 if shots is None else 1e-2 - assert qml.math.allclose(out, qml.math.cos(x), atol=tol, rtol=0) - - @pytest.mark.parametrize( - "x, y", [(0.2, np.array([0.4])), (np.array([0.1, 5.1]), np.array([0.1, -0.3]))] - ) - def test_with_multiple_pars(self, x, y, shots, mock_default_qubit): - """Test that broadcasting on a circuit with a - single parametrized operation works.""" - dev = mock_default_qubit(wires=2, shots=shots) - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(x, y): - qml.RX(x, wires=0) - qml.RX(y, wires=1) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - - out = circuit(x, y) - expected = qml.math.stack([qml.math.cos(x) * qml.math.ones_like(y), -qml.math.sin(y)]) - - assert circuit.device.num_executions == len(y) - tol = 1e-10 if shots is None else 1e-2 - - assert qml.math.allclose(out[0], expected[0], atol=tol, rtol=0) - assert qml.math.allclose(out[1], expected[1], atol=tol, rtol=0) - - @pytest.mark.parametrize( - "x, y", [(0.2, np.array([0.4])), (np.array([0.1, 5.1]), np.array([0.1, -0.3]))] - ) - def test_with_Hamiltonian(self, x, y, shots, mock_default_qubit): - """Test that broadcasting on a circuit with a - single parametrized operation works.""" - dev = mock_default_qubit(wires=2, shots=shots) - - Ham = qml.Hamiltonian([0.3, 0.9], [qml.PauliZ(0), qml.PauliY(1)]) - Ham.compute_grouping() - - @qml.qnode(dev, diff_method="parameter-shift") - def circuit(x, y): - qml.RX(x, wires=0) - qml.RX(y, wires=1) - return qml.expval(Ham) - - out = circuit(x, y) - expected = 0.3 * qml.math.cos(x) * qml.math.ones_like(y) - 0.9 * qml.math.sin(y) - - assert circuit.device.num_executions == len(y) - tol = 1e-10 if shots is None else 1e-2 - assert qml.math.allclose(out, expected, atol=tol, rtol=0) diff --git a/tests/devices/test_default_qutrit.py b/tests/devices/test_default_qutrit.py index ee860f01af9..6d718c45f18 100644 --- a/tests/devices/test_default_qutrit.py +++ b/tests/devices/test_default_qutrit.py @@ -1466,7 +1466,6 @@ def circuit(x): expected = -np.sin(p) - assert circuit.gradient_fn == "backprop" assert np.isclose(circuit(p), expected, atol=tol, rtol=0) def test_correct_state(self, tol, use_jit): @@ -1658,7 +1657,6 @@ def circuit(x): expected = -np.sin(p) - assert circuit.gradient_fn == "backprop" assert np.isclose(circuit(p), expected, atol=tol, rtol=0) def test_correct_state(self, tol): @@ -1839,7 +1837,6 @@ def circuit(x): expected = -np.sin(p) - assert circuit.gradient_fn == "backprop" assert np.isclose(circuit(p), expected, atol=tol, rtol=0) def test_correct_state(self, tol): diff --git a/tests/devices/test_default_qutrit_mixed.py b/tests/devices/test_default_qutrit_mixed.py index f2613421904..5178e1c800a 100644 --- a/tests/devices/test_default_qutrit_mixed.py +++ b/tests/devices/test_default_qutrit_mixed.py @@ -1720,7 +1720,7 @@ def test_differentiation_jax( # pylint: disable=too-many-arguments if use_jit: diff_func = jax.jit(diff_func) jac = jax.jacobian(diff_func, args_to_diff)(relaxations, misclassifications) - assert np.allclose(jac, expected) + assert qml.math.allclose(jac, expected, rtol=0.05) @pytest.mark.torch @pytest.mark.parametrize("relaxations, misclassifications, expected", diff_parameters) diff --git a/tests/devices/test_device.py b/tests/devices/test_legacy_device.py similarity index 99% rename from tests/devices/test_device.py rename to tests/devices/test_legacy_device.py index c05a5e76309..c08d9411dca 100644 --- a/tests/devices/test_device.py +++ b/tests/devices/test_legacy_device.py @@ -22,7 +22,7 @@ import pytest import pennylane as qml -from pennylane import Device +from pennylane.devices import LegacyDevice as Device from pennylane.wires import Wires mock_device_paulis = ["PauliX", "PauliY", "PauliZ"] @@ -188,6 +188,12 @@ def get_device(wires=1): yield get_device +def test_deprecated_access(): + """Test that accessing via top-level is deprecated.""" + with pytest.warns(qml.PennyLaneDeprecationWarning, match="Device will no longer be accessible"): + qml.Device # pylint: disable=pointless-statement + + # pylint: disable=pointless-statement def test_invalid_attribute_in_devices_raises_error(): with pytest.raises(AttributeError, match="'pennylane.devices' has no attribute 'blabla'"): @@ -1151,7 +1157,7 @@ class TestGrouping: """Tests for the use_grouping option for devices.""" # pylint: disable=too-few-public-methods, unused-argument, missing-function-docstring, missing-class-docstring - class SomeDevice(qml.Device): + class SomeDevice(qml.devices.LegacyDevice): name = "" short_name = "" pennylane_requires = "" diff --git a/tests/devices/test_legacy_facade.py b/tests/devices/test_legacy_facade.py index dac63c2864e..4af1ae1e77a 100644 --- a/tests/devices/test_legacy_facade.py +++ b/tests/devices/test_legacy_facade.py @@ -21,7 +21,6 @@ import pytest import pennylane as qml -from pennylane.devices.default_qubit_autograd import DefaultQubitAutograd from pennylane.devices.execution_config import ExecutionConfig from pennylane.devices.legacy_facade import ( LegacyDeviceFacade, @@ -401,80 +400,7 @@ class BackpropDevice(DummyDevice): dev = LegacyDeviceFacade(BackpropDevice(wires=2, shots=None)) - x = qml.numpy.array(0.1) - tape = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.Z(0))]) - assert dev.supports_derivatives(qml.devices.ExecutionConfig(gradient_method="backprop")) - assert dev._create_temp_device((tape,)) is dev.target_device config = qml.devices.ExecutionConfig(gradient_method="backprop", use_device_gradient=True) assert dev.preprocess(config)[1] is config # unchanged - - def test_backprop_has_passthru_devices(self): - """Test that backprop is supported if the device has passthru devices.""" - - class BackpropDevice(DummyDevice): - - _capabilities = {"passthru_devices": {"autograd": "default.qubit.autograd"}} - - dev = LegacyDeviceFacade(BackpropDevice(shots=None)) - - x = qml.numpy.array(0.1) - tape = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.Z(0))]) - assert dev.supports_derivatives() - assert dev.supports_derivatives(ExecutionConfig(gradient_method="backprop")) - assert dev.supports_derivatives(ExecutionConfig(gradient_method="backprop"), tape) - - config = qml.devices.ExecutionConfig(gradient_method="backprop", use_device_gradient=True) - assert dev.preprocess(config)[1] is config # unchanged - - with pytest.warns(qml.PennyLaneDeprecationWarning, match="switching of devices"): - tmp_device = dev._create_temp_device((tape,)) - assert tmp_device.short_name == "default.qubit.autograd" - - def test_backprop_passthru_device_self(self): - """Test that the temporary device is the original device if the passthru device is itself.""" - - class BackpropSelfDevice(DummyDevice): - - short_name = "BackpropSelfDevice" - - _capabilities = {"passthru_devices": {"autograd": "BackpropSelfDevice"}} - - dev = LegacyDeviceFacade(BackpropSelfDevice(wires=2)) - - x = qml.numpy.array(0.1) - tape = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.Z(0))]) - tmp_dev = dev._create_temp_device((tape,)) - assert tmp_dev is dev.target_device - - def test_passthru_device_does_not_exist(self): - """Test that if backprop is requested for a device that does not support it, a device error is raised.""" - - x = qml.numpy.array(0.1) - tape = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.Z(0))]) - - dev = LegacyDeviceFacade(DummyDevice(wires=2)) - config = qml.devices.ExecutionConfig(gradient_method="backprop") - with pytest.raises(qml.DeviceError, match=r"does not support backpropagation"): - dev.execute(tape, config) - - @pytest.mark.parametrize("dev_class", (qml.devices.DefaultQubitLegacy, DefaultQubitAutograd)) - def test_backprop_device_substitution(self, dev_class): - """Test that default.qubit.legacy is substituted for a backprop device during backprop execution.""" - - with pytest.warns(qml.PennyLaneDeprecationWarning, match="use 'default.qubit'"): - dq_legacy = dev_class(wires=2) - dev = LegacyDeviceFacade(dq_legacy) - - def f(x): - tape = qml.tape.QuantumScript([qml.RX(x, 0)], [qml.expval(qml.Z(0))]) - return dev.execute(tape, qml.devices.ExecutionConfig(gradient_method="backprop")) - - assert qml.math.allclose(dq_legacy.state, np.array([1, 0, 0, 0])) - - with dev.tracker: - g = qml.grad(f)(qml.numpy.array(0.5)) - assert qml.math.allclose(g, -np.sin(0.5)) - assert dev.tracker.totals["executions"] == 1 - assert not qml.math.allclose(dq_legacy.state, np.array([1, 0, 0, 0])) diff --git a/tests/test_qubit_device.py b/tests/devices/test_qubit_device.py similarity index 99% rename from tests/test_qubit_device.py rename to tests/devices/test_qubit_device.py index 1f02976527a..9edc522a408 100644 --- a/tests/test_qubit_device.py +++ b/tests/devices/test_qubit_device.py @@ -21,8 +21,8 @@ import pytest import pennylane as qml -from pennylane import QubitDevice from pennylane import numpy as pnp +from pennylane.devices import QubitDevice from pennylane.measurements import ( Expectation, ExpectationMP, @@ -165,6 +165,12 @@ def _working_get_batch_size(tensor, expected_shape, expected_size): return None +def test_deprecated_access(): + """Test that accessing via top-level is deprecated.""" + with pytest.warns(qml.PennyLaneDeprecationWarning, match="Device will no longer be accessible"): + qml.QubitDevice # pylint: disable=pointless-statement + + def test_notimplemented_circuit_hash(mock_qubit_device): """Test that the circuit hash property is not implemented""" dev = mock_qubit_device() @@ -1508,11 +1514,7 @@ class TestResourcesTracker: Resources(2, 6, {"Hadamard": 3, "RX": 2, "CNOT": 1}, {1: 5, 2: 1}, 4, Shots((10, 10, 50))), ) # Resources(wires, gates, gate_types, gate_sizes, depth, shots) - devices = ( - "default.qubit.legacy", - "default.qubit.autograd", - "default.qubit.jax", - ) + devices = ("default.qubit.legacy",) @pytest.mark.all_interfaces @pytest.mark.parametrize("dev_name", devices) @@ -1659,7 +1661,7 @@ def test_generate_basis_states(): def test_samples_to_counts_all_outomces(): """Test that _samples_to_counts can handle counts with all outcomes.""" - class DummyQubitDevice(qml.QubitDevice): + class DummyQubitDevice(qml.devices.QubitDevice): author = None name = "bla" @@ -1680,7 +1682,7 @@ def apply(self, operations, **kwargs): def test_no_adjoint_jacobian_errors(): """Test that adjoint_jacobian errors with batching and shot vectors""" - class DummyQubitDevice(qml.QubitDevice): + class DummyQubitDevice(qml.devices.QubitDevice): author = None name = "bla" diff --git a/tests/test_qutrit_device.py b/tests/devices/test_qutrit_device.py similarity index 99% rename from tests/test_qutrit_device.py rename to tests/devices/test_qutrit_device.py index 6291a3e9d71..8799d75234b 100644 --- a/tests/test_qutrit_device.py +++ b/tests/devices/test_qutrit_device.py @@ -22,8 +22,8 @@ from scipy.stats import unitary_group import pennylane as qml -from pennylane import QubitDevice, QutritDevice from pennylane import numpy as pnp +from pennylane.devices import QubitDevice, QutritDevice from pennylane.measurements import ( Counts, CountsMP, @@ -142,6 +142,12 @@ def get_qutrit_device(wires=1): # TODO: Add tests for expval, var after observables are added +def test_deprecated_access(): + """Test that accessing via top-level is deprecated.""" + with pytest.warns(qml.PennyLaneDeprecationWarning, match="Device will no longer be accessible"): + qml.QutritDevice # pylint: disable=pointless-statement + + class TestOperations: """Tests the logic related to operations""" diff --git a/tests/fermi/test_fermionic.py b/tests/fermi/test_fermionic.py index 368a6bb1fb4..dc61295115a 100644 --- a/tests/fermi/test_fermionic.py +++ b/tests/fermi/test_fermionic.py @@ -21,17 +21,37 @@ import pennylane as qml from pennylane import numpy as pnp -from pennylane.fermi.fermionic import FermiSentence, FermiWord, _to_string, from_string +from pennylane.fermi.fermionic import ( + FermiA, + FermiC, + FermiSentence, + FermiWord, + _to_string, + from_string, +) # pylint: disable=too-many-public-methods fw1 = FermiWord({(0, 0): "+", (1, 1): "-"}) +fw1_dag = FermiWord({(0, 1): "+", (1, 0): "-"}) + fw2 = FermiWord({(0, 0): "+", (1, 0): "-"}) +fw2_dag = FermiWord({(0, 0): "+", (1, 0): "-"}) + fw3 = FermiWord({(0, 0): "+", (1, 3): "-", (2, 0): "+", (3, 4): "-"}) +fw3_dag = FermiWord({(0, 4): "+", (1, 0): "-", (2, 3): "+", (3, 0): "-"}) + fw4 = FermiWord({}) +fw4_dag = FermiWord({}) + fw5 = FermiWord({(0, 10): "+", (1, 30): "-", (2, 0): "+", (3, 400): "-"}) +fw5_dag = FermiWord({(0, 400): "+", (1, 0): "-", (2, 30): "+", (3, 10): "-"}) + fw6 = FermiWord({(0, 10): "+", (1, 30): "+", (2, 0): "-", (3, 400): "-"}) +fw6_dag = FermiWord({(0, 400): "+", (1, 0): "+", (2, 30): "-", (3, 10): "-"}) + fw7 = FermiWord({(0, 10): "-", (1, 30): "+", (2, 0): "-", (3, 400): "+"}) +fw7_dag = FermiWord({(0, 400): "-", (1, 0): "+", (2, 30): "-", (3, 10): "+"}) class TestFermiWord: @@ -147,6 +167,24 @@ def test_to_mat_error(self): with pytest.raises(ValueError, match="n_orbitals cannot be smaller than 2"): fw1.to_mat(n_orbitals=1) + tup_fw_dag = ( + (fw1, fw1_dag), + (fw2, fw2_dag), + (fw3, fw3_dag), + (fw4, fw4_dag), + (fw5, fw5_dag), + (fw6, fw6_dag), + (fw7, fw7_dag), + (FermiA(0), FermiC(0)), + (FermiC(0), FermiA(0)), + (FermiA(1), FermiC(1)), + (FermiC(1), FermiA(1)), + ) + + @pytest.mark.parametrize("fw, fw_dag", tup_fw_dag) + def test_adjoint(self, fw, fw_dag): + assert fw.adjoint() == fw_dag + class TestFermiWordArithmetic: WORDS_MUL = ( @@ -458,13 +496,29 @@ def test_array_must_not_exceed_length_1(self, method_name): fs1 = FermiSentence({fw1: 1.23, fw2: 4j, fw3: -0.5}) +fs1_dag = FermiSentence({fw1_dag: 1.23, fw2_dag: -4j, fw3_dag: -0.5}) + fs2 = FermiSentence({fw1: -1.23, fw2: -4j, fw3: 0.5}) +fs2_dag = FermiSentence({fw1_dag: -1.23, fw2_dag: 4j, fw3_dag: 0.5}) + fs1_hamiltonian = FermiSentence({fw1: 1.23, fw2: 4, fw3: -0.5}) +fs1_hamiltonian_dag = FermiSentence({fw1_dag: 1.23, fw2_dag: 4, fw3_dag: -0.5}) + fs2_hamiltonian = FermiSentence({fw1: -1.23, fw2: -4, fw3: 0.5}) +fs2_hamiltonian_dag = FermiSentence({fw1_dag: -1.23, fw2_dag: -4, fw3_dag: 0.5}) + fs3 = FermiSentence({fw3: -0.5, fw4: 1}) +fs3_dag = FermiSentence({fw3_dag: -0.5, fw4_dag: 1}) + fs4 = FermiSentence({fw4: 1}) +fs4_dag = FermiSentence({fw4_dag: 1}) + fs5 = FermiSentence({}) +fs5_dag = FermiSentence({}) + fs6 = FermiSentence({fw1: 1.2, fw2: 3.1}) +fs6_dag = FermiSentence({fw1_dag: 1.2, fw2_dag: 3.1}) + fs7 = FermiSentence( { FermiWord({(0, 0): "+", (1, 1): "-"}): 1.23, # a+(0) a(1) @@ -652,6 +706,21 @@ def test_to_mat_error(self): with pytest.raises(ValueError, match="n_orbitals cannot be smaller than 3"): fs7.to_mat(n_orbitals=2) + fs_dag_tup = ( + (fs1, fs1_dag), + (fs2, fs2_dag), + (fs3, fs3_dag), + (fs4, fs4_dag), + (fs5, fs5_dag), + (fs6, fs6_dag), + (fs1_hamiltonian, fs1_hamiltonian_dag), + (fs2_hamiltonian, fs2_hamiltonian_dag), + ) + + @pytest.mark.parametrize("fs, fs_dag", fs_dag_tup) + def test_adjoint(self, fs, fs_dag): + assert fs.adjoint() == fs_dag + class TestFermiSentenceArithmetic: tup_fs_mult = ( # computed by hand diff --git a/tests/gradients/core/test_adjoint_metric_tensor.py b/tests/gradients/core/test_adjoint_metric_tensor.py index 4f917caeea4..23910fc0862 100644 --- a/tests/gradients/core/test_adjoint_metric_tensor.py +++ b/tests/gradients/core/test_adjoint_metric_tensor.py @@ -282,8 +282,7 @@ def circuit(*params): @pytest.mark.jax @pytest.mark.skip("JAX does not support forward pass execution of the metric tensor.") - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_correct_output_tape_jax(self, dev_name, ansatz, params): + def test_correct_output_tape_jax(self, ansatz, params): """Test that the output is correct when using JAX and calling the adjoint metric tensor directly on a tape.""" @@ -291,7 +290,7 @@ def test_correct_output_tape_jax(self, dev_name, ansatz, params): expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params) j_params = tuple(jax.numpy.array(p) for p in params) - dev = qml.device(dev_name, wires=self.num_wires) + dev = qml.device("default.qubit", wires=self.num_wires) @qml.qnode(dev, interface="jax") def circuit(*params): @@ -311,8 +310,7 @@ def circuit(*params): @pytest.mark.torch @pytest.mark.parametrize("interface", interfaces) - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_correct_output_tape_torch(self, ansatz, params, interface, dev_name): + def test_correct_output_tape_torch(self, ansatz, params, interface): """Test that the output is correct when using Torch and calling the adjoint metric tensor directly on a tape.""" @@ -320,7 +318,7 @@ def test_correct_output_tape_torch(self, ansatz, params, interface, dev_name): expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params) t_params = tuple(torch.tensor(p, requires_grad=True) for p in params) - dev = qml.device(dev_name, wires=self.num_wires) + dev = qml.device("default.qubit", wires=self.num_wires) @qml.qnode(dev, interface=interface) def circuit(*params): @@ -340,8 +338,7 @@ def circuit(*params): @pytest.mark.tf @pytest.mark.parametrize("interface", interfaces) - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_correct_output_tape_tf(self, ansatz, params, interface, dev_name): + def test_correct_output_tape_tf(self, ansatz, params, interface): """Test that the output is correct when using TensorFlow and calling the adjoint metric tensor directly on a tape.""" @@ -349,7 +346,7 @@ def test_correct_output_tape_tf(self, ansatz, params, interface, dev_name): expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params) t_params = tuple(tf.Variable(p) for p in params) - dev = qml.device(dev_name, wires=self.num_wires) + dev = qml.device("default.qubit", wires=self.num_wires) @qml.qnode(dev, interface=interface) def circuit(*params): @@ -431,8 +428,7 @@ def circuit(*params): @pytest.mark.torch @pytest.mark.parametrize("ansatz, params", list(zip(fubini_ansatze, fubini_params))) @pytest.mark.parametrize("interface", interfaces) - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_correct_output_qnode_torch(self, ansatz, params, interface, dev_name): + def test_correct_output_qnode_torch(self, ansatz, params, interface): """Test that the output is correct when using Torch and calling the adjoint metric tensor on a QNode.""" @@ -440,7 +436,7 @@ def test_correct_output_qnode_torch(self, ansatz, params, interface, dev_name): expected = autodiff_metric_tensor(ansatz, self.num_wires)(*params) t_params = tuple(torch.tensor(p, requires_grad=True, dtype=torch.float64) for p in params) - dev = qml.device(dev_name, wires=self.num_wires) + dev = qml.device("default.qubit", wires=self.num_wires) @qml.qnode(dev, interface=interface) def circuit(*params): diff --git a/tests/gradients/core/test_hadamard_gradient.py b/tests/gradients/core/test_hadamard_gradient.py index 89e05fd6fab..ab225109fa9 100644 --- a/tests/gradients/core/test_hadamard_gradient.py +++ b/tests/gradients/core/test_hadamard_gradient.py @@ -1032,12 +1032,10 @@ class TestHadamardTestGradDiff: """Test that the transform is differentiable""" @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, dev_name): + def test_autograd(self): """Tests that the output of the hadamard gradient transform can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=3) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=3) params = np.array([0.543, -0.654], requires_grad=True) def cost_fn_hadamard(x): @@ -1050,7 +1048,7 @@ def cost_fn_hadamard(x): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.hadamard_grad(tape) - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) return qml.math.stack(jac) def cost_fn_param_shift(x): @@ -1063,7 +1061,7 @@ def cost_fn_param_shift(x): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.param_shift(tape) - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) return qml.math.stack(jac) res_hadamard = qml.jacobian(cost_fn_hadamard)(params) @@ -1071,14 +1069,12 @@ def cost_fn_param_shift(x): assert np.allclose(res_hadamard, res_param_shift) @pytest.mark.tf - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf(self, dev_name): + def test_tf(self): """Tests that the output of the hadamard gradient transform can be differentiated using TF, yielding second derivatives.""" import tensorflow as tf - dev = qml.device(dev_name, wires=3) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=3) params = tf.Variable([0.543, -0.654], dtype=tf.float64) with tf.GradientTape() as t_h: @@ -1091,7 +1087,7 @@ def test_tf(self, dev_name): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.hadamard_grad(tape) - jac_h = fn(execute_fn(tapes)) + jac_h = fn(dev.execute(tapes)) jac_h = qml.math.stack(jac_h) with tf.GradientTape() as t_p: @@ -1104,7 +1100,7 @@ def test_tf(self, dev_name): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.param_shift(tape) - jac_p = fn(execute_fn(tapes)) + jac_p = fn(dev.execute(tapes)) jac_p = qml.math.stack(jac_p) res_hadamard = t_h.jacobian(jac_h, params) @@ -1113,14 +1109,12 @@ def test_tf(self, dev_name): assert np.allclose(res_hadamard, res_param_shift) @pytest.mark.torch - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_torch(self, dev_name): + def test_torch(self): """Tests that the output of the hadamard gradient transform can be differentiated using Torch, yielding second derivatives.""" import torch - dev = qml.device(dev_name, wires=3) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=3) params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True) def cost_h(x): @@ -1133,8 +1127,7 @@ def cost_h(x): tape = qml.tape.QuantumScript.from_queue(q) tapes, fn = qml.gradients.hadamard_grad(tape) - jac = fn(execute_fn(tapes)) - return jac + return fn(dev.execute(tapes)) def cost_p(x): with qml.queuing.AnnotatedQueue() as q: @@ -1146,8 +1139,7 @@ def cost_p(x): tape = qml.tape.QuantumScript.from_queue(q) tapes, fn = qml.gradients.param_shift(tape) - jac = fn(execute_fn(tapes)) - return jac + return fn(dev.execute(tapes)) res_hadamard = torch.autograd.functional.jacobian(cost_h, params) res_param_shift = torch.autograd.functional.jacobian(cost_p, params) @@ -1156,15 +1148,13 @@ def cost_p(x): assert np.allclose(res_hadamard[1].detach(), res_param_shift[1].detach()) @pytest.mark.jax - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_jax(self, dev_name): + def test_jax(self): """Tests that the output of the hadamard gradient transform can be differentiated using JAX, yielding second derivatives.""" import jax from jax import numpy as jnp - dev = qml.device(dev_name, wires=3) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=3) params = jnp.array([0.543, -0.654]) def cost_h(x): @@ -1178,8 +1168,7 @@ def cost_h(x): tape.trainable_params = {0, 1} tapes, fn = qml.gradients.hadamard_grad(tape) - jac = fn(execute_fn(tapes)) - return jac + return fn(dev.execute(tapes)) def cost_p(x): with qml.queuing.AnnotatedQueue() as q: @@ -1192,8 +1181,7 @@ def cost_p(x): tape.trainable_params = {0, 1} tapes, fn = qml.gradients.hadamard_grad(tape) - jac = fn(execute_fn(tapes)) - return jac + return fn(dev.execute(tapes)) res_hadamard = jax.jacobian(cost_h)(params) res_param_shift = jax.jacobian(cost_p)(params) diff --git a/tests/gradients/core/test_jvp.py b/tests/gradients/core/test_jvp.py index bf1d63c60b6..e8ad4e1d614 100644 --- a/tests/gradients/core/test_jvp.py +++ b/tests/gradients/core/test_jvp.py @@ -17,6 +17,7 @@ import pennylane as qml from pennylane import numpy as np from pennylane.gradients import param_shift +from pennylane.measurements.shots import Shots _x = np.arange(12).reshape((2, 3, 2)) @@ -653,12 +654,10 @@ class TestJVPGradients: # Include batch_dim!=None cases once #4462 is resolved @pytest.mark.autograd @pytest.mark.parametrize("batch_dim", [None]) # , 1, 3]) - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, tol, dev_name, batch_dim): + def test_autograd(self, tol, batch_dim): """Tests that the output of the JVP transform can be differentiated using autograd.""" - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = np.array([0.543, -0.654], requires_grad=True) if batch_dim is not None: params = np.outer(np.arange(1, 1 + batch_dim), params, requires_grad=True) @@ -671,8 +670,7 @@ def cost_fn(params, tangent): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.jvp(tape, tangent, param_shift) - jvp = fn(execute_fn(tapes)) - return jvp + return fn(dev.execute(tapes)) res = cost_fn(params, tangent) exp = expected_jvp(params, tangent) @@ -685,14 +683,12 @@ def cost_fn(params, tangent): # Include batch_dim!=None cases once #4462 is resolved @pytest.mark.torch @pytest.mark.parametrize("batch_dim", [None]) # , 1, 3]) - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_torch(self, tol, dev_name, batch_dim): + def test_torch(self, tol, batch_dim): """Tests that the output of the JVP transform can be differentiated using Torch.""" import torch - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params_np = np.array([0.543, -0.654], requires_grad=True) if batch_dim is not None: @@ -708,8 +704,7 @@ def cost_fn(params, tangent): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.jvp(tape, tangent, param_shift) - jvp = fn(execute_fn(tapes)) - return jvp + return fn(dev.execute(tapes)) res = cost_fn(params, tangent) exp = expected_jvp(params_np, tangent_np) @@ -723,14 +718,12 @@ def cost_fn(params, tangent): @pytest.mark.tf @pytest.mark.slow @pytest.mark.parametrize("batch_dim", [None]) # , 1, 3]) - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf(self, tol, dev_name, batch_dim): + def test_tf(self, tol, batch_dim): """Tests that the output of the JVP transform can be differentiated using Tensorflow.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params_np = np.array([0.543, -0.654], requires_grad=True) if batch_dim is not None: params_np = np.outer(np.arange(1, 1 + batch_dim), params_np, requires_grad=True) @@ -745,8 +738,7 @@ def cost_fn(params, tangent): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.jvp(tape, tangent, param_shift) - jvp = fn(execute_fn(tapes)) - return jvp + return fn(dev.execute(tapes)) with tf.GradientTape() as t: res = cost_fn(params, tangent) @@ -761,15 +753,13 @@ def cost_fn(params, tangent): # Include batch_dim!=None cases once #4462 is resolved @pytest.mark.jax @pytest.mark.parametrize("batch_dim", [None]) # , 1, 3]) - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_jax(self, tol, dev_name, batch_dim): + def test_jax(self, tol, batch_dim): """Tests that the output of the JVP transform can be differentiated using JAX.""" import jax from jax import numpy as jnp - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit") params_np = np.array([0.543, -0.654], requires_grad=True) if batch_dim is not None: params_np = np.outer(np.arange(1, 1 + batch_dim), params_np, requires_grad=True) @@ -784,8 +774,7 @@ def cost_fn(params, tangent): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.jvp(tape, tangent, param_shift) - jvp = fn(execute_fn(tapes)) - return jvp + return fn(dev.execute(tapes)) res = cost_fn(params, tangent) exp = expected_jvp(params_np, tangent_np) @@ -799,7 +788,8 @@ def cost_fn(params, tangent): class TestBatchJVP: """Tests for the batch JVP function""" - def test_one_tape_no_trainable_parameters(self): + @pytest.mark.parametrize("shots", [Shots(None), Shots(10), Shots([20, 10])]) + def test_one_tape_no_trainable_parameters(self, shots): """A tape with no trainable parameters will simply return None""" dev = qml.device("default.qubit", wires=2) @@ -808,14 +798,14 @@ def test_one_tape_no_trainable_parameters(self): qml.CNOT(wires=[0, 1]) qml.expval(qml.PauliZ(0)) - tape1 = qml.tape.QuantumScript.from_queue(q1) + tape1 = qml.tape.QuantumScript.from_queue(q1, shots=shots) with qml.queuing.AnnotatedQueue() as q2: qml.RX(0.4, wires=0) qml.RX(0.6, wires=0) qml.CNOT(wires=[0, 1]) qml.expval(qml.PauliZ(0)) - tape2 = qml.tape.QuantumScript.from_queue(q2) + tape2 = qml.tape.QuantumScript.from_queue(q2, shots=shots) tape1.trainable_params = {} tape2.trainable_params = {0, 1} @@ -823,16 +813,17 @@ def test_one_tape_no_trainable_parameters(self): tangents = [np.array([1.0, 1.0]), np.array([1.0, 1.0])] v_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, param_shift) - assert len(v_tapes) == 4 # Even though there are 3 parameters, only two contribute # to the JVP, so only 2*2=4 quantum evals + assert len(v_tapes) == 4 res = fn(dev.execute(v_tapes)) assert qml.math.allclose(res[0], np.array(0.0)) assert res[1] is not None - def test_all_tapes_no_trainable_parameters(self): + @pytest.mark.parametrize("shots", [Shots(None), Shots(10), Shots([20, 10])]) + def test_all_tapes_no_trainable_parameters(self, shots): """If all tapes have no trainable parameters all outputs will be None""" with qml.queuing.AnnotatedQueue() as q1: @@ -840,14 +831,14 @@ def test_all_tapes_no_trainable_parameters(self): qml.CNOT(wires=[0, 1]) qml.expval(qml.PauliZ(0)) - tape1 = qml.tape.QuantumScript.from_queue(q1) + tape1 = qml.tape.QuantumScript.from_queue(q1, shots=shots) with qml.queuing.AnnotatedQueue() as q2: qml.RX(0.4, wires=0) qml.RX(0.6, wires=0) qml.CNOT(wires=[0, 1]) qml.expval(qml.PauliZ(0)) - tape2 = qml.tape.QuantumScript.from_queue(q2) + tape2 = qml.tape.QuantumScript.from_queue(q2, shots=shots) tape1.trainable_params = set() tape2.trainable_params = set() diff --git a/tests/gradients/core/test_pulse_gradient.py b/tests/gradients/core/test_pulse_gradient.py index 28dfa294673..b427f5ee23b 100644 --- a/tests/gradients/core/test_pulse_gradient.py +++ b/tests/gradients/core/test_pulse_gradient.py @@ -846,7 +846,6 @@ def test_raises_for_invalid_reorder_fn(self, reorder_fn): @pytest.mark.jax -@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) class TestStochPulseGrad: """Test working cases of stoch_pulse_grad.""" @@ -880,7 +879,7 @@ def sine(p, t): ), ), ) - def test_all_zero_grads(self, dev_name, ops, arg, exp_shapes): # pylint:disable=unused-argument + def test_all_zero_grads(self, ops, arg, exp_shapes): # pylint:disable=unused-argument """Test that a zero gradient is returned when all trainable parameters are identified to have zero gradient in advance.""" import jax @@ -902,7 +901,7 @@ def test_all_zero_grads(self, dev_name, ops, arg, exp_shapes): # pylint:disable assert qml.math.allclose(r, np.zeros(exp_shape)) jax.clear_caches() - def test_some_zero_grads(self, dev_name): + def test_some_zero_grads(self): """Test that a zero gradient is returned for trainable parameters that are identified to have a zero gradient in advance.""" import jax @@ -918,7 +917,7 @@ def test_some_zero_grads(self, dev_name): tapes, fn = stoch_pulse_grad(tape, num_split_times=3) assert len(tapes) == 2 * 3 - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) res = fn(qml.execute(tapes, dev, None)) assert isinstance(res, tuple) and len(res) == 2 assert qml.math.allclose(res[0][0], np.zeros(5)) @@ -927,7 +926,7 @@ def test_some_zero_grads(self, dev_name): @pytest.mark.parametrize("num_split_times", [1, 3]) @pytest.mark.parametrize("t", [2.0, 3, (0.5, 0.6), (0.1, 0.9, 1.2)]) - def test_constant_ry(self, dev_name, num_split_times, t): + def test_constant_ry(self, num_split_times, t): """Test that the derivative of a pulse generated by a constant Hamiltonian, which is a Pauli word, is computed correctly.""" import jax @@ -940,7 +939,7 @@ def test_constant_ry(self, dev_name, num_split_times, t): op = qml.evolve(ham_single_q_const)(params, t) tape = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0))]) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) # Effective rotation parameter p = params[0] * delta_t r = qml.execute([tape], dev, None) @@ -952,7 +951,7 @@ def test_constant_ry(self, dev_name, num_split_times, t): assert qml.math.isclose(res, -2 * jnp.sin(2 * p) * delta_t) jax.clear_caches() - def test_constant_ry_argnum(self, dev_name): + def test_constant_ry_argnum(self): """Test that the derivative of a pulse generated by a constant Hamiltonian, which is a Pauli word, is computed correctly if it is not the only operation in a tape but selected via `argnum`.""" @@ -968,7 +967,7 @@ def test_constant_ry_argnum(self, dev_name): tape = qml.tape.QuantumScript([qml.RY(y, 0), op], [qml.expval(qml.PauliZ(0))]) tape.trainable_params = [1] - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) # Effective rotation parameter p = params[0] * t r = qml.execute([tape], dev, None) @@ -983,7 +982,7 @@ def test_constant_ry_argnum(self, dev_name): @pytest.mark.parametrize("num_split_times", [1, 3]) @pytest.mark.parametrize("t", [2.0, 3, (0.5, 0.6), (0.1, 0.9, 1.2)]) - def test_constant_ry_rescaled(self, dev_name, num_split_times, t): + def test_constant_ry_rescaled(self, num_split_times, t): """Test that the derivative of a pulse generated by a constant Hamiltonian, which is a Pauli sentence, is computed correctly.""" import jax @@ -998,7 +997,7 @@ def test_constant_ry_rescaled(self, dev_name, num_split_times, t): op = qml.evolve(ham_single_q_const)(params, t) tape = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0))]) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) # Prefactor due to the generator being a Pauli sentence prefactor = np.sqrt(0.85) # Effective rotation parameter @@ -1013,7 +1012,7 @@ def test_constant_ry_rescaled(self, dev_name, num_split_times, t): jax.clear_caches() @pytest.mark.parametrize("t", [0.02, (0.5, 0.6)]) - def test_sin_envelope_rz_expval(self, dev_name, t): + def test_sin_envelope_rz_expval(self, t): """Test that the derivative of a pulse with a sine wave envelope is computed correctly when returning an expectation value.""" import jax @@ -1021,7 +1020,7 @@ def test_sin_envelope_rz_expval(self, dev_name, t): T = t if isinstance(t, tuple) else (0, t) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) params = [jnp.array([2.3, -0.245])] ham = self.sine * qml.PauliZ(0) @@ -1052,7 +1051,7 @@ def test_sin_envelope_rz_expval(self, dev_name, t): jax.clear_caches() @pytest.mark.parametrize("t", [0.02, (0.5, 0.6)]) - def test_sin_envelope_rx_probs(self, dev_name, t): + def test_sin_envelope_rx_probs(self, t): """Test that the derivative of a pulse with a sine wave envelope is computed correctly when returning probabilities.""" import jax @@ -1060,7 +1059,7 @@ def test_sin_envelope_rx_probs(self, dev_name, t): T = t if isinstance(t, tuple) else (0, t) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) params = [jnp.array([2.3, -0.245])] ham = self.sine * qml.PauliX(0) @@ -1093,7 +1092,7 @@ def test_sin_envelope_rx_probs(self, dev_name, t): jax.clear_caches() @pytest.mark.parametrize("t", [0.02, (0.5, 0.6)]) - def test_sin_envelope_rx_expval_probs(self, dev_name, t): + def test_sin_envelope_rx_expval_probs(self, t): """Test that the derivative of a pulse with a sine wave envelope is computed correctly when returning expectation.""" import jax @@ -1101,7 +1100,7 @@ def test_sin_envelope_rx_expval_probs(self, dev_name, t): T = t if isinstance(t, tuple) else (0, t) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) params = [jnp.array([2.3, -0.245])] ham = self.sine * qml.PauliX(0) @@ -1138,7 +1137,7 @@ def test_sin_envelope_rx_expval_probs(self, dev_name, t): jax.clear_caches() @pytest.mark.parametrize("t", [0.02, (0.5, 0.6)]) - def test_pwc_envelope_rx(self, dev_name, t): + def test_pwc_envelope_rx(self, t): """Test that the derivative of a pulse generated by a piecewise constant Hamiltonian is computed correctly.""" import jax @@ -1146,7 +1145,7 @@ def test_pwc_envelope_rx(self, dev_name, t): T = t if isinstance(t, tuple) else (0, t) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) params = [jnp.array([0.24, 0.9, -0.1, 2.3, -0.245])] op = qml.evolve(qml.pulse.pwc(t) * qml.PauliZ(0))(params, t) tape = qml.tape.QuantumScript([qml.Hadamard(0), op], [qml.expval(qml.PauliX(0))]) @@ -1168,7 +1167,7 @@ def test_pwc_envelope_rx(self, dev_name, t): jax.clear_caches() @pytest.mark.parametrize("t", [2.0, 3, (0.5, 0.6)]) - def test_constant_commuting(self, dev_name, t): + def test_constant_commuting(self, t): """Test that the derivative of a pulse generated by two constant commuting Hamiltonians is computed correctly.""" import jax @@ -1182,7 +1181,7 @@ def test_constant_commuting(self, dev_name, t): ) tape = qml.tape.QuantumScript([op], [qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))]) - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) r = qml.execute([tape], dev, None) # Effective rotation parameters p = [_p * (T[1] - T[0]) for _p in params] @@ -1199,7 +1198,7 @@ def test_constant_commuting(self, dev_name, t): jax.clear_caches() @pytest.mark.slow - def test_advanced_pulse(self, dev_name): + def test_advanced_pulse(self): """Test the derivative of a more complex pulse.""" import jax import jax.numpy as jnp @@ -1213,7 +1212,7 @@ def test_advanced_pulse(self, dev_name): * qml.dot([1.0, 0.4], [qml.PauliY(0) @ qml.PauliY(1), qml.PauliX(0) @ qml.PauliX(1)]) ) params = [jnp.array(1.51), jnp.array(-0.371), jnp.array([0.2, 0.2, -0.4])] - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) @qml.qnode(dev, interface="jax") def qnode(params): @@ -1237,7 +1236,7 @@ def qnode(params): assert all(qml.math.allclose(r, e, rtol=0.4) for r, e in zip(res, exp_grad)) jax.clear_caches() - def test_randomness(self, dev_name): + def test_randomness(self): """Test that the derivative of a pulse is exactly the same when reusing a seed and that it differs when using a different seed.""" import jax @@ -1268,7 +1267,7 @@ def test_randomness(self, dev_name): qml.assert_equal(op_a_0, op_a_1) qml.assert_equal(op_a_0, op_b) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) res_a_0 = fn_a_0(qml.execute(tapes_a_0, dev, None)) res_a_1 = fn_a_1(qml.execute(tapes_a_1, dev, None)) res_b = fn_b(qml.execute(tapes_b, dev, None)) @@ -1277,7 +1276,7 @@ def test_randomness(self, dev_name): assert not res_a_0 == res_b jax.clear_caches() - def test_two_pulses(self, dev_name): + def test_two_pulses(self): """Test that the derivatives of two pulses in a circuit are computed correctly.""" import jax import jax.numpy as jnp @@ -1288,7 +1287,7 @@ def test_two_pulses(self, dev_name): ham_1 = qml.dot([0.3, jnp.polyval], [qml.PauliZ(0), qml.PauliY(0) @ qml.PauliY(1)]) params_0 = [jnp.array(1.51), jnp.array(-0.371)] params_1 = [jnp.array([0.2, 0.2, -0.4])] - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) @qml.qnode(dev, interface="jax") def qnode(params_0, params_1): @@ -1318,13 +1317,13 @@ def qnode(params_0, params_1): (qml.Hamiltonian([0.25, 1.2], [qml.PauliX(0), qml.PauliX(0) @ qml.PauliZ(1)]), 8, 1.45), ], ) - def test_with_jit(self, dev_name, generator, exp_num_tapes, prefactor): + def test_with_jit(self, generator, exp_num_tapes, prefactor): """Test that the stochastic parameter-shift rule works with JITting.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=len(generator.wires)) + dev = qml.device("default.qubit", wires=len(generator.wires)) T = (0.2, 0.5) ham_single_q_const = qml.dot([qml.pulse.constant], [generator]) meas = [qml.expval(qml.PauliZ(0))] @@ -1349,7 +1348,7 @@ def fun(params): jax.clear_caches() @pytest.mark.parametrize("shots", [None, 100]) - def test_shots_attribute(self, dev_name, shots): # pylint:disable=unused-argument + def test_shots_attribute(self, shots): # pylint:disable=unused-argument """Tests that the shots attribute is copied to the new tapes""" tape = qml.tape.QuantumTape([], [qml.expval(qml.PauliZ(0)), qml.probs([1, 2])], shots=shots) with pytest.warns(UserWarning, match="Attempted to compute the gradient of a tape with no"): @@ -1359,14 +1358,13 @@ def test_shots_attribute(self, dev_name, shots): # pylint:disable=unused-argume @pytest.mark.jax -@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) class TestStochPulseGradQNode: """Test that stoch_pulse_grad integrates correctly with QNodes.""" - def test_raises_for_application_to_qnodes(self, dev_name): + def test_raises_for_application_to_qnodes(self): """Test that an error is raised when applying ``stoch_pulse_grad`` to a QNode directly.""" - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) ham_single_q_const = qml.pulse.constant * qml.PauliY(0) @qml.qnode(dev, interface="jax") @@ -1380,14 +1378,14 @@ def circuit(params): # TODO: include the following tests when #4225 is resolved. @pytest.mark.skip("Applying this gradient transform to QNodes directly is not supported.") - def test_qnode_expval_single_par(self, dev_name): + def test_qnode_expval_single_par(self): """Test that a simple qnode that returns an expectation value can be differentiated with stoch_pulse_grad.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = 0.2 ham_single_q_const = qml.pulse.constant * qml.PauliY(0) @@ -1398,7 +1396,6 @@ def circuit(params): params = jnp.array(0.4) with qml.Tracker(dev) as tracker: - _match = "stochastic pulse parameter-shift .* scalar pulse parameters." grad = stoch_pulse_grad(circuit, num_split_times=2)(params) p = params * T @@ -1408,21 +1405,19 @@ def circuit(params): @pytest.mark.jax -@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) class TestStochPulseGradIntegration: """Test that stoch_pulse_grad integrates correctly with QNodes and ML interfaces.""" @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1), ([100, 99], 0.1)]) @pytest.mark.parametrize("num_split_times", [1, 2]) - def test_simple_qnode_expval(self, dev_name, num_split_times, shots, tol): + def test_simple_qnode_expval(self, num_split_times, shots, tol): """Test that a simple qnode that returns an expectation value can be differentiated with stoch_pulse_grad.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - key = "prng_key" if dev_name == "default.qubit.jax" else "seed" - dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)}) + dev = qml.device("default.qubit", wires=1, shots=shots, seed=jax.random.PRNGKey(74)) T = 0.2 ham_single_q_const = qml.pulse.constant * qml.PauliY(0) @@ -1443,15 +1438,14 @@ def circuit(params): @pytest.mark.slow @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1), ([100, 99], 0.1)]) @pytest.mark.parametrize("num_split_times", [1, 2]) - def test_simple_qnode_expval_two_evolves(self, dev_name, num_split_times, shots, tol): + def test_simple_qnode_expval_two_evolves(self, num_split_times, shots, tol): """Test that a simple qnode that returns an expectation value can be differentiated with stoch_pulse_grad.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - key = "prng_key" if dev_name == "default.qubit.jax" else "seed" - dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)}) + dev = qml.device("default.qubit", wires=1, shots=shots, seed=jax.random.PRNGKey(74)) T_x = 0.1 T_y = 0.2 ham_x = qml.pulse.constant * qml.PauliX(0) @@ -1475,15 +1469,14 @@ def circuit(params): @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1), ([100, 99], 0.1)]) @pytest.mark.parametrize("num_split_times", [1, 2]) - def test_simple_qnode_probs(self, dev_name, num_split_times, shots, tol): + def test_simple_qnode_probs(self, num_split_times, shots, tol): """Test that a simple qnode that returns an probabilities can be differentiated with stoch_pulse_grad.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - key = "prng_key" if dev_name == "default.qubit.jax" else "seed" - dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)}) + dev = qml.device("default.qubit", wires=1, shots=shots, seed=jax.random.PRNGKey(74)) T = 0.2 ham_single_q_const = qml.pulse.constant * qml.PauliY(0) @@ -1503,15 +1496,14 @@ def circuit(params): @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1), ([100, 100], 0.1)]) @pytest.mark.parametrize("num_split_times", [1, 2]) - def test_simple_qnode_probs_expval(self, dev_name, num_split_times, shots, tol): + def test_simple_qnode_probs_expval(self, num_split_times, shots, tol): """Test that a simple qnode that returns an probabilities can be differentiated with stoch_pulse_grad.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - key = "prng_key" if dev_name == "default.qubit.jax" else "seed" - dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)}) + dev = qml.device("default.qubit", wires=1, shots=shots, seed=jax.random.PRNGKey(74)) T = 0.2 ham_single_q_const = qml.pulse.constant * qml.PauliY(0) @@ -1538,13 +1530,13 @@ def circuit(params): @pytest.mark.xfail @pytest.mark.parametrize("num_split_times", [1, 2]) @pytest.mark.parametrize("time_interface", ["python", "numpy", "jax"]) - def test_simple_qnode_jit(self, dev_name, num_split_times, time_interface): + def test_simple_qnode_jit(self, num_split_times, time_interface): """Test that a simple qnode can be differentiated with stoch_pulse_grad.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = {"python": 0.2, "numpy": np.array(0.2), "jax": jnp.array(0.2)}[time_interface] ham_single_q_const = qml.pulse.constant * qml.PauliY(0) @@ -1563,7 +1555,7 @@ def circuit(params, T=None): jax.clear_caches() @pytest.mark.slow - def test_advanced_qnode(self, dev_name): + def test_advanced_qnode(self): """Test that an advanced qnode can be differentiated with stoch_pulse_grad.""" import jax import jax.numpy as jnp @@ -1571,7 +1563,7 @@ def test_advanced_qnode(self, dev_name): jax.config.update("jax_enable_x64", True) params = [jnp.array(0.21), jnp.array(-0.171), jnp.array([0.05, 0.03, -0.1])] - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) ham = ( qml.pulse.constant * qml.PauliX(0) + (lambda p, t: jnp.sin(p * t)) * qml.PauliZ(0) @@ -1595,11 +1587,7 @@ def ansatz(params): with qml.Tracker(dev) as tracker: grad_pulse_grad = jax.grad(qnode_pulse_grad)(params) - assert ( - tracker.totals["executions"] == (1 + 2 * 3 * num_split_times) - if dev_name == "default.qubit.jax" - else 1 - ) + assert tracker.totals["executions"] == 1 + 2 * 3 * num_split_times grad_backprop = jax.grad(qnode_backprop)(params) assert all( @@ -1607,7 +1595,7 @@ def ansatz(params): ) jax.clear_caches() - def test_multi_return_broadcasting_multi_shots_raises(self, dev_name): + def test_multi_return_broadcasting_multi_shots_raises(self): """Test that a simple qnode that returns an expectation value and probabilities can be differentiated with stoch_pulse_grad with use_broadcasting.""" import jax @@ -1615,8 +1603,7 @@ def test_multi_return_broadcasting_multi_shots_raises(self, dev_name): jax.config.update("jax_enable_x64", True) shots = [100, 100] - key = "prng_key" if dev_name == "default.qubit.jax" else "seed" - dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)}) + dev = qml.device("default.qubit", wires=1, shots=shots, seed=jax.random.PRNGKey(74)) T = 0.2 ham_single_q_const = qml.pulse.constant * qml.PauliY(0) @@ -1639,15 +1626,14 @@ def circuit(params): # TODO: delete error test above and uncomment the following test case once #2690 is resolved. @pytest.mark.parametrize("shots, tol", [(None, 1e-4), (100, 0.1)]) # , ([100, 100], 0.1)]) @pytest.mark.parametrize("num_split_times", [1, 2]) - def test_qnode_probs_expval_broadcasting(self, dev_name, num_split_times, shots, tol): + def test_qnode_probs_expval_broadcasting(self, num_split_times, shots, tol): """Test that a simple qnode that returns an expectation value and probabilities can be differentiated with stoch_pulse_grad with use_broadcasting.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - key = "prng_key" if dev_name == "default.qubit.jax" else "seed" - dev = qml.device(dev_name, wires=1, shots=shots, **{key: jax.random.PRNGKey(74)}) + dev = qml.device("default.qubit", wires=1, shots=shots, seed=jax.random.PRNGKey(74)) T = 0.2 ham_single_q_const = qml.pulse.constant * qml.PauliY(0) @@ -1676,13 +1662,13 @@ def circuit(params): jax.clear_caches() @pytest.mark.parametrize("num_split_times", [1, 2]) - def test_broadcasting_coincides_with_nonbroadcasting(self, dev_name, num_split_times): + def test_broadcasting_coincides_with_nonbroadcasting(self, num_split_times): """Test that using broadcasting or not does not change the result.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) T = 0.2 def f(p, t): @@ -1720,7 +1706,7 @@ def ansatz(params): assert qml.math.allclose(j0, j1) jax.clear_caches() - def test_with_drive_exact(self, dev_name): + def test_with_drive_exact(self): """Test that a HardwareHamiltonian only containing a drive is differentiated correctly for a constant amplitude and zero frequency and phase.""" import jax @@ -1729,7 +1715,7 @@ def test_with_drive_exact(self, dev_name): H = qml.pulse.transmon_drive(qml.pulse.constant, 0.0, 0.0, wires=[0]) atol = 1e-5 - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) def ansatz(params): qml.evolve(H, atol=atol)(params, t=timespan) @@ -1745,7 +1731,7 @@ def ansatz(params): assert qml.math.allclose(res, exact, atol=6e-5) jax.clear_caches() - def test_with_drive_approx(self, dev_name): + def test_with_drive_approx(self): """Test that a HardwareHamiltonian only containing a drive is differentiated approximately correctly for a constant phase and zero frequency.""" import jax @@ -1754,7 +1740,7 @@ def test_with_drive_approx(self, dev_name): H = qml.pulse.transmon_drive(1 / (2 * np.pi), qml.pulse.constant, 0.0, wires=[0]) atol = 1e-5 - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) def ansatz(params): qml.evolve(H, atol=atol)(params, t=timespan) @@ -1780,7 +1766,7 @@ def ansatz(params): @pytest.mark.slow @pytest.mark.parametrize("num_params", [1, 2]) - def test_with_two_drives(self, dev_name, num_params): + def test_with_two_drives(self, num_params): """Test that a HardwareHamiltonian only containing two drives is differentiated approximately correctly. The two cases of the parametrization test the cases where reordered parameters @@ -1799,7 +1785,7 @@ def test_with_two_drives(self, dev_name, num_params): amps[0], qml.pulse.constant, 0.0, wires=[0] ) + qml.pulse.rydberg_drive(amps[1], qml.pulse.constant, 0.0, wires=[1]) atol = 1e-5 - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) def ansatz(params): qml.evolve(H, atol=atol)(params, t=timespan) @@ -1824,19 +1810,18 @@ def ansatz(params): @pytest.mark.jax -@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) class TestStochPulseGradDiff: """Test that stoch_pulse_grad is differentiable.""" # pylint: disable=too-few-public-methods @pytest.mark.slow - def test_jax(self, dev_name): + def test_jax(self): """Test that stoch_pulse_grad is differentiable with JAX.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = 0.5 ham_single_q_const = qml.pulse.constant * qml.PauliY(0) diff --git a/tests/gradients/core/test_pulse_odegen.py b/tests/gradients/core/test_pulse_odegen.py index 302bd12d4fd..3a7dbd4bddc 100644 --- a/tests/gradients/core/test_pulse_odegen.py +++ b/tests/gradients/core/test_pulse_odegen.py @@ -997,20 +997,18 @@ def test_all_zero_diff_methods_multiple_returns_tape(self): @pytest.mark.jax -@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) class TestPulseOdegenTape: """Test that differentiating tapes with ``pulse_odegen`` works.""" @pytest.mark.parametrize("shots, tol", [(None, 1e-7), (1000, 0.05), ([1000, 100], 0.05)]) - def test_single_pulse_single_term(self, dev_name, shots, tol): + def test_single_pulse_single_term(self, shots, tol): """Test that a single pulse with a single Hamiltonian term is differentiated correctly.""" import jax import jax.numpy as jnp prng_key = jax.random.PRNGKey(8251) - key = "prng_key" if dev_name == "default.qubit.jax" else "seed" - dev = qml.device(dev_name, wires=1, shots=shots, **{key: prng_key}) + dev = qml.device("default.qubit", wires=1, shots=shots, seed=prng_key) H = jnp.polyval * X(0) x = jnp.array([0.4, 0.2, 0.1]) @@ -1038,16 +1036,15 @@ def test_single_pulse_single_term(self, dev_name, shots, tol): @pytest.mark.slow @pytest.mark.parametrize("shots, tol", [(None, 1e-7), ([1000, 100], 0.05)]) - def test_single_pulse_multi_term(self, dev_name, shots, tol): + def test_single_pulse_multi_term(self, shots, tol): """Test that a single pulse with multiple Hamiltonian terms is differentiated correctly.""" import jax import jax.numpy as jnp prng_key = jax.random.PRNGKey(8251) - dev = qml.device(dev_name, wires=1, shots=None) - key = "prng_key" if dev_name == "default.qubit.jax" else "seed" - dev_shots = qml.device(dev_name, wires=1, shots=shots, **{key: prng_key}) + dev = qml.device("default.qubit", wires=1, shots=None) + dev_shots = qml.device("default.qubit", wires=1, shots=shots, seed=prng_key) H = 0.1 * Z(0) + jnp.polyval * X(0) + qml.pulse.constant * Y(0) x = jnp.array([0.4, 0.2, 0.1]) @@ -1079,13 +1076,13 @@ def circuit(par): assert all(qml.math.allclose(g, e, atol=tol) for g, e in zip(grad, exp_grad)) @pytest.mark.parametrize("argnum", (0, [0], 1, [1])) - def test_single_pulse_multi_term_argnum(self, dev_name, argnum): + def test_single_pulse_multi_term_argnum(self, argnum): """Test that a single pulse with multiple Hamiltonian terms is differentiated correctly when setting ``argnum``.""" import jax import jax.numpy as jnp - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) H = jnp.polyval * X(0) + qml.pulse.constant * X(0) x = jnp.array([0.4, 0.2, 0.1]) @@ -1119,16 +1116,15 @@ def test_single_pulse_multi_term_argnum(self, dev_name, argnum): @pytest.mark.slow @pytest.mark.parametrize("shots, tol", [(None, 1e-7), ([1000, 100], 0.05)]) - def test_multi_pulse(self, dev_name, shots, tol): + def test_multi_pulse(self, shots, tol): """Test that a single pulse with multiple Hamiltonian terms is differentiated correctly.""" import jax import jax.numpy as jnp prng_key = jax.random.PRNGKey(8251) - dev = qml.device(dev_name, wires=1, shots=None) - key = "prng_key" if dev_name == "default.qubit.jax" else "seed" - dev_shots = qml.device(dev_name, wires=1, shots=shots, **{key: prng_key}) + dev = qml.device("default.qubit", wires=1, shots=None) + dev_shots = qml.device("default.qubit", wires=1, shots=shots, seed=prng_key) H0 = 0.1 * Z(0) + jnp.polyval * X(0) H1 = 0.2 * Y(0) + qml.pulse.constant * Y(0) + jnp.polyval * Z(0) @@ -1164,15 +1160,14 @@ def circuit(par): @pytest.mark.jax -@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) class TestPulseOdegenQNode: """Test that pulse_odegen integrates correctly with QNodes.""" - def test_raises_for_application_to_qnodes(self, dev_name): + def test_raises_for_application_to_qnodes(self): """Test that an error is raised when applying ``stoch_pulse_grad`` to a QNode directly.""" - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) ham_single_q_const = qml.pulse.constant * qml.PauliY(0) @qml.qnode(dev, interface="jax") @@ -1186,14 +1181,14 @@ def circuit(params): # TODO: include the following tests when #4225 is resolved. @pytest.mark.skip("Applying this gradient transform to QNodes directly is not supported.") - def test_qnode_expval_single_par(self, dev_name): + def test_qnode_expval_single_par(self): """Test that a simple qnode that returns an expectation value can be differentiated with pulse_odegen.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = 0.2 ham_single_q_const = qml.pulse.constant * Y(0) @@ -1212,14 +1207,14 @@ def circuit(params): assert tracker.totals["executions"] == 2 # two shifted tapes @pytest.mark.skip("Applying this gradient transform to QNodes directly is not supported.") - def test_qnode_expval_probs_single_par(self, dev_name): + def test_qnode_expval_probs_single_par(self): """Test that a simple qnode that returns an expectation value can be differentiated with pulse_odegen.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = 0.2 ham_single_q_const = jnp.polyval * Y(0) @@ -1244,14 +1239,14 @@ def circuit(params): assert qml.math.allclose(j, e) @pytest.mark.skip("Applying this gradient transform to QNodes directly is not supported.") - def test_qnode_probs_expval_multi_par(self, dev_name): + def test_qnode_probs_expval_multi_par(self): """Test that a simple qnode that returns probabilities can be differentiated with pulse_odegen.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1, shots=None) + dev = qml.device("default.qubit", wires=1, shots=None) T = 0.2 ham_single_q_const = jnp.polyval * Y(0) + qml.pulse.constant * Y(0) @@ -1284,18 +1279,17 @@ def circuit(params, c): @pytest.mark.jax -@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) class TestPulseOdegenIntegration: """Test that pulse_odegen integrates correctly with QNodes.""" - def test_simple_qnode_expval(self, dev_name): + def test_simple_qnode_expval(self): """Test that a simple qnode that returns an expectation value can be differentiated with pulse_odegen.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = 0.2 ham_single_q_const = qml.pulse.constant * Y(0) @@ -1312,14 +1306,14 @@ def circuit(params): assert qml.math.allclose(grad, exp_grad) assert tracker.totals["executions"] == 1 + 2 # one forward pass, two shifted tapes - def test_simple_qnode_expval_two_evolves(self, dev_name): + def test_simple_qnode_expval_two_evolves(self): """Test that a simple qnode that returns an expectation value can be differentiated with pulse_odegen.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T_x = 0.1 T_y = 0.2 ham_x = qml.pulse.constant * X(0) @@ -1338,14 +1332,14 @@ def circuit(params): exp_grad = [[-2 * jnp.sin(2 * (p_x + p_y)) * T_x], [-2 * jnp.sin(2 * (p_x + p_y)) * T_y]] assert qml.math.allclose(grad, exp_grad) - def test_simple_qnode_probs(self, dev_name): + def test_simple_qnode_probs(self): """Test that a simple qnode that returns probabilities can be differentiated with pulse_odegen.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = 0.2 ham_single_q_const = qml.pulse.constant * Y(0) @@ -1360,14 +1354,14 @@ def circuit(params): exp_jac = jnp.array([-1, 1]) * jnp.sin(2 * p) * T assert qml.math.allclose(jac, exp_jac) - def test_simple_qnode_probs_expval(self, dev_name): + def test_simple_qnode_probs_expval(self): """Test that a simple qnode that returns probabilities can be differentiated with pulse_odegen.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = 0.2 ham_single_q_const = jnp.polyval * Y(0) @@ -1389,13 +1383,13 @@ def circuit(params): @pytest.mark.xfail @pytest.mark.parametrize("time_interface", ["python", "numpy", "jax"]) - def test_simple_qnode_jit(self, dev_name, time_interface): + def test_simple_qnode_jit(self, time_interface): """Test that a simple qnode can be differentiated with pulse_odegen.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = {"python": 0.2, "numpy": np.array(0.2), "jax": jnp.array(0.2)}[time_interface] ham_single_q_const = qml.pulse.constant * Y(0) @@ -1411,7 +1405,7 @@ def circuit(params, T=None): assert qml.math.isclose(jit_grad, exp_grad) @pytest.mark.slow - def test_advanced_qnode(self, dev_name): + def test_advanced_qnode(self): """Test that an advanced qnode can be differentiated with pulse_odegen.""" import jax import jax.numpy as jnp @@ -1419,7 +1413,7 @@ def test_advanced_qnode(self, dev_name): jax.config.update("jax_enable_x64", True) params = [jnp.array(0.21), jnp.array(-0.171), jnp.array([0.05, 0.03, -0.1])] - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) ham = ( qml.pulse.constant * X(0) + (lambda p, t: jnp.sin(p * t)) * Z(0) @@ -1448,14 +1442,14 @@ def ansatz(params): ) @pytest.mark.parametrize("argnums", [[0, 1], 0, 1]) - def test_simple_qnode_expval_multiple_params(self, dev_name, argnums): + def test_simple_qnode_expval_multiple_params(self, argnums): """Test that a simple qnode with two parameters can be differentiated with pulse_odegen and `argnums` works as expected.""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = 0.2 ham1 = qml.pulse.constant * Y(0) ham2 = qml.pulse.constant * Y(0) @@ -1481,7 +1475,6 @@ def circuit(param1, param2): @pytest.mark.jax -@pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) class TestPulseOdegenDiff: """Test that pulse_odegen is differentiable, i.e. that computing the derivative with pulse_odegen is differentiable a second time, @@ -1489,14 +1482,14 @@ class TestPulseOdegenDiff: # pylint: disable=too-few-public-methods @pytest.mark.slow - def test_jax(self, dev_name): + def test_jax(self): """Test that pulse_odegen is differentiable, allowing to compute the Hessian, with JAX..""" import jax import jax.numpy as jnp jax.config.update("jax_enable_x64", True) - dev = qml.device(dev_name, wires=1) + dev = qml.device("default.qubit", wires=1) T = 0.5 ham_single_q_const = qml.pulse.constant * Y(0) diff --git a/tests/gradients/core/test_vjp.py b/tests/gradients/core/test_vjp.py index 15e652f3eb2..6917207f715 100644 --- a/tests/gradients/core/test_vjp.py +++ b/tests/gradients/core/test_vjp.py @@ -371,12 +371,10 @@ class TestVJPGradients: """Gradient tests for the vjp function""" @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, dev_name, tol): + def test_autograd(self, tol): """Tests that the output of the VJP transform can be differentiated using autograd.""" - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = np.array([0.543, -0.654], requires_grad=True) def cost_fn(x, dy): @@ -386,8 +384,7 @@ def cost_fn(x, dy): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.vjp(tape, dy, param_shift) - vjp = fn(execute_fn(tapes)) - return vjp + return fn(dev.execute(tapes)) dy = np.array([-1.0, 0.0, 0.0, 1.0], requires_grad=False) res = cost_fn(params, dy) @@ -397,13 +394,12 @@ def cost_fn(x, dy): assert np.allclose(res, qml.jacobian(expected)(params), atol=tol, rtol=0) @pytest.mark.torch - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_torch(self, dev_name, tol): + def test_torch(self, tol): """Tests that the output of the VJP transform can be differentiated using Torch.""" import torch - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) params_np = np.array([0.543, -0.654], requires_grad=True) params = torch.tensor(params_np, requires_grad=True, dtype=torch.float64) @@ -427,14 +423,12 @@ def test_torch(self, dev_name, tol): @pytest.mark.tf @pytest.mark.slow - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf(self, dev_name, tol): + def test_tf(self, tol): """Tests that the output of the VJP transform can be differentiated using TF.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params_np = np.array([0.543, -0.654], requires_grad=True) params = tf.Variable(params_np, dtype=tf.float64) @@ -447,7 +441,7 @@ def test_tf(self, dev_name, tol): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.vjp(tape, dy, param_shift) - vjp = fn(execute_fn(tapes)) + vjp = fn(dev.execute(tapes)) assert np.allclose(vjp, expected(params), atol=tol, rtol=0) @@ -492,15 +486,13 @@ def test_tf(self, dev_name, tol): @pytest.mark.jax @pytest.mark.slow - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_jax(self, dev_name, tol): + def test_jax(self, tol): """Tests that the output of the VJP transform can be differentiated using JAX.""" import jax from jax import numpy as jnp - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params_np = np.array([0.543, -0.654], requires_grad=True) params = jnp.array(params_np) @@ -512,8 +504,7 @@ def cost_fn(x): dy = jax.numpy.array([-1.0, 0.0, 0.0, 1.0]) tape.trainable_params = {0, 1} tapes, fn = qml.gradients.vjp(tape, dy, param_shift) - vjp = fn(execute_fn(tapes)) - return vjp + return fn(dev.execute(tapes)) res = cost_fn(params) assert np.allclose(res, expected(params), atol=tol, rtol=0) diff --git a/tests/gradients/finite_diff/test_finite_difference.py b/tests/gradients/finite_diff/test_finite_difference.py index dc7a0c4b951..29045701566 100644 --- a/tests/gradients/finite_diff/test_finite_difference.py +++ b/tests/gradients/finite_diff/test_finite_difference.py @@ -857,12 +857,10 @@ class TestFiniteDiffGradients: """Test that the transform is differentiable""" @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, dev_name, approx_order, strategy, tol): + def test_autograd(self, approx_order, strategy, tol): """Tests that the output of the finite-difference transform can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = np.array([0.543, -0.654], requires_grad=True) def cost_fn(x): @@ -875,8 +873,7 @@ def cost_fn(x): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy) - jac = np.array(fn(execute_fn(tapes))) - return jac + return np.array(fn(dev.execute(tapes))) res = qml.jacobian(cost_fn)(params) x, y = params @@ -890,12 +887,10 @@ def cost_fn(x): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd_ragged(self, dev_name, approx_order, strategy, tol): + def test_autograd_ragged(self, approx_order, strategy, tol): """Tests that the output of the finite-difference transform of a ragged tape can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = np.array([0.543, -0.654], requires_grad=True) def cost_fn(x): @@ -909,7 +904,7 @@ def cost_fn(x): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy) - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) return jac[1][0] x, y = params @@ -919,14 +914,12 @@ def cost_fn(x): @pytest.mark.tf @pytest.mark.slow - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf(self, dev_name, approx_order, strategy, tol): + def test_tf(self, approx_order, strategy, tol): """Tests that the output of the finite-difference transform can be differentiated using TF, yielding second derivatives.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = tf.Variable([0.543, -0.654], dtype=tf.float64) with tf.GradientTape(persistent=True) as t: @@ -939,7 +932,7 @@ def test_tf(self, dev_name, approx_order, strategy, tol): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy) - jac_0, jac_1 = fn(execute_fn(tapes)) + jac_0, jac_1 = fn(dev.execute(tapes)) x, y = 1.0 * params @@ -955,14 +948,12 @@ def test_tf(self, dev_name, approx_order, strategy, tol): assert np.allclose([res_0, res_1], expected, atol=tol, rtol=0) @pytest.mark.tf - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf_ragged(self, dev_name, approx_order, strategy, tol): + def test_tf_ragged(self, approx_order, strategy, tol): """Tests that the output of the finite-difference transform of a ragged tape can be differentiated using TF, yielding second derivatives.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = tf.Variable([0.543, -0.654], dtype=tf.float64) with tf.GradientTape(persistent=True) as t: @@ -977,7 +968,7 @@ def test_tf_ragged(self, dev_name, approx_order, strategy, tol): tape.trainable_params = {0, 1} tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy) - jac_01 = fn(execute_fn(tapes))[1][0] + jac_01 = fn(dev.execute(tapes))[1][0] x, y = 1.0 * params @@ -988,14 +979,12 @@ def test_tf_ragged(self, dev_name, approx_order, strategy, tol): assert np.allclose(res_01[0], expected, atol=tol, rtol=0) @pytest.mark.torch - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_torch(self, dev_name, approx_order, strategy, tol): + def test_torch(self, approx_order, strategy, tol): """Tests that the output of the finite-difference transform can be differentiated using Torch, yielding second derivatives.""" import torch - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True) def cost_fn(params): @@ -1007,7 +996,7 @@ def cost_fn(params): tape = qml.tape.QuantumScript.from_queue(q) tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy) - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) return jac hess = torch.autograd.functional.jacobian(cost_fn, params) @@ -1025,15 +1014,13 @@ def cost_fn(params): assert np.allclose(hess[1].detach().numpy(), expected[1], atol=tol, rtol=0) @pytest.mark.jax - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_jax(self, dev_name, approx_order, strategy, tol): + def test_jax(self, approx_order, strategy, tol): """Tests that the output of the finite-difference transform can be differentiated using JAX, yielding second derivatives.""" import jax from jax import numpy as jnp - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = jnp.array([0.543, -0.654]) def cost_fn(x): @@ -1046,8 +1033,7 @@ def cost_fn(x): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1} tapes, fn = finite_diff(tape, n=1, approx_order=approx_order, strategy=strategy) - jac = fn(execute_fn(tapes)) - return jac + return fn(dev.execute(tapes)) res = jax.jacobian(cost_fn)(params) assert isinstance(res, tuple) @@ -1061,14 +1047,11 @@ def cost_fn(x): assert np.allclose(res, expected, atol=tol, rtol=0) @pytest.mark.jax - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_jax_probs( - self, dev_name, approx_order, strategy, tol - ): # pylint: disable=unused-argument + def test_jax_probs(self, approx_order, strategy, tol): # pylint: disable=unused-argument """Tests that the output of the finite-difference transform is similar using or not diff method on the QNode.""" import jax - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) x = jax.numpy.array(0.543) @qml.qnode(dev) diff --git a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py index 09fb355084f..b86e829703b 100644 --- a/tests/gradients/finite_diff/test_finite_difference_shot_vec.py +++ b/tests/gradients/finite_diff/test_finite_difference_shot_vec.py @@ -844,12 +844,10 @@ class TestFiniteDiffGradients: """Test that the transform is differentiable""" @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, dev_name, approx_order, strategy): + def test_autograd(self, approx_order, strategy): """Tests that the output of the finite-difference transform can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = np.array([0.543, -0.654], requires_grad=True) def cost_fn(x): @@ -868,8 +866,7 @@ def cost_fn(x): strategy=strategy, h=h_val, ) - jac = np.array(fn(execute_fn(tapes))) - return jac + return np.array(fn(dev.execute(tapes))) all_res = qml.jacobian(cost_fn)(params) @@ -888,12 +885,10 @@ def cost_fn(x): assert np.allclose(res, expected, atol=0.3, rtol=0) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd_ragged(self, dev_name, approx_order, strategy): + def test_autograd_ragged(self, approx_order, strategy): """Tests that the output of the finite-difference transform of a ragged tape can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = np.array([0.543, -0.654], requires_grad=True) def cost_fn(x): @@ -913,7 +908,7 @@ def cost_fn(x): strategy=strategy, h=h_val, ) - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) return jac[1][0] x, y = params @@ -928,14 +923,12 @@ def cost_fn(x): @pytest.mark.tf @pytest.mark.slow - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf(self, dev_name, approx_order, strategy): + def test_tf(self, approx_order, strategy): """Tests that the output of the finite-difference transform can be differentiated using TF, yielding second derivatives.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = tf.Variable([0.543, -0.654], dtype=tf.float64) with tf.GradientTape(persistent=True) as t: @@ -954,7 +947,7 @@ def test_tf(self, dev_name, approx_order, strategy): strategy=strategy, h=h_val, ) - jac_0, jac_1 = fn(execute_fn(tapes)) + jac_0, jac_1 = fn(dev.execute(tapes)) x, y = 1.0 * params @@ -970,14 +963,12 @@ def test_tf(self, dev_name, approx_order, strategy): assert np.allclose([res_0, res_1], expected, atol=0.3, rtol=0) @pytest.mark.tf - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf_ragged(self, dev_name, approx_order, strategy): + def test_tf_ragged(self, approx_order, strategy): """Tests that the output of the finite-difference transform of a ragged tape can be differentiated using TF, yielding second derivatives.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = tf.Variable([0.543, -0.654], dtype=tf.float64) with tf.GradientTape(persistent=True) as t: @@ -998,7 +989,7 @@ def test_tf_ragged(self, dev_name, approx_order, strategy): h=h_val, ) - jac_01 = fn(execute_fn(tapes))[1][0] + jac_01 = fn(dev.execute(tapes))[1][0] x, y = 1.0 * params @@ -1009,14 +1000,12 @@ def test_tf_ragged(self, dev_name, approx_order, strategy): assert np.allclose(res_01[0], expected, atol=0.3, rtol=0) @pytest.mark.torch - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_torch(self, dev_name, approx_order, strategy): + def test_torch(self, approx_order, strategy): """Tests that the output of the finite-difference transform can be differentiated using Torch, yielding second derivatives.""" import torch - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True) def cost_fn(params): @@ -1034,8 +1023,7 @@ def cost_fn(params): strategy=strategy, h=h_val, ) - jac = fn(execute_fn(tapes)) - return jac + return fn(dev.execute(tapes)) hess = torch.autograd.functional.jacobian(cost_fn, params) @@ -1052,15 +1040,13 @@ def cost_fn(params): assert np.allclose(hess[1].detach().numpy(), expected[1], atol=0.3, rtol=0) @pytest.mark.jax - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_jax(self, dev_name, approx_order, strategy): + def test_jax(self, approx_order, strategy): """Tests that the output of the finite-difference transform can be differentiated using JAX, yielding second derivatives.""" import jax from jax import numpy as jnp - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = jnp.array([0.543, -0.654]) def cost_fn(x): @@ -1079,8 +1065,7 @@ def cost_fn(x): strategy=strategy, h=h_val, ) - jac = fn(execute_fn(tapes)) - return jac + return fn(dev.execute(tapes)) all_res = jax.jacobian(cost_fn)(params) diff --git a/tests/gradients/finite_diff/test_spsa_gradient.py b/tests/gradients/finite_diff/test_spsa_gradient.py index fdc6e9051dd..d8f19dcf826 100644 --- a/tests/gradients/finite_diff/test_spsa_gradient.py +++ b/tests/gradients/finite_diff/test_spsa_gradient.py @@ -985,12 +985,10 @@ class TestSpsaGradientDifferentiation: """Test that the transform is differentiable""" @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, dev_name, sampler, num_directions, atol): + def test_autograd(self, sampler, num_directions, atol): """Tests that the output of the SPSA gradient transform can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = np.array([0.543, -0.654], requires_grad=True) rng = np.random.default_rng(42) @@ -1006,7 +1004,7 @@ def cost_fn(x): tapes, fn = spsa_grad( tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng ) - jac = np.array(fn(execute_fn(tapes))) + jac = np.array(fn(dev.execute(tapes))) if sampler is coordinate_sampler: jac *= 2 return jac @@ -1023,12 +1021,10 @@ def cost_fn(x): assert np.allclose(res, expected, atol=atol, rtol=0) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd_ragged(self, dev_name, sampler, num_directions, atol): + def test_autograd_ragged(self, sampler, num_directions, atol): """Tests that the output of the SPSA gradient transform of a ragged tape can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = np.array([0.543, -0.654], requires_grad=True) rng = np.random.default_rng(42) @@ -1045,7 +1041,7 @@ def cost_fn(x): tapes, fn = spsa_grad( tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng ) - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) if sampler is coordinate_sampler: jac = tuple(tuple(2 * _j for _j in _jac) for _jac in jac) return jac[1][0] @@ -1057,14 +1053,12 @@ def cost_fn(x): @pytest.mark.tf @pytest.mark.slow - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf(self, dev_name, sampler, num_directions, atol): + def test_tf(self, sampler, num_directions, atol): """Tests that the output of the SPSA gradient transform can be differentiated using TF, yielding second derivatives.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = tf.Variable([0.543, -0.654], dtype=tf.float64) rng = np.random.default_rng(42) @@ -1080,7 +1074,7 @@ def test_tf(self, dev_name, sampler, num_directions, atol): tapes, fn = spsa_grad( tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng ) - jac_0, jac_1 = fn(execute_fn(tapes)) + jac_0, jac_1 = fn(dev.execute(tapes)) if sampler is coordinate_sampler: jac_0 *= 2 jac_1 *= 2 @@ -1100,14 +1094,12 @@ def test_tf(self, dev_name, sampler, num_directions, atol): @pytest.mark.tf @pytest.mark.slow - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf_ragged(self, dev_name, sampler, num_directions, atol): + def test_tf_ragged(self, sampler, num_directions, atol): """Tests that the output of the SPSA gradient transform of a ragged tape can be differentiated using TF, yielding second derivatives.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = tf.Variable([0.543, -0.654], dtype=tf.float64) rng = np.random.default_rng(42) @@ -1125,7 +1117,7 @@ def test_tf_ragged(self, dev_name, sampler, num_directions, atol): tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng ) - jac_01 = fn(execute_fn(tapes))[1][0] + jac_01 = fn(dev.execute(tapes))[1][0] if sampler is coordinate_sampler: jac_01 *= 2 @@ -1138,14 +1130,12 @@ def test_tf_ragged(self, dev_name, sampler, num_directions, atol): assert np.allclose(res_01[0], expected, atol=atol, rtol=0) @pytest.mark.torch - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_torch(self, dev_name, sampler, num_directions, atol): + def test_torch(self, sampler, num_directions, atol): """Tests that the output of the SPSA gradient transform can be differentiated using Torch, yielding second derivatives.""" import torch - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True) rng = np.random.default_rng(42) @@ -1160,7 +1150,7 @@ def cost_fn(params): tapes, fn = spsa_grad( tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng ) - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) if sampler is coordinate_sampler: jac = tuple(2 * _jac for _jac in jac) return jac @@ -1180,15 +1170,13 @@ def cost_fn(params): assert np.allclose(hess[1].detach().numpy(), expected[1], atol=atol, rtol=0) @pytest.mark.jax - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_jax(self, dev_name, sampler, num_directions, atol): + def test_jax(self, sampler, num_directions, atol): """Tests that the output of the SPSA gradient transform can be differentiated using JAX, yielding second derivatives.""" import jax from jax import numpy as jnp - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = jnp.array([0.543, -0.654]) rng = np.random.default_rng(42) @@ -1204,7 +1192,7 @@ def cost_fn(x): tapes, fn = spsa_grad( tape, n=1, num_directions=num_directions, sampler=sampler, sampler_rng=rng ) - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) if sampler is coordinate_sampler: jac = tuple(2 * _jac for _jac in jac) return jac diff --git a/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py b/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py index 83e145f8204..46f8aa1288e 100644 --- a/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py +++ b/tests/gradients/finite_diff/test_spsa_gradient_shot_vec.py @@ -939,12 +939,10 @@ class TestSpsaGradientDifferentiation: """Test that the transform is differentiable""" @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, dev_name, approx_order, strategy): + def test_autograd(self, approx_order, strategy): """Tests that the output of the SPSA gradient transform can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = np.array([0.543, -0.654], requires_grad=True) rng = np.random.default_rng(42) @@ -965,8 +963,7 @@ def cost_fn(x): h=h_val, sampler_rng=rng, ) - jac = np.array(fn(execute_fn(tapes))) - return jac + return np.array(fn(dev.execute(tapes))) all_res = qml.jacobian(cost_fn)(params) @@ -985,12 +982,10 @@ def cost_fn(x): assert np.allclose(res, expected, atol=spsa_shot_vec_tol, rtol=0) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd_ragged(self, dev_name, approx_order, strategy): + def test_autograd_ragged(self, approx_order, strategy): """Tests that the output of the SPSA gradient transform of a ragged tape can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = np.array([0.543, -0.654], requires_grad=True) rng = np.random.default_rng(42) @@ -1012,7 +1007,7 @@ def cost_fn(x): h=h_val, sampler_rng=rng, ) - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) return jac[1][0] x, y = params @@ -1027,14 +1022,12 @@ def cost_fn(x): @pytest.mark.tf @pytest.mark.slow - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf(self, dev_name, approx_order, strategy): + def test_tf(self, approx_order, strategy): """Tests that the output of the SPSA gradient transform can be differentiated using TF, yielding second derivatives.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = tf.Variable([0.543, -0.654], dtype=tf.float64) rng = np.random.default_rng(42) @@ -1055,7 +1048,7 @@ def test_tf(self, dev_name, approx_order, strategy): h=h_val, sampler_rng=rng, ) - jac_0, jac_1 = fn(execute_fn(tapes)) + jac_0, jac_1 = fn(dev.execute(tapes)) x, y = 1.0 * params @@ -1071,14 +1064,12 @@ def test_tf(self, dev_name, approx_order, strategy): assert np.allclose([res_0, res_1], expected, atol=spsa_shot_vec_tol, rtol=0) @pytest.mark.tf - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf_ragged(self, dev_name, approx_order, strategy): + def test_tf_ragged(self, approx_order, strategy): """Tests that the output of the SPSA gradient transform of a ragged tape can be differentiated using TF, yielding second derivatives.""" import tensorflow as tf - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = tf.Variable([0.543, -0.654], dtype=tf.float64) rng = np.random.default_rng(42) @@ -1101,7 +1092,7 @@ def test_tf_ragged(self, dev_name, approx_order, strategy): sampler_rng=rng, ) - jac_01 = fn(execute_fn(tapes))[1][0] + jac_01 = fn(dev.execute(tapes))[1][0] x, y = 1.0 * params @@ -1112,14 +1103,12 @@ def test_tf_ragged(self, dev_name, approx_order, strategy): assert np.allclose(res_01[0], expected, atol=spsa_shot_vec_tol, rtol=0) @pytest.mark.torch - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_torch(self, dev_name, approx_order, strategy): + def test_torch(self, approx_order, strategy): """Tests that the output of the SPSA gradient transform can be differentiated using Torch, yielding second derivatives.""" import torch - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = torch.tensor([0.543, -0.654], dtype=torch.float64, requires_grad=True) rng = np.random.default_rng(42) @@ -1139,8 +1128,7 @@ def cost_fn(params): h=h_val, sampler_rng=rng, ) - jac = fn(execute_fn(tapes)) - return jac + return fn(dev.execute(tapes)) hess = torch.autograd.functional.jacobian(cost_fn, params) @@ -1157,15 +1145,13 @@ def cost_fn(params): assert np.allclose(hess[1].detach().numpy(), expected[1], atol=spsa_shot_vec_tol, rtol=0) @pytest.mark.jax - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_jax(self, dev_name, approx_order, strategy): + def test_jax(self, approx_order, strategy): """Tests that the output of the SPSA gradient transform can be differentiated using JAX, yielding second derivatives.""" import jax from jax import numpy as jnp - dev = qml.device(dev_name, wires=2, shots=many_shots_shot_vector) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2, shots=many_shots_shot_vector) params = jnp.array([0.543, -0.654]) rng = np.random.default_rng(42) @@ -1186,8 +1172,7 @@ def cost_fn(x): h=h_val, sampler_rng=rng, ) - jac = fn(execute_fn(tapes)) - return jac + return fn(dev.execute(tapes)) all_res = jax.jacobian(cost_fn)(params) diff --git a/tests/gradients/parameter_shift/test_parameter_shift.py b/tests/gradients/parameter_shift/test_parameter_shift.py index c857ff815c3..e5a00b2cdf0 100644 --- a/tests/gradients/parameter_shift/test_parameter_shift.py +++ b/tests/gradients/parameter_shift/test_parameter_shift.py @@ -1603,12 +1603,10 @@ def cost_fn(params): # assert np.allclose(jac[1, 1, 1], -2 * np.cos(2 * y), atol=tol, rtol=0) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_fallback_single_meas(self, dev_name, mocker): + def test_fallback_single_meas(self, mocker): """Test that fallback gradient functions are correctly used for a single measurement.""" spy = mocker.spy(qml.gradients, "finite_diff") - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) x = 0.543 y = -0.654 @@ -1628,7 +1626,7 @@ def cost_fn(params): spy.assert_called() assert spy.call_args[1]["argnum"] == {1} - return fn(execute_fn(tapes)) + return fn(dev.execute(tapes)) res = cost_fn(params) @@ -1645,12 +1643,10 @@ def cost_fn(params): @pytest.mark.autograd @pytest.mark.parametrize("RX, RY, argnum", [(RX_with_F, qml.RY, 0), (qml.RX, RY_with_F, 1)]) - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_fallback_probs(self, dev_name, RX, RY, argnum, mocker): + def test_fallback_probs(self, RX, RY, argnum, mocker): """Test that fallback gradient functions are correctly used with probs""" spy = mocker.spy(qml.gradients, "finite_diff") - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) x = 0.543 y = -0.654 @@ -1672,7 +1668,7 @@ def cost_fn(params): spy.assert_called() assert spy.call_args[1]["argnum"] == {argnum} - return fn(execute_fn(tapes)) + return fn(dev.execute(tapes)) res = cost_fn(params) @@ -1729,15 +1725,13 @@ def cost_fn(params): assert np.allclose(res[1][1], probs_expected[:, 1]) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_all_fallback(self, dev_name, mocker, tol): + def test_all_fallback(self, mocker, tol): """Test that *only* the fallback logic is called if no parameters support the parameter-shift rule""" spy_fd = mocker.spy(qml.gradients, "finite_diff") spy_ps = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift") - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) x = 0.543 y = -0.654 @@ -1755,7 +1749,7 @@ def test_all_fallback(self, dev_name, mocker, tol): spy_fd.assert_called() spy_ps.assert_not_called() - res = fn(execute_fn(tapes)) + res = fn(dev.execute(tapes)) assert isinstance(res, tuple) assert res[0].shape == () @@ -2889,16 +2883,14 @@ def test_variance_gradients_agree_finite_differences(self, tol): assert np.allclose(g, grad_F2[idx1][idx2], atol=tol, rtol=0) @pytest.mark.jax - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_fallback(self, dev_name, mocker, tol): + def test_fallback(self, mocker, tol): """Test that fallback gradient functions are correctly used""" import jax from jax import numpy as jnp spy = mocker.spy(qml.gradients, "finite_diff") - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) x = 0.543 y = -0.654 @@ -2920,7 +2912,7 @@ def cost_fn(params): spy.assert_called() assert spy.call_args[1]["argnum"] == {1} - return fn(execute_fn(tapes)) + return fn(dev.execute(tapes)) res = cost_fn(params) assert len(res) == 2 and isinstance(res, tuple) @@ -2934,15 +2926,13 @@ def cost_fn(params): assert np.allclose(jac[1][1][1], -2 * np.cos(2 * y), atol=tol, rtol=0) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_all_fallback(self, dev_name, mocker, tol): + def test_all_fallback(self, mocker, tol): """Test that *only* the fallback logic is called if no parameters support the parameter-shift rule""" spy_fd = mocker.spy(qml.gradients, "finite_diff") spy_ps = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift") - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) x = 0.543 y = -0.654 @@ -2960,7 +2950,7 @@ def test_all_fallback(self, dev_name, mocker, tol): spy_fd.assert_called() spy_ps.assert_not_called() - res = fn(execute_fn(tapes)) + res = fn(dev.execute(tapes)) assert len(res) == 2 assert res[0].shape == () assert res[1].shape == () @@ -3298,12 +3288,10 @@ class TestParamShiftGradients: @pytest.mark.autograd # TODO: support Hessian with the new return types @pytest.mark.skip - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, dev_name, tol, broadcast, expected): + def test_autograd(self, tol, broadcast, expected): """Tests that the output of the parameter-shift transform can be differentiated using autograd, yielding second derivatives.""" - dev = qml.device(dev_name, wires=2) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=2) params = np.array([0.543, -0.654], requires_grad=True) exp_num_tapes, exp_batch_sizes = expected @@ -3319,7 +3307,7 @@ def cost_fn(x): tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast) assert len(tapes) == exp_num_tapes assert [t.batch_size for t in tapes] == exp_batch_sizes - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) return jac res = qml.jacobian(cost_fn)(params) @@ -3575,8 +3563,7 @@ def cost_fn(weights, coeffs1, coeffs2, dev=None, broadcast=False): tape = qml.tape.QuantumScript.from_queue(q) tape.trainable_params = {0, 1, 2, 3, 4, 5} tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast) - execute_fn = dev.batch_execute if isinstance(dev, qml.Device) else dev.execute - jac = fn(execute_fn(tapes)) + jac = fn(dev.execute(tapes)) return jac @staticmethod @@ -3598,14 +3585,13 @@ def cost_fn_expected(weights, coeffs1, coeffs2): ] @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, dev_name, tol, broadcast): + def test_autograd(self, tol, broadcast): """Test gradient of multiple trainable Hamiltonian coefficients using autograd""" 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(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) res = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) @@ -3619,8 +3605,7 @@ def test_autograd(self, dev_name, tol, broadcast): # assert np.allclose(res[2][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0) @pytest.mark.tf - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf(self, dev_name, tol, broadcast): + def test_tf(self, tol, broadcast): """Test gradient of multiple trainable Hamiltonian coefficients using tf""" import tensorflow as tf @@ -3629,15 +3614,8 @@ def test_tf(self, dev_name, tol, broadcast): coeffs2 = tf.Variable([0.7], dtype=tf.float64) weights = tf.Variable([0.4, 0.5], dtype=tf.float64) - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) - # Old op math with old device API does not support broadcasting - if broadcast and "tf" in dev_name and not qml.operation.active_new_opmath(): - with pytest.raises( - NotImplementedError, match="Hamiltonians .* together with parameter broadcasting" - ): - _ = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast) - return with tf.GradientTape() as _: jac = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast) @@ -3654,8 +3632,7 @@ def test_tf(self, dev_name, tol, broadcast): # assert np.allclose(hess[1][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0) @pytest.mark.torch - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_torch(self, dev_name, tol, broadcast): + def test_torch(self, tol, broadcast): """Test gradient of multiple trainable Hamiltonian coefficients using torch""" import torch @@ -3664,15 +3641,8 @@ def test_torch(self, dev_name, tol, broadcast): coeffs2 = torch.tensor([0.7], dtype=torch.float64, requires_grad=True) weights = torch.tensor([0.4, 0.5], dtype=torch.float64, requires_grad=True) - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) - # Old op math with old device API does not support broadcasting - if broadcast and "torch" in dev_name and not qml.operation.active_new_opmath(): - with pytest.raises( - NotImplementedError, match="Hamiltonians .* together with parameter broadcasting" - ): - res = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast) - return res = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast) expected = self.cost_fn_expected( weights.detach().numpy(), coeffs1.detach().numpy(), coeffs2.detach().numpy() @@ -3688,8 +3658,7 @@ def test_torch(self, dev_name, tol, broadcast): # assert np.allclose(hess[2][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0) @pytest.mark.jax - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_jax(self, dev_name, tol, broadcast): + def test_jax(self, tol, broadcast): """Test gradient of multiple trainable Hamiltonian coefficients using JAX""" import jax @@ -3699,15 +3668,8 @@ def test_jax(self, dev_name, tol, broadcast): coeffs1 = jnp.array([0.1, 0.2, 0.3]) coeffs2 = jnp.array([0.7]) weights = jnp.array([0.4, 0.5]) - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) - # Old op math with old device API does not support broadcasting - if broadcast and "jax" in dev_name and not qml.operation.active_new_opmath(): - with pytest.raises( - NotImplementedError, match="Hamiltonians .* together with parameter broadcasting" - ): - res = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast) - return res = self.cost_fn(weights, coeffs1, coeffs2, dev, broadcast) expected = self.cost_fn_expected(weights, coeffs1, coeffs2) assert np.allclose(res, np.array(expected), atol=tol, rtol=0) diff --git a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py index 96ea7e5ef1e..6d9212dceb9 100644 --- a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py +++ b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py @@ -766,14 +766,10 @@ def test_variance_gradients_agree_finite_differences(self, broadcast): assert np.allclose(g, grad_F2[idx1][idx2], atol=finite_diff_tol, rtol=0) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_fallback(self, dev_name, mocker, broadcast): + def test_fallback(self, mocker, broadcast): """Test that fallback gradient functions are correctly used""" - if broadcast and dev_name == "default.qubit.autograd": - pytest.xfail(reason="Return types + autograd + broadcasting does not work") spy = mocker.spy(qml.gradients, "finite_diff") - dev = qml.device(dev_name, wires=3, shots=fallback_shot_vec) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=3, shots=fallback_shot_vec) x = 0.543 y = -0.654 @@ -798,7 +794,7 @@ def cost_fn(params): spy.assert_called() assert spy.call_args[1]["argnum"] == {1} - return fn(execute_fn(tapes)) + return fn(dev.execute(tapes)) all_res = cost_fn(params) @@ -827,15 +823,11 @@ def cost_fn(params): # assert np.allclose(jac[1, 1, 1], -2 * np.cos(2 * y), atol=shot_vec_tol, rtol=0) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_fallback_single_meas(self, dev_name, mocker, broadcast): + def test_fallback_single_meas(self, mocker, broadcast): """Test that fallback gradient functions are correctly used for a single measurement.""" - if broadcast and dev_name == "default.qubit.autograd": - pytest.xfail(reason="Return types + autograd + broadcasting does not work") spy = mocker.spy(qml.gradients, "finite_diff") shot_vec = tuple([1000000] * 4) - dev = qml.device(dev_name, wires=3, shots=shot_vec) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=3, shots=shot_vec) x = 0.543 y = -0.654 @@ -856,7 +848,7 @@ def cost_fn(params): spy.assert_called() assert spy.call_args[1]["argnum"] == {1} - return fn(execute_fn(tapes)) + return fn(dev.execute(tapes)) all_res = cost_fn(params) @@ -872,20 +864,13 @@ def cost_fn(params): assert np.allclose(res[0], expval_expected[0], atol=finite_diff_tol) assert np.allclose(res[1], expval_expected[1], atol=finite_diff_tol) - @pytest.mark.autograd @pytest.mark.parametrize("RX, RY, argnum", [(RX_with_F, qml.RY, 0), (qml.RX, RY_with_F, 1)]) - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) def test_fallback_probs( - self, dev_name, RX, RY, argnum, mocker, broadcast + self, RX, RY, argnum, mocker, broadcast ): # pylint:disable=too-many-arguments """Test that fallback gradient functions are correctly used with probs""" - if broadcast and dev_name == "default.qubit.autograd": - pytest.xfail( - reason="Return types + autograd + old device API + broadcasting does not work" - ) spy = mocker.spy(qml.gradients, "finite_diff") - dev = qml.device(dev_name, wires=3, shots=fallback_shot_vec) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=3, shots=fallback_shot_vec) x = 0.543 y = -0.654 @@ -908,7 +893,7 @@ def cost_fn(params): spy.assert_called() assert spy.call_args[1]["argnum"] == {argnum} - return fn(execute_fn(tapes)) + return fn(dev.execute(tapes)) all_res = cost_fn(params) assert isinstance(all_res, tuple) @@ -969,8 +954,7 @@ def cost_fn(params): assert np.allclose(res[1][1], probs_expected[:, 1], atol=finite_diff_tol) @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_all_fallback(self, dev_name, mocker, broadcast): + def test_all_fallback(self, mocker, broadcast): """Test that *only* the fallback logic is called if no parameters support the parameter-shift rule""" if broadcast: @@ -978,8 +962,7 @@ def test_all_fallback(self, dev_name, mocker, broadcast): spy_fd = mocker.spy(qml.gradients, "finite_diff") spy_ps = mocker.spy(qml.gradients.parameter_shift, "expval_param_shift") - dev = qml.device(dev_name, wires=3, shots=fallback_shot_vec) - execute_fn = dev.execute if dev_name == "default.qubit" else dev.batch_execute + dev = qml.device("default.qubit", wires=3, shots=fallback_shot_vec) x = 0.543 y = -0.654 @@ -998,7 +981,7 @@ def test_all_fallback(self, dev_name, mocker, broadcast): spy_fd.assert_called() spy_ps.assert_not_called() - all_res = fn(execute_fn(tapes)) + all_res = fn(dev.execute(tapes)) assert len(all_res) == len(fallback_shot_vec) assert isinstance(all_res, tuple) @@ -2245,8 +2228,7 @@ def cost_fn(weights, coeffs1, coeffs2, dev=None, broadcast=False): tape = qml.tape.QuantumScript.from_queue(q, shots=dev.shots) tape.trainable_params = {0, 1, 2, 3, 4, 5} tapes, fn = qml.gradients.param_shift(tape, broadcast=broadcast) - execute_fn = dev.batch_execute if isinstance(dev, qml.Device) else dev.execute - return fn(execute_fn(tapes)) + return fn(dev.execute(tapes)) @staticmethod def cost_fn_expected(weights, coeffs1, coeffs2): @@ -2268,15 +2250,14 @@ def cost_fn_expected(weights, coeffs1, coeffs2): @pytest.mark.xfail(reason="TODO") @pytest.mark.autograd - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.autograd"]) - def test_autograd(self, dev_name, broadcast, tol): + def test_autograd(self, broadcast, tol): """Test gradient of multiple trainable Hamiltonian coefficients using autograd""" 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) shot_vec = many_shots_shot_vector - dev = qml.device(dev_name, wires=2, shots=shot_vec) + dev = qml.device("default.qubit", wires=2, shots=shot_vec) if broadcast: with pytest.raises( @@ -2297,8 +2278,7 @@ def test_autograd(self, dev_name, broadcast, tol): @pytest.mark.xfail(reason="TODO") @pytest.mark.tf - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_tf(self, dev_name, broadcast, tol): + def test_tf(self, broadcast, tol): """Test gradient of multiple trainable Hamiltonian coefficients using tf""" import tensorflow as tf @@ -2308,7 +2288,7 @@ def test_tf(self, dev_name, broadcast, tol): weights = tf.Variable([0.4, 0.5], dtype=tf.float64) shot_vec = many_shots_shot_vector - dev = qml.device(dev_name, wires=2, shots=shot_vec) + dev = qml.device("default.qubit", wires=2, shots=shot_vec) if broadcast: with pytest.raises( @@ -2335,8 +2315,7 @@ def test_tf(self, dev_name, broadcast, tol): # TODO: Torch support for param-shift @pytest.mark.torch @pytest.mark.xfail - @pytest.mark.parametrize("dev_name", ["default.qubit"]) - def test_torch(self, dev_name, broadcast, tol): + def test_torch(self, broadcast, tol): """Test gradient of multiple trainable Hamiltonian coefficients using torch""" import torch @@ -2345,7 +2324,7 @@ def test_torch(self, dev_name, broadcast, tol): coeffs2 = torch.tensor([0.7], dtype=torch.float64, requires_grad=True) weights = torch.tensor([0.4, 0.5], dtype=torch.float64, requires_grad=True) - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) if broadcast: with pytest.raises( @@ -2367,8 +2346,7 @@ def test_torch(self, dev_name, broadcast, tol): assert np.allclose(hess[2][:, -1], np.zeros([2, 1, 1]), atol=tol, rtol=0) @pytest.mark.jax - @pytest.mark.parametrize("dev_name", ["default.qubit", "default.qubit.jax"]) - def test_jax(self, dev_name, broadcast, tol): + def test_jax(self, broadcast, tol): """Test gradient of multiple trainable Hamiltonian coefficients using JAX""" import jax @@ -2378,7 +2356,7 @@ def test_jax(self, dev_name, broadcast, tol): coeffs1 = jnp.array([0.1, 0.2, 0.3]) coeffs2 = jnp.array([0.7]) weights = jnp.array([0.4, 0.5]) - dev = qml.device(dev_name, wires=2) + dev = qml.device("default.qubit", wires=2) if broadcast: with pytest.raises( diff --git a/tests/interfaces/legacy_devices_integration/test_autograd_legacy.py b/tests/interfaces/legacy_devices_integration/test_autograd_legacy.py deleted file mode 100644 index 03072f2515d..00000000000 --- a/tests/interfaces/legacy_devices_integration/test_autograd_legacy.py +++ /dev/null @@ -1,1367 +0,0 @@ -# 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_compute_derivatives") - - 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_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)]) - res = qml.execute([qscript], dev, interface=None) - 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.target_device._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.target_device._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 == 1 - - # With caching, only 1 evaluation required. - dev.target_device._num_executions = 0 - jac_fn = qml.jacobian(cost) - jac_fn(params, cache=True) - assert dev.num_executions == 1 - - 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"] == [1, 2] - - 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 - or execute_kwargs["gradient_kwargs"]["method"] == "adjoint_jacobian" - ): - pytest.skip("Adjoint differentiation does not support samples") - - shots = 10 - - 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, shots=10) - return qml.execute([tape], device, **execute_kwargs)[0] - - 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) - - -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/legacy_devices_integration/test_autograd_qnode_legacy.py b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py deleted file mode 100644 index eb65a655877..00000000000 --- a/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py +++ /dev/null @@ -1,2356 +0,0 @@ -# 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. -# 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 a QNode""" -# 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 - -import pennylane as qml -from pennylane import numpy as np -from pennylane import qnode - -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], -] - -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], -] - -pytestmark = pytest.mark.autograd - -TOL_FOR_SPSA = 1.0 -SEED_FOR_SPSA = 32651 -H_FOR_SPSA = 0.01 - - -@pytest.mark.parametrize( - "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_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.target_device, "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) - 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) - - 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) - - 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, interface=interface, diff_method=diff_method) - 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) - 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): - """Test jacobian calculation""" - num_wires = 2 - 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} - 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) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - - res = circuit(a, b) - - def cost(x, y): - return autograd.numpy.hstack(circuit(x, y)) - - assert circuit.qtape.trainable_params == [0, 1] - assert isinstance(res, tuple) - assert len(res) == 2 - - 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) - assert isinstance(res, tuple) and len(res) == 2 - expected = ([-np.sin(a), np.sin(a) * np.sin(b)], [0, -np.cos(a) * np.cos(b)]) - assert isinstance(res[0], np.ndarray) - assert res[0].shape == (2,) - assert np.allclose(res[0], expected[0], atol=tol, rtol=0) - - assert isinstance(res[1], np.ndarray) - 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): - """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 - ) - - 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) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliY(1)) - - def cost(x, y): - return autograd.numpy.hstack(circuit(x, y)) - - jac_fn = qml.jacobian(cost) - res = jac_fn(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], atol=tol, rtol=0) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - # call the Jacobian with new parameters - a = np.array(0.6, requires_grad=True) - b = np.array(0.832, requires_grad=True) - - res = jac_fn(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], 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): - """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 = {} - - 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) - 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): - """Test changing the trainability of parameters changes the - number of differentiation requests made""" - if diff_method != "parameter-shift": - pytest.skip("Test only supports parameter-shift") - - 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) - 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)) - - def loss(a, b): - return np.sum(autograd.numpy.hstack(circuit(a, b))) - - grad_fn = qml.grad(loss) - res = grad_fn(a, b) - - # the tape has reported both arguments as trainable - assert circuit.qtape.trainable_params == [0, 1] - - expected = [-np.sin(a) + np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] - assert np.allclose(res, expected, atol=tol, rtol=0) - - # make the second QNode argument a constant - a = np.array(0.54, requires_grad=True) - b = np.array(0.8, requires_grad=False) - - res = grad_fn(a, b) - - # the tape has reported only the first argument as trainable - assert circuit.qtape.trainable_params == [0] - - 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, 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 - ) - def circuit(a, b, c): - qml.RY(a * c, wires=0) - qml.RZ(b, wires=0) - qml.RX(c + c**2 + np.sin(a), wires=0) - return qml.expval(qml.PauliZ(0)) - - res = qml.jacobian(circuit)(a, b, c) - - assert circuit.qtape.trainable_params == [0, 2] - tape_params = np.array(circuit.qtape.get_parameters()) - assert np.all(tape_params == [a * c, c + c**2 + np.sin(a)]) - - assert isinstance(res, tuple) and len(res) == 2 - assert res[0].shape == () - assert res[1].shape == () - - 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 - ) - def circuit(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1)) - - a = np.array(0.1, requires_grad=False) - b = np.array(0.2, requires_grad=False) - - res = circuit(a, b) - - if diff_method == "finite-diff": - assert circuit.qtape.trainable_params == [] - - assert len(res) == 2 - assert isinstance(res, tuple) - - 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(cost0)(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(cost1)(a, b) - - assert grad == tuple() - - 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 - ) - def circuit(U, a): - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - return qml.expval(qml.PauliZ(0)) - - res = circuit(U, a) - - if diff_method == "finite-diff": - assert circuit.qtape.trainable_params == [1] - - res = qml.grad(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 - ): - """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) - def circuit(data1): - qml.templates.AmplitudeEmbedding(data1, wires=[0, 1]) - return qml.expval(qml.PauliZ(0)) - - grad_fn = qml.grad(circuit, argnum=0) - data1 = np.array([0, 1, 1, 0], requires_grad=False) / np.sqrt(2) - - 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): - """Test that operation and nested tape expansion - is differentiable""" - 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=20) - kwargs = {**kwargs, **spsa_kwargs} - tol = TOL_FOR_SPSA - - 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(dev_name, wires=2) - a = np.array(0.1, requires_grad=False) - p = np.array([0.1, 0.2, 0.3], requires_grad=True) - - @qnode(dev, **kwargs) - def circuit(a, p): - qml.RX(a, wires=0) - U3(p[0], p[1], p[2], wires=0) - return qml.expval(qml.PauliX(0)) - - res = 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 np.allclose(res, expected, atol=tol, rtol=0) - - res = qml.grad(circuit)(a, p) - - assert isinstance(res, np.ndarray) - assert len(res) == 3 - - 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) - - -class TestShotsIntegration: - """Test that the QNode correctly changes shot value, and - remains differentiable.""" - - 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) - - @qnode(dev, diff_method=qml.gradients.param_shift) - 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.target_device, "sample") - - # 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() - - # 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 not dev.shots - 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 = 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) - 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: fix the shot vectors issue - 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 all(r.shape == (3,) for r in res) - - 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) - ) - - 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) - 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) - # 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 - - # 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_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_name, diff_method, grad_on_execution, 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("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) - - @qnode(dev, **kwargs) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[1]) - - res = qml.jacobian(circuit)(x, y) - assert isinstance(res, tuple) and len(res) == 2 - - expected = ( - np.array([-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2]), - np.array([-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]), - ) - 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 - ): - """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("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) - - @qnode(dev, **kwargs) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[0]), qml.probs(wires=[1]) - - res = circuit(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=tol, rtol=0) - - def cost(x, y): - return autograd.numpy.hstack(circuit(x, y)) - - res = qml.jacobian(cost)(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 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): - """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 - ) - 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) - - @qnode(dev, **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=[1]) - - res = circuit(x, y) - assert isinstance(res, tuple) - expected = [np.cos(x), [(1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(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 cost(x, y): - return autograd.numpy.hstack(circuit(x, y)) - - res = qml.jacobian(cost)(x, y) - assert isinstance(res, tuple) - assert 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_ragged_differentiation_variance( - self, interface, dev_name, diff_method, grad_on_execution, 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("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) - - @qnode(dev, **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.probs(wires=[1]) - - res = circuit(x, y) - - expected_var = np.array(np.sin(x) ** 2) - expected_probs = np.array( - [(1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2] - ) - - assert isinstance(res, tuple) - assert len(res) == 2 - - 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 np.allclose(res[1], expected_probs, atol=tol, rtol=0) - - def cost(x, y): - return autograd.numpy.hstack(circuit(x, y)) - - jac = qml.jacobian(cost)(x, y) - assert isinstance(res, tuple) and len(res) == 2 - - expected = ( - np.array([np.sin(2 * 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 isinstance(jac, tuple) - assert len(jac) == 2 - - 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 np.allclose(jac[1], expected[1], atol=tol, rtol=0) - - 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) - - class Template(qml.templates.StronglyEntanglingLayers): - def decomposition(self): - return [qml.templates.StronglyEntanglingLayers(*self.parameters, self.wires)] - - @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) - def circuit2(data, weights): - qml.templates.AngleEmbedding(data, wires=[0, 1]) - Template(weights, wires=[0, 1]) - return qml.expval(qml.PauliX(0)) - - def cost(w1, w2): - c1 = circuit1(w1) - c2 = circuit2(c1, w2) - return np.sum(c2) ** 2 - - w1 = qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=3) - w2 = qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=4) - - weights = [ - np.random.random(w1, requires_grad=True), - np.random.random(w2, requires_grad=True), - ] - - grad_fn = qml.grad(cost) - res = grad_fn(*weights) - - assert len(res) == 2 - - 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) - 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) - - @qnode(dev1, **kwargs) - def circuit1(a, b, c): - qml.RX(a, wires=0) - qml.RX(b, wires=1) - qml.RX(c, wires=2) - qml.CNOT(wires=[0, 1]) - 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) - - @qnode(dev2, interface=interface, diff_method=diff_method) - def circuit2(data, weights): - qml.RX(data[0], wires=0) - qml.RX(data[1], wires=1) - qml.CNOT(wires=[0, 1]) - qml.RZ(weights[0], wires=0) - qml.RZ(weights[1], wires=1) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.PauliX(0) @ qml.PauliY(1)) - - def cost(a, b, c, weights): - return circuit2(circuit1(a, b, c), weights) - - grad_fn = qml.grad(cost) - - # Set the first parameter of circuit1 as non-differentiable. - a = np.array(0.4, requires_grad=False) - - # The remaining free parameters are all differentiable. - b = np.array(0.5, requires_grad=True) - c = np.array(0.1, requires_grad=True) - weights = np.array([0.2, 0.3], requires_grad=True) - - res = grad_fn(a, b, c, weights) - - # Output should have shape [dcost/db, dcost/dc, dcost/dw], - # where b,c are scalars, and w is a vector of length 2. - assert len(res) == 3 - assert res[0].shape == tuple() # scalar - assert res[1].shape == tuple() # scalar - assert res[2].shape == (2,) # vector - - cacbsc = np.cos(a) * np.cos(b) * np.sin(c) - - expected = np.array( - [ - # analytic expression for dcost/db - -np.cos(a) - * np.sin(b) - * np.sin(c) - * np.cos(cacbsc) - * np.sin(weights[0]) - * np.sin(np.cos(a)), - # analytic expression for dcost/dc - np.cos(a) - * np.cos(b) - * np.cos(c) - * np.cos(cacbsc) - * np.sin(weights[0]) - * np.sin(np.cos(a)), - # analytic expression for dcost/dw[0] - np.sin(cacbsc) * np.cos(weights[0]) * np.sin(np.cos(a)), - # analytic expression for dcost/dw[1] - 0, - ] - ) - - # np.hstack 'flattens' the ragged gradient array allowing it - # to be compared with the expected result - assert np.allclose(np.hstack(res), expected, atol=tol, rtol=0) - - if diff_method != "backprop": - # Check that the gradient was computed - # for all parameters in circuit2 - assert circuit2.qtape.trainable_params == [0, 1, 2, 3] - - # Check that the parameter-shift rule was not applied - # 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): - """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, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = np.array([1.0, 2.0], requires_grad=True) - res = circuit(x) - g = qml.grad(circuit)(x) - g2 = qml.grad(lambda x: np.sum(qml.grad(circuit)(x)))(x) - - a, b = x - - expected_res = np.cos(a) * np.cos(b) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - - expected_g2 = [ - -np.cos(a) * np.cos(b) + np.sin(a) * np.sin(b), - np.sin(a) * np.sin(b) - np.cos(a) * np.cos(b), - ] - - if diff_method in {"finite-diff"}: - tol = 10e-2 - - assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - - 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, - max_diff=2, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = np.array([1.0, 2.0], requires_grad=True) - res = circuit(x) - - a, b = x - - expected_res = np.cos(a) * np.cos(b) - - assert isinstance(res, np.ndarray) - assert res.shape == () - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - grad_fn = qml.grad(circuit) - g = grad_fn(x) - - expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)] - - assert isinstance(g, np.ndarray) - assert g.shape == (2,) - assert np.allclose(g, expected_g, atol=tol, rtol=0) - - hess = qml.jacobian(grad_fn)(x) - - expected_hess = [ - [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)], - [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)], - ] - - assert isinstance(hess, np.ndarray) - assert hess.shape == (2, 2) - - if diff_method in {"finite-diff"}: - tol = 10e-2 - - 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 - ): - """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, - max_diff=2, - ) - def circuit(x): - qml.RY(x[0], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = np.array([1.0, 2.0], requires_grad=True) - res = circuit(x) - - a, _ = x - - expected_res = np.cos(a) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - grad_fn = qml.grad(circuit) - - hess = qml.jacobian(grad_fn)(x) - - expected_hess = [ - [-np.cos(a), 0], - [0, 0], - ] - - if diff_method in {"finite-diff"}: - tol = 10e-2 - - 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): - """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, - max_diff=2, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.probs(wires=0) - - x = np.array([1.0, 2.0], requires_grad=True) - res = circuit(x) - - a, b = 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 np.allclose(res, expected_res, atol=tol, rtol=0) - - jac_fn = qml.jacobian(circuit) - jac = jac_fn(x) - - expected_res = [ - [-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)], - ] - - assert isinstance(jac, np.ndarray) - assert jac.shape == (2, 2) - assert np.allclose(jac, expected_res, atol=tol, rtol=0) - - hess = qml.jacobian(jac_fn)(x) - - expected_hess = [ - [ - [-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)], - ], - ] - - assert isinstance(hess, np.ndarray) - assert hess.shape == (2, 2, 2) - - if diff_method in {"finite-diff"}: - tol = 10e-2 - - 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 - ): - """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, - max_diff=2, - ) - def circuit(x): - qml.RX(x[0], wires=0) - qml.RY(x[1], wires=0) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0)) - - def cost_fn(x): - return x @ autograd.numpy.hstack(circuit(x)) - - x = np.array([0.76, -0.87], requires_grad=True) - res = cost_fn(x) - - a, b = x - - expected_res = x @ [np.cos(a) * np.cos(b), np.cos(a) * np.cos(b)] - - assert isinstance(res, np.ndarray) - assert res.shape == () - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - hess = qml.jacobian(qml.grad(cost_fn))(x) - - expected_hess = [ - [ - -(np.cos(b) * ((a + b) * np.cos(a) + 2 * np.sin(a))), - -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b), - ], - [ - -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b), - -(np.cos(a) * ((a + b) * np.cos(b) + 2 * np.sin(b))), - ], - ] - - assert hess.shape == (2, 2) - - if diff_method in {"finite-diff"}: - tol = 10e-2 - - 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 - ): - """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, - max_diff=2, - ) - def circuit(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=0) - return qml.probs(wires=0) - - a = np.array(1.0, requires_grad=True) - b = np.array(2.0, requires_grad=True) - res = 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 np.allclose(res, expected_res, atol=tol, rtol=0) - - jac_fn = qml.jacobian(circuit) - g = jac_fn(a, b) - assert isinstance(g, tuple) and len(g) == 2 - - expected_g = ( - [-0.5 * np.sin(a) * np.cos(b), 0.5 * np.sin(a) * np.cos(b)], - [-0.5 * np.cos(a) * np.sin(b), 0.5 * np.cos(a) * np.sin(b)], - ) - assert g[0].shape == (2,) - assert np.allclose(g[0], expected_g[0], atol=tol, rtol=0) - - 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] - 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 - assert isinstance(hess_b, tuple) and len(hess_b) == 2 - - exp_hess_a = ( - [-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)], - ) - exp_hess_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)], - ) - - if diff_method in {"finite-diff"}: - tol = 10e-2 - - for hess, exp_hess in zip([hess_a, hess_b], [exp_hess_a, exp_hess_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): - """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, - max_diff=2, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - qml.RY(x[0], wires=1) - qml.RX(x[1], wires=1) - return qml.expval(qml.PauliZ(0)), qml.probs(wires=1) - - x = np.array([1.0, 2.0], requires_grad=True) - - 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)) - - def cost_fn(x): - return autograd.numpy.hstack(circuit(x)) - - jac_fn = qml.jacobian(cost_fn) - - hess = qml.jacobian(jac_fn)(x) - expected_hess = [ - [ - [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)], - [np.sin(a) * np.sin(b), -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)], - ], - [ - [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)], - ], - ] - - if diff_method in {"finite-diff"}: - tol = 10e-2 - - assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - - 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") - - 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 - ) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.state() - - def cost_fn(x, y): - res = circuit(x, y) - assert res.dtype is np.dtype("complex128") - probs = np.abs(res) ** 2 - return probs[0] + probs[2] - - res = cost_fn(x, y) - - if diff_method not in {"backprop"}: - pytest.skip("Test only supports backprop") - - res = qml.jacobian(cost_fn)(x, y) - expected = np.array([-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2]) - assert isinstance(res, tuple) - assert len(res) == 2 - - assert isinstance(res[0], np.ndarray) - assert res[0].shape == () - assert isinstance(res[1], np.ndarray) - assert res[1].shape == () - 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): - """Test that the variance of a projector is correctly returned""" - kwargs = dict( - diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - 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) - - @qnode(dev, **kwargs) - def circuit(x, y): - qml.RX(x, wires=0) - qml.RY(y, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.var(qml.Projector(P, wires=0) @ qml.PauliX(1)) - - 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 np.allclose(res, expected, atol=tol, rtol=0) - - jac = qml.jacobian(circuit)(x, y) - expected = np.array( - [ - [ - 0.5 * np.sin(x) * (np.cos(x / 2) ** 2 + np.cos(2 * y) * np.sin(x / 2) ** 2), - -2 * np.cos(y) * np.sin(x / 2) ** 4 * np.sin(y), - ] - ] - ) - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], np.ndarray) - assert jac[0].shape == () - - assert isinstance(jac[1], np.ndarray) - assert jac[1].shape == () - - assert np.allclose(jac, expected, atol=tol, rtol=0) - - -@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) - - # 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), - ] - ] - ) - 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)) - - spy = mocker.spy(dev.target_device, "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_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") - 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) - - 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) - def circuit(x, y): - qml.Hadamard(wires=0) - PhaseShift(x, wires=0) - PhaseShift(2 * y, wires=0) - return qml.expval(qml.PauliX(0)) - - x = np.array(0.5, requires_grad=True) - y = np.array(0.7, requires_grad=False) - 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 - ): - """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, - ) - if diff_method in ["adjoint", "hadamard"]: - pytest.skip("The diff method requested does not yet support Hamiltonians") - elif diff_method == "spsa": - 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) - def circuit(data, weights, coeffs): - weights = weights.reshape(1, -1) - qml.templates.AngleEmbedding(data, wires=[0, 1]) - qml.templates.BasicEntanglerLayers(weights, wires=[0, 1]) - return qml.expval(qml.Hamiltonian(coeffs, obs)) - - d = np.array([0.1, 0.2], requires_grad=False) - w = np.array([0.654, -0.734], requires_grad=True) - c = np.array([-0.6543, 0.24, 0.54], requires_grad=True) - - # test output - 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) - - # test gradients - 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]), - ] - expected_c = [0, -np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]), np.cos(d[1] + w[1])] - assert np.allclose(grad[0], expected_w, atol=tol) - 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 == "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) - else: - 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) - 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) - - @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 - ): - """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 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, - ) - 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( - dev, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - max_diff=max_diff, - **gradient_kwargs, - ) - def circuit(data, weights, coeffs): - weights = weights.reshape(1, -1) - qml.templates.AngleEmbedding(data, wires=[0, 1]) - qml.templates.BasicEntanglerLayers(weights, wires=[0, 1]) - H = qml.Hamiltonian(coeffs, obs) - H.compute_grouping() - return qml.expval(H) - - d = np.array([0.1, 0.2], requires_grad=False) - w = np.array([0.654, -0.734], requires_grad=True) - c = np.array([-0.6543, 0.24, 0.54], requires_grad=True) - - # test output - 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 = 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]), - ] - expected_c = [0, -np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]), np.cos(d[1] + w[1])] - assert np.allclose(grad[0], expected_w, atol=tol) - assert np.allclose(grad[1], expected_c, atol=tol) - - # test second-order derivatives - if diff_method == "parameter-shift" and max_diff == 2: - 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) - 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""" - - def test_backprop_error(self): - """Test that sampling in backpropagation grad_on_execution raises an error""" - 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" - ): - 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 = qml.device("default.qubit.legacy", wires=2, shots=n_sample) - - @qnode(dev) - def circuit(): - qml.RX(0.54, wires=0) - return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - - res = circuit() - - assert isinstance(res, tuple) - assert len(res) == 2 - - assert res[0].shape == (10,) - assert isinstance(res[0], np.ndarray) - - assert res[1].shape == (10,) - assert isinstance(res[1], np.ndarray) - - 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") - 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() - - 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 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) - - @qnode(dev) - def circuit(): - qml.RX(0.54, wires=0) - - return qml.sample(qml.PauliZ(0)) - - result = circuit() - - assert isinstance(result, np.ndarray) - 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(dev) - def circuit(): - return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1)), qml.sample(qml.PauliZ(2)) - - 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,) - assert isinstance(result[0], np.ndarray) - - assert result[1].shape == (10,) - assert isinstance(result[1], np.ndarray) - - assert result[2].shape == (10,) - assert isinstance(result[2], np.ndarray) - - -@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.""" - - # 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 - - 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) - return qml.expval(qml.PauliZ(0)) - - a = np.array(0.1, requires_grad=True) - - grad = qml.grad(circuit)(a) - - import sys - - 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 - - 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) - return qml.expval(qml.PauliZ(0)) - - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - grad = qml.grad(circuit)(a, b) - - assert isinstance(grad, tuple) - assert len(grad) == 2 - assert grad[0].shape == () - assert grad[1].shape == () - - def test_grad_single_measurement_multiple_param_array( - 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 - - 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) - return qml.expval(qml.PauliZ(0)) - - a = np.array([0.1, 0.2], requires_grad=True) - - grad = qml.grad(circuit)(a) - - assert isinstance(grad, np.ndarray) - assert len(grad) == 2 - assert grad.shape == (2,) - - def test_jacobian_single_measurement_param_probs( - 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") - - 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) - return qml.probs(wires=[0, 1]) - - a = np.array(0.1, requires_grad=True) - - jac = qml.jacobian(circuit)(a) - - assert isinstance(jac, np.ndarray) - assert jac.shape == (4,) - - def test_jacobian_single_measurement_probs_multiple_param( - 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") - - 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) - return qml.probs(wires=[0, 1]) - - a = np.array(0.1, requires_grad=True) - b = np.array(0.2, requires_grad=True) - - jac = qml.jacobian(circuit)(a, b) - - assert isinstance(jac, tuple) - - assert isinstance(jac[0], np.ndarray) - assert jac[0].shape == (4,) - - assert isinstance(jac[1], np.ndarray) - assert jac[1].shape == (4,) - - def test_jacobian_single_measurement_probs_multiple_param_single_array( - 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") - - 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) - return qml.probs(wires=[0, 1]) - - a = np.array([0.1, 0.2], requires_grad=True) - jac = qml.jacobian(circuit)(a) - - assert isinstance(jac, np.ndarray) - assert jac.shape == (4, 2) - - def test_jacobian_multiple_measurement_single_param( - self, dev_name, diff_method, grad_on_execution - ): - """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) - 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, requires_grad=True) - - def cost(x): - return anp.hstack(circuit(x)) - - jac = qml.jacobian(cost)(a) - - assert isinstance(jac, np.ndarray) - assert jac.shape == (5,) - - def test_jacobian_multiple_measurement_multiple_param( - self, dev_name, diff_method, grad_on_execution - ): - """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) - 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(x, y): - return anp.hstack(circuit(x, y)) - - jac = qml.jacobian(cost)(a, b) - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], np.ndarray) - assert jac[0].shape == (5,) - - assert isinstance(jac[1], np.ndarray) - assert jac[1].shape == (5,) - - def test_jacobian_multiple_measurement_multiple_param_array( - self, dev_name, diff_method, grad_on_execution - ): - """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) - 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], requires_grad=True) - - def cost(x): - return anp.hstack(circuit(x)) - - jac = qml.jacobian(cost)(a) - - assert isinstance(jac, np.ndarray) - assert jac.shape == (5, 2) - - 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.") - - 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) - 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): - return anp.hstack(qml.grad(circuit)(x, y)) - - hess = qml.jacobian(cost)(par_0, par_1) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], np.ndarray) - assert hess[0].shape == (2,) - - 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): - """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) - 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 = qml.jacobian(qml.grad(circuit))(params) - - assert isinstance(hess, np.ndarray) - assert hess.shape == (2, 2) - - 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.") - elif diff_method == "hadamard": - pytest.skip("Hadamard gradient does not support variances.") - - 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) - 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): - return anp.hstack(qml.grad(circuit)(x, y)) - - hess = qml.jacobian(cost)(par_0, par_1) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], np.ndarray) - assert hess[0].shape == (2,) - - 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): - """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) - 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 = qml.jacobian(qml.grad(circuit))(params) - - 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): - """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) - 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 circuit_stack(x, y): - return anp.hstack(circuit(x, y)) - - def cost(x, y): - return anp.hstack(qml.jacobian(circuit_stack)(x, y)) - - hess = qml.jacobian(cost)(par_0, par_1) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], np.ndarray) - assert hess[0].shape == (6,) - - assert isinstance(hess[1], np.ndarray) - assert hess[1].shape == (6,) - - def test_hessian_expval_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 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) - 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): - return anp.hstack(circuit(x)) - - hess = qml.jacobian(qml.jacobian(cost))(params) - - assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) - - 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.") - elif diff_method == "hadamard": - pytest.skip("Hadamard gradient does not support variances.") - - 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) - 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]) - - def circuit_stack(x, y): - return anp.hstack(circuit(x, y)) - - def cost(x, y): - return anp.hstack(qml.jacobian(circuit_stack)(x, y)) - - hess = qml.jacobian(cost)(par_0, par_1) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], np.ndarray) - assert hess[0].shape == (6,) - - 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): - """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) - 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]) - - def cost(x): - return anp.hstack(circuit(x)) - - hess = qml.jacobian(qml.jacobian(cost))(params) - - assert isinstance(hess, np.ndarray) - assert hess.shape == (3, 2, 2) - - -@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 = qml.device(dev_name, wires=1) - - @qml.qnode(dev, interface="autograd") - def circuit(): - qml.Hadamard(wires=0) - return qml.state() - - res = circuit() - assert isinstance(res, np.tensor) 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 deleted file mode 100644 index 24d1d824ced..00000000000 --- a/tests/interfaces/legacy_devices_integration/test_autograd_qnode_shot_vector_legacy.py +++ /dev/null @@ -1,658 +0,0 @@ -# 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_jax_jit_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_jit_legacy.py deleted file mode 100644 index eaa0cd95cf8..00000000000 --- a/tests/interfaces/legacy_devices_integration/test_jax_jit_legacy.py +++ /dev/null @@ -1,899 +0,0 @@ -# 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.target_device, "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.target_device._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, 1 evaluation required. - jax.grad(cost)(params, cache=None) - assert dev.num_executions == 1 - - # With caching, also 1 evaluation required. - dev.target_device._num_executions = 0 - jac_fn = jax.grad(cost) - jac_fn(params, cache=True) - assert dev.num_executions == 1 - - -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.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) - 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) - - -@pytest.mark.xfail(reason="Need to figure out how to handle this case in a less ambiguous manner") -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/legacy_devices_integration/test_jax_jit_qnode_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_jit_qnode_legacy.py deleted file mode 100644 index 414b6c15506..00000000000 --- a/tests/interfaces/legacy_devices_integration/test_jax_jit_qnode_legacy.py +++ /dev/null @@ -1,2955 +0,0 @@ -# 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. -"""Integration tests for using the JAX-JIT interface with a QNode""" -# pylint: disable=too-many-arguments,too-few-public-methods -import pytest - -import pennylane as qml -from pennylane import numpy as np -from pennylane import qnode - -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_qubit_device_and_diff_method = [ - ["auto"] + inner_list for inner_list in qubit_device_and_diff_method -] + [["jax-jit"] + inner_list for inner_list in qubit_device_and_diff_method] - -pytestmark = pytest.mark.jax - -jax = pytest.importorskip("jax") -jax.config.update("jax_enable_x64", True) - -TOL_FOR_SPSA = 1.0 -SEED_FOR_SPSA = 32651 -H_FOR_SPSA = 0.05 - - -@pytest.mark.parametrize( - "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_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 - ) - 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) - jax.jit(circuit)(a) - - assert circuit.interface == interface - - # the tape is able to deduce trainable parameters - assert circuit.qtape.trainable_params == [0] - - # gradients should work - grad = jax.jit(jax.grad(circuit))(a) - assert isinstance(grad, jax.Array) - assert grad.shape == () - - 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": - pytest.skip("Test only supports parameter-shift") - - 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, - ) - def circuit(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliY(1)])) - - grad_fn = jax.jit(jax.grad(circuit, argnums=[0, 1])) - res = grad_fn(a, b) - - # the tape has reported both arguments as trainable - assert circuit.qtape.trainable_params == [0, 1] - - expected = [-np.sin(a) + np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] - assert np.allclose(res, expected, atol=tol, rtol=0) - - # make the second QNode argument a constant - grad_fn = jax.grad(circuit, argnums=0) - res = grad_fn(a, b) - - # the tape has reported only the first argument as trainable - assert circuit.qtape.trainable_params == [0] - - 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): - """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 - ) - def circuit(a, b, c): - qml.RY(a * c, wires=0) - qml.RZ(b, wires=0) - qml.RX(c + c**2 + jax.numpy.sin(a), wires=0) - return qml.expval(qml.PauliZ(0)) - - res = jax.grad(circuit, argnums=[0, 2])(a, b, c) - - if diff_method == "finite-diff": - assert circuit.qtape.trainable_params == [0, 2] - - assert len(res) == 2 - - 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 - ) - def circuit(U, a): - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - return qml.expval(qml.PauliZ(0)) - - res = jax.grad(circuit, argnums=1)(U, a) - assert np.allclose(res, np.sin(a), atol=tol, rtol=0) - - 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): - """Test that operation and nested tape expansion - is differentiable""" - - gradient_kwargs = {} - if diff_method == "spsa": - gradient_kwargs = { - "sampler_rng": SEED_FOR_SPSA, - "num_directions": 10, - } - tol = TOL_FOR_SPSA - - 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), - ] - - 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]) - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **gradient_kwargs, - ) - def circuit(a, p): - qml.RX(a, wires=0) - U3(p[0], p[1], p[2], wires=0) - return qml.expval(qml.PauliX(0)) - - res = jax.jit(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 np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.jit(jax.grad(circuit, argnums=1))(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=tol, rtol=0) - - 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, - 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) - return qml.expval(qml.PauliZ(0)) - - if diff_method in {"finite-diff", "parameter-shift", "spsa"} and interface == "jax-jit": - # No jax.jacobian support for call - pytest.xfail(reason="batching rules are implemented only for id_tap, not for call.") - - jax.jit(jax.jacobian(circuit))(a) - - -@pytest.mark.parametrize( - "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_name, diff_method, grad_on_execution, interface, tol): - """Test jacobian calculation""" - - gradient_kwargs = {} - - if diff_method == "spsa": - 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, - **gradient_kwargs, - ) - 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)) - - res = jax.jit(circuit)(a, b) - - assert circuit.qtape.trainable_params == [0, 1] - assert isinstance(res, tuple) - assert len(res) == 2 - - expected = [np.cos(a), -np.cos(a) * np.sin(b)] - assert np.allclose(res[0], expected[0], atol=tol, rtol=0) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - res = jax.jit(jax.jacobian(circuit, argnums=[0, 1]))(a, b) - assert circuit.qtape.trainable_params == [0, 1] - - expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - assert isinstance(res, tuple) - assert len(res) == 2 - - assert isinstance(res[0], tuple) - assert isinstance(res[0][0], jax.numpy.ndarray) - assert res[0][0].shape == () - assert np.allclose(res[0][0], expected[0][0], atol=tol, rtol=0) - assert isinstance(res[0][1], jax.numpy.ndarray) - assert res[0][1].shape == () - assert np.allclose(res[0][1], expected[0][1], atol=tol, rtol=0) - - assert isinstance(res[1], tuple) - assert isinstance(res[1][0], jax.numpy.ndarray) - assert res[1][0].shape == () - assert np.allclose(res[1][0], expected[1][0], atol=tol, rtol=0) - assert isinstance(res[1][1], jax.numpy.ndarray) - 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): - """Test jacobian calculation when no prior circuit evaluation has been performed""" - - gradient_kwargs = {} - if diff_method == "spsa": - 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, - **gradient_kwargs, - ) - 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)) - - jac_fn = jax.jit(jax.jacobian(circuit, argnums=[0, 1])) - - res = jac_fn(a, b) - - assert isinstance(res, tuple) - assert len(res) == 2 - - 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) - - # call the Jacobian with new parameters - a = jax.numpy.array(0.6) - b = jax.numpy.array(0.832) - - res = jac_fn(a, b) - - assert isinstance(res, tuple) - assert len(res) == 2 - - 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) - - 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""" - gradient_kwargs = {} - 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) - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **gradient_kwargs, - ) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[1]) - - res = jax.jit(jax.jacobian(circuit, argnums=[0, 1]))(x, y) - - expected = np.array( - [ - [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2], - [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2], - ] - ) - - assert isinstance(res, tuple) - assert len(res) == 2 - - assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == (2,) - - assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) - - 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): - """Tests correct output shape and evaluation for a tape - with multiple prob outputs""" - gradient_kwargs = {} - 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) - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **gradient_kwargs, - ) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[0]), qml.probs(wires=[1, 2]) - - res = circuit(x, y) - - assert isinstance(res, tuple) - assert len(res) == 2 - - expected = [ - [np.cos(x / 2) ** 2, np.sin(x / 2) ** 2], - [(1 + np.cos(x) * np.cos(y)) / 2, 0, (1 - np.cos(x) * np.cos(y)) / 2, 0], - ] - - assert isinstance(res[0], jax.numpy.ndarray) - 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,) # pylint:disable=comparison-with-callable - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - jac = jax.jit(jax.jacobian(circuit, argnums=[0, 1]))(x, y) - expected_0 = np.array( - [ - [-np.sin(x) / 2, np.sin(x) / 2], - [0, 0], - ] - ) - - expected_1 = np.array( - [ - [-np.cos(y) * np.sin(x) / 2, 0, np.sin(x) * np.cos(y) / 2, 0], - [-np.cos(x) * np.sin(y) / 2, 0, np.cos(x) * np.sin(y) / 2, 0], - ] - ) - - assert isinstance(jac, tuple) - assert isinstance(jac[0], tuple) - - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == (2,) - assert np.allclose(jac[0][0], expected_0[0], atol=tol, rtol=0) - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == (2,) - assert np.allclose(jac[0][1], expected_0[1], atol=tol, rtol=0) - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == (4,) - - assert np.allclose(jac[1][0], expected_1[0], atol=tol, rtol=0) - assert isinstance(jac[1][1], jax.numpy.ndarray) - 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): - """Tests correct output shape and evaluation for a tape - with prob and expval outputs""" - gradient_kwargs = {} - 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) - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **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=[1]) - - res = jax.jit(circuit)(x, y) - expected = [np.cos(x), [(1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2]] - - assert isinstance(res, tuple) - assert len(res) == 2 - - assert isinstance(res[0], jax.numpy.ndarray) - 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,) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - jac = jax.jit(jax.jacobian(circuit, argnums=[0, 1]))(x, y) - expected = [ - [-np.sin(x), 0], - [ - [-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2], - [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], - ], - ] - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert np.allclose(jac[0][0], expected[0][0], atol=tol, rtol=0) - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - assert np.allclose(jac[0][1], expected[0][1], atol=tol, rtol=0) - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == (2,) - assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - assert isinstance(jac[1][1], jax.numpy.ndarray) - assert jac[1][1].shape == (2,) - 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 - ): - """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 - 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) - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **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=[1]) - - jac = jax.jit(jax.jacobian(circuit, argnums=[0]))(x, y) - - expected = [ - [-np.sin(x), 0], - [ - [-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2], - [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], - ], - ] - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 1 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert np.allclose(jac[0][0], expected[0][0], atol=tol, rtol=0) - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 1 - assert isinstance(jac[1][0], jax.numpy.ndarray) - 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): - """Tests correct output shape and evaluation for a tape - with prob and variance outputs""" - gradient_kwargs = {} - 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": 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) - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **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.probs(wires=[1]) - - res = jax.jit(circuit)(x, y) - - expected = [ - np.sin(x) ** 2, - [(1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2], - ] - - assert isinstance(res[0], jax.numpy.ndarray) - 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,) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - jac = jax.jit(jax.jacobian(circuit, argnums=[0, 1]))(x, y) - expected = [ - [2 * np.cos(x) * np.sin(x), 0], - [ - [-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2], - [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], - ], - ] - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert np.allclose(jac[0][0], expected[0][0], atol=tol, rtol=0) - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - assert np.allclose(jac[0][1], expected[0][1], atol=tol, rtol=0) - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == (2,) - assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - assert isinstance(jac[1][1], jax.numpy.ndarray) - assert jac[1][1].shape == (2,) - assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) - - -@pytest.mark.parametrize("interface", ["auto", "jax", "jax-jit"]) -class TestShotsIntegration: - """Test that the QNode correctly changes shot value, and - 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) - - @jax.jit - @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)), 1) - - 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(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.expval(qml.PauliY(1)) - - spy = mocker.spy(dev.target_device, "sample") - - # 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() - - # 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 not dev.shots - 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(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)) - - # 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 == qml.measurements.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): - """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") - 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)) - - # 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 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" - - 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_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_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.") - - 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 - ) - def circuit(): - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - - res = jax.jit(circuit)() - - assert isinstance(res, tuple) - - assert isinstance(res[0], jax.Array) - assert res[0].shape == (10,) - assert isinstance(res[1], jax.Array) - assert res[1].shape == (10,) - - 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.") - - 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 - ) - def circuit(): - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - return qml.counts(qml.PauliZ(0)), qml.counts(qml.PauliX(1)) - - if interface == "jax-jit": - with pytest.raises( - NotImplementedError, match="The JAX-JIT interface doesn't support qml.counts." - ): - jax.jit(circuit)() - else: - res = jax.jit(circuit)() - - assert isinstance(res, tuple) - - assert isinstance(res[0], dict) - assert len(res[0]) == 2 - assert isinstance(res[1], dict) - assert len(res[1]) == 2 - - def test_chained_qnodes(self, dev_name, diff_method, grad_on_execution, 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) - - 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 - ) - 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 - ) - def circuit2(data, weights): - qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) - Template(weights, wires=[0, 1]) - return qml.expval(qml.PauliX(0)) - - def cost(weights): - w1, w2 = weights - c1 = circuit1(w1) - c2 = circuit2(c1, w2) - return jax.numpy.sum(c2) ** 2 - - w1 = qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=3) - w2 = qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=4) - - weights = [ - jax.numpy.array(np.random.random(w1)), - jax.numpy.array(np.random.random(w2)), - ] - - grad_fn = jax.jit(jax.grad(cost)) - res = grad_fn(weights) - - assert len(res) == 2 - - -@pytest.mark.parametrize( - "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_name, diff_method, grad_on_execution, interface, tol): - """Test second derivative calculation of a scalar-valued QNode""" - - gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not second derivative.") - elif diff_method == "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, - max_diff=2, - **gradient_kwargs, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = jax.numpy.array([1.0, 2.0]) - res = circuit(x) - g = jax.jit(jax.grad(circuit))(x) - g2 = jax.jit(jax.grad(lambda x: jax.numpy.sum(jax.grad(circuit)(x))))(x) - - a, b = x - - expected_res = np.cos(a) * np.cos(b) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - - expected_g2 = [ - -np.cos(a) * np.cos(b) + np.sin(a) * np.sin(b), - np.sin(a) * np.sin(b) - np.cos(a) * np.cos(b), - ] - if diff_method == "finite-diff": - assert np.allclose(g2, expected_g2, atol=10e-2, rtol=0) - else: - assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - - 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": - pytest.skip("Adjoint does not support second derivative.") - elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "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, - max_diff=2, - **gradient_kwargs, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = jax.numpy.array([1.0, 2.0]) - res = jax.jit(circuit)(x) - - a, b = x - - expected_res = np.cos(a) * np.cos(b) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - grad_fn = jax.jit(jax.grad(circuit)) - g = grad_fn(x) - - expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - - hess = jax.jit(jax.jacobian(grad_fn))(x) - - expected_hess = [ - [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)], - [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)], - ] - if diff_method == "finite-diff": - assert np.allclose(hess, expected_hess, atol=10e-2, rtol=0) - 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): - """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": - gradient_kwargs = { - "h": H_FOR_SPSA, - "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, - max_diff=2, - **gradient_kwargs, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.probs(wires=0) - - x = jax.numpy.array([1.0, 2.0]) - res = circuit(x) - - a, b = x - - expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - jac_fn = jax.jit(jax.jacobian(circuit)) - g = jac_fn(x) - - expected_g = [ - [-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)], - ] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - - hess = jax.jit(jax.jacobian(jac_fn))(x) - - expected_hess = [ - [ - [-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)], - ], - ] - if diff_method == "finite-diff": - assert np.allclose(hess, expected_hess, atol=10e-2, rtol=0) - else: - 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 - ): - """Test hessian calculation of a vector valued QNode with post-processing""" - gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support second derivative.") - elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "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, - max_diff=2, - **gradient_kwargs, - ) - def circuit(x): - qml.RX(x[0], wires=0) - qml.RY(x[1], wires=0) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0)) - - def cost_fn(x): - return x @ jax.numpy.array(circuit(x)) - - x = jax.numpy.array([0.76, -0.87]) - res = jax.jit(cost_fn)(x) - - a, b = x - - expected_res = x @ jax.numpy.array([np.cos(a) * np.cos(b), np.cos(a) * np.cos(b)]) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - grad_fn = jax.jit(jax.grad(cost_fn)) - g = grad_fn(x) - - expected_g = [ - np.cos(b) * (np.cos(a) - (a + b) * np.sin(a)), - np.cos(a) * (np.cos(b) - (a + b) * np.sin(b)), - ] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - hess = jax.jit(jax.jacobian(grad_fn))(x) - - expected_hess = [ - [ - -(np.cos(b) * ((a + b) * np.cos(a) + 2 * np.sin(a))), - -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b), - ], - [ - -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b), - -(np.cos(a) * ((a + b) * np.cos(b) + 2 * np.sin(b))), - ], - ] - - if diff_method == "finite-diff": - assert np.allclose(hess, expected_hess, atol=10e-2, rtol=0) - else: - 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 - ): - """Test hessian calculation of a vector valued QNode that has separate input arguments""" - gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support second derivative.") - elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "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, - max_diff=2, - **gradient_kwargs, - ) - def circuit(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=0) - return qml.probs(wires=0) - - a = jax.numpy.array(1.0) - b = jax.numpy.array(2.0) - res = 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 np.allclose(res, expected_res, atol=tol, rtol=0) - - jac_fn = jax.jit(jax.jacobian(circuit, argnums=[0, 1])) - g = jac_fn(a, b) - - expected_g = np.array( - [ - [-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)], - ] - ) - 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( - [ - [ - [-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.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)], - ], - ] - ) - if diff_method == "finite-diff": - assert np.allclose(hess, expected_hess, atol=10e-2, rtol=0) - else: - assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - - 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 - - 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 - ) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.state() - - def cost_fn(x, y): - res = circuit(x, y) - assert res.dtype is np.dtype("complex128") # pylint:disable=no-member - probs = jax.numpy.abs(res) ** 2 - return probs[0] + probs[2] - - res = jax.jit(cost_fn)(x, y) - - if diff_method not in {"backprop"}: - pytest.skip("Test only supports backprop") - - res = jax.jit(jax.grad(cost_fn, argnums=[0, 1]))(x, y) - expected = np.array([-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2]) - 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): - """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("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 - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **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.Projector(P, wires=0) @ qml.PauliX(1)) - - res = jax.jit(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 np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.jit(jax.grad(circuit, argnums=[0, 1]))(x, y) - expected = np.array( - [ - 0.5 * np.sin(x) * (np.cos(x / 2) ** 2 + np.cos(2 * y) * np.sin(x / 2) ** 2), - -2 * np.cos(y) * np.sin(x / 2) ** 4 * np.sin(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.target_device, "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 - with the JAX interface""" - - @pytest.mark.parametrize("max_diff", [1, 2]) - def test_gradient_expansion_trainable_only( - 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 - - 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, - interface=interface, - ) - def circuit(x, y): - qml.Hadamard(wires=0) - PhaseShift(x, wires=0) - PhaseShift(2 * y, wires=0) - return qml.expval(qml.PauliX(0)) - - 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 - ): - """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 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") - elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "num_directions": 20, - "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), - } - 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)] - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - max_diff=max_diff, - **gradient_kwargs, - ) - def circuit(data, weights, coeffs): - weights = weights.reshape(1, -1) - qml.templates.AngleEmbedding(data, wires=[0, 1]) - qml.templates.BasicEntanglerLayers(weights, wires=[0, 1]) - return qml.expval(qml.Hamiltonian(coeffs, obs)) - - d = jax.numpy.array([0.1, 0.2]) - w = jax.numpy.array([0.654, -0.734]) - c = jax.numpy.array([-0.6543, 0.24, 0.54]) - - # test output - 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() - - # test gradients - 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]), - ] - expected_c = [0, -np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]), np.cos(d[1] + w[1])] - assert np.allclose(grad[0], expected_w, atol=tol) - assert np.allclose(grad[1], expected_c, atol=tol) - - # TODO: Add parameter shift when the bug with trainable params and hamiltonian_grad is solved. - # test second-order derivatives - if diff_method in "backprop" and max_diff == 2: - grad2_c = jax.jacobian(jax.grad(circuit, argnums=[2]), argnums=[2])(d, w, c) - assert np.allclose(grad2_c, 0, atol=tol) - - grad2_w_c = jax.jacobian(jax.grad(circuit, argnums=[1]), argnums=[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]), - -np.sin(d[1] + w[1]), - ] - 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 - ): - """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 diff_method in ("adjoint", "backprop", "finite-diff"): - pytest.skip("The adjoint and backprop methods do not yet support sampling") - elif diff_method == "hadamard": - pytest.skip("The Hadamard method does not yet support Hamiltonians") - elif diff_method == "spsa": - 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( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - max_diff=max_diff, - **gradient_kwargs, - ) - def circuit(data, weights, coeffs): - weights = weights.reshape(1, -1) - qml.templates.AngleEmbedding(data, wires=[0, 1]) - qml.templates.BasicEntanglerLayers(weights, wires=[0, 1]) - H = qml.Hamiltonian(coeffs, obs) - H.compute_grouping() - return qml.expval(H) - - d = jax.numpy.array([0.1, 0.2]) - w = jax.numpy.array([0.654, -0.734]) - c = jax.numpy.array([-0.6543, 0.24, 0.54]) - - # test output - 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) - 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]), - ] - expected_c = [0, -np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]), np.cos(d[1] + w[1])] - assert np.allclose(grad[0], expected_w, atol=tol) - assert np.allclose(grad[1], expected_c, atol=tol) - - # TODO: Fix hamiltonian grad for parameter shift and jax - # # test second-order derivatives - # if diff_method == "parameter-shift" and max_diff == 2: - - # grad2_c = jax.jacobian(jax.grad(circuit, argnum=2), argnum=2)(d, w, c) - # assert np.allclose(grad2_c, 0, atol=tol) - - # grad2_w_c = jax.jacobian(jax.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]), - # -np.sin(d[1] + w[1]), - # ] - # 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 - ): - """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.""" - 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( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - 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) - - res = _measure_operator() - return res - - 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_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.""" - 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 - 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( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - 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) - - 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) - - -jacobian_fn = [jax.jacobian, jax.jacrev, jax.jacfwd] - - -@pytest.mark.parametrize("jacobian", jacobian_fn) -@pytest.mark.parametrize( - "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_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 diff_method == "spsa": - gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} - tol = TOL_FOR_SPSA - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **gradient_kwargs, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = jax.numpy.array([1.0, 2.0]) - res = circuit(x) - g = jax.jit(jacobian(circuit))(x) - - a, b = x - - expected_res = np.cos(a) * np.cos(b) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - - @pytest.mark.filterwarnings( - "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): - """Test that the jax device works with qml.Hermitian and jitting even - when shots>0. - - Note: before a fix, the cases of shots=10 and shots=1000 were failing due - 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.") - - if diff_method == "adjoint": - pytest.skip("Computing the gradient for observables is not supported with adjoint.") - - 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 - ) - def circ(projector): - return qml.expval(qml.Hermitian(projector, wires=range(2))) - - result = jax.jit(circ)(projector) - assert jax.numpy.allclose(result, 1) - - @pytest.mark.filterwarnings( - "ignore:Requested adjoint differentiation to be computed with finite shots." - ) - @pytest.mark.parametrize("shots", [10, 1000]) - def test_probs_obs_none( - 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 - ) - def circuit(): - return qml.probs(wires=0) - - assert jax.numpy.allclose(circuit(), jax.numpy.array([1.0, 0.0])) - - # @pytest.mark.xfail( - # 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 - ): - """Test derivative calculation of a scalar valued QNode with respect - to a subset of arguments""" - if diff_method == "spsa" and not grad_on_execution: - pytest.xfail(reason="incorrect jacobian results") - - 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 - ) - def circuit(a, b, c): - qml.RY(a, wires=0) - qml.RX(b, wires=0) - qml.RZ(c, wires=0) - return qml.expval(qml.PauliZ(0)) - - res = jax.jit(circuit)(a, b, 0.0) - expected_res = np.cos(a) * np.cos(b) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - g = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b, 0.0) - expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)] - 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 - ): - """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.") - elif diff_method == "spsa": - gradient_kwargs = {"h": H_FOR_SPSA, "sampler_rng": np.random.default_rng(SEED_FOR_SPSA)} - tol = TOL_FOR_SPSA - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **gradient_kwargs, - ) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[1]) - - def cost(x, y, idx): - res = circuit(x, y) - return res[idx] # pylint:disable=unsubscriptable-object - - x = jax.numpy.array(1.0) - y = jax.numpy.array(2.0) - expected_g = ( - np.array([-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2]), - np.array([-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2]), - ) - - idx = 0 - g0 = jax.jit(jacobian(cost, argnums=0))(x, y, idx) - g1 = jax.jit(jacobian(cost, argnums=1))(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) - - idx = 1 - g0 = jax.jit(jacobian(cost, argnums=0))(x, y, idx) - g1 = jax.jit(jacobian(cost, argnums=1))(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) - - def test_matrix_parameter( - self, dev_name, diff_method, grad_on_execution, 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) - - @qml.qnode( - dev, diff_method=diff_method, interface=interface, grad_on_execution=grad_on_execution - ) - def circ(p, U): - qml.QubitUnitary(U, wires=0) - qml.RY(p, wires=0) - return qml.expval(qml.PauliZ(0)) - - p = jax.numpy.array(0.1) - U = jax.numpy.array([[0, 1], [1, 0]]) - 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))) - res = jac_fn(p, U) - assert np.allclose(res, np.sin(p), atol=tol, rtol=0) - - -@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 -) -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 - ): - """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 - ) - 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) - - grad = jax.jit(jacobian(circuit))(a) - - 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 - ): - """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 - ) - 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) - - grad = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) - - assert isinstance(grad, tuple) - assert len(grad) == 2 - assert grad[0].shape == () - assert grad[1].shape == () - - def test_grad_single_measurement_multiple_param_array( - 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 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 - ) - 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]) - - 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_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 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) - - @qnode( - 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 = jax.numpy.array(0.1) - - 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_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 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) - - @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=0) - return qml.probs(wires=[0, 1]) - - a = jax.numpy.array(0.1) - b = jax.numpy.array(0.2) - - jac = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) - - assert isinstance(jac, tuple) - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == (4,) - - assert isinstance(jac[1], jax.numpy.ndarray) - 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 - ): - """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 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 - ) - 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 = 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_name, diff_method, grad_on_execution, jacobian, shots, interface - ): - """The jacobian 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 = 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 - ) - 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 = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) - - assert isinstance(jac, tuple) - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == () - assert isinstance(jac[1][1], jax.numpy.ndarray) - assert jac[1][1].shape == () - - def test_jacobian_expval_expval_multiple_params_array( - 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") - 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 - ) - 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 = jax.jit(jacobian(circuit))(a) - - assert isinstance(jac, tuple) - assert len(jac) == 2 # measurements - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == (2,) - - assert isinstance(jac[1], jax.numpy.ndarray) - assert jac[1].shape == (2,) - - def test_jacobian_var_var_multiple_params( - 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 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") - - 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 - ) - 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 = jax.jit(jacobian(circuit, argnums=[0, 1]))(par_0, par_1) - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == () - assert isinstance(jac[1][1], jax.numpy.ndarray) - assert jac[1][1].shape == () - - def test_jacobian_var_var_multiple_params_array( - 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 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") - - dev = qml.device(dev_name, wires=2, shots=shots) - - @qnode( - 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.var(qml.PauliZ(0) @ qml.PauliX(1)), qml.var(qml.PauliZ(0)) - - a = jax.numpy.array([0.1, 0.2]) - - jac = jax.jit(jacobian(circuit))(a) - - assert isinstance(jac, tuple) - assert len(jac) == 2 # measurements - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == (2,) - - assert isinstance(jac[1], jax.numpy.ndarray) - assert jac[1].shape == (2,) - - def test_jacobian_multiple_measurement_single_param( - 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") - 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 - ) - 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 = jax.jit(jacobian(circuit))(a) - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == () - - assert isinstance(jac[1], jax.numpy.ndarray) - assert jac[1].shape == (4,) - - def test_jacobian_multiple_measurement_multiple_param( - 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") - - 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 - ) - 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 = jax.jit(jacobian(circuit, argnums=[0, 1]))(a, b) - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == (4,) - assert isinstance(jac[1][1], jax.numpy.ndarray) - 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 - ): - """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) - - @qnode( - 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 = jax.numpy.array([0.1, 0.2]) - - jac = jax.jit(jacobian(circuit))(a) - - assert isinstance(jac, tuple) - assert len(jac) == 2 # measurements - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == (2,) - - assert isinstance(jac[1], jax.numpy.ndarray) - assert jac[1].shape == (4, 2) - - -hessian_fn = [ - jax.hessian, - lambda fn, argnums=0: jax.jacrev(jax.jacfwd(fn, argnums=argnums), argnums=argnums), - lambda fn, argnums=0: jax.jacfwd(jax.jacrev(fn, argnums=argnums), argnums=argnums), -] - - -@pytest.mark.parametrize("hessian", hessian_fn) -@pytest.mark.parametrize( - "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_name, diff_method, hessian, 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.") - - 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, - grad_on_execution=grad_on_execution, - ) - 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.jit(hessian(circuit, argnums=[0, 1]))(par_0, par_1) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], tuple) - assert len(hess[0]) == 2 - assert isinstance(hess[0][0], jax.numpy.ndarray) - assert isinstance(hess[0][1], jax.numpy.ndarray) - assert hess[0][0].shape == () - assert hess[0][1].shape == () - - assert isinstance(hess[1], tuple) - assert len(hess[1]) == 2 - assert isinstance(hess[1][0], jax.numpy.ndarray) - assert isinstance(hess[1][1], jax.numpy.ndarray) - assert hess[1][0].shape == () - assert hess[1][1].shape == () - - def test_hessian_expval_multiple_param_array( - 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 == "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( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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.jit(hessian(circuit))(params) - - assert isinstance(hess, jax.numpy.ndarray) - assert hess.shape == (2, 2) - - def test_hessian_var_multiple_params( - self, dev_name, diff_method, hessian, 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.") - elif diff_method == "hadamard": - pytest.skip("Test does not supports hadamard because of var.") - - par_0 = jax.numpy.array(0.1, dtype=jax.numpy.float64) - par_1 = jax.numpy.array(0.2, dtype=jax.numpy.float64) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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.jit(hessian(circuit, argnums=[0, 1]))(par_0, par_1) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], tuple) - assert len(hess[0]) == 2 - assert isinstance(hess[0][0], jax.numpy.ndarray) - assert isinstance(hess[0][1], jax.numpy.ndarray) - assert hess[0][0].shape == () - assert hess[0][1].shape == () - - assert isinstance(hess[1], tuple) - assert len(hess[1]) == 2 - assert isinstance(hess[1][0], jax.numpy.ndarray) - assert isinstance(hess[1][1], jax.numpy.ndarray) - assert hess[1][0].shape == () - assert hess[1][1].shape == () - - def test_hessian_var_multiple_param_array( - 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 == "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( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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.jit(hessian(circuit))(params) - - assert isinstance(hess, jax.numpy.ndarray) - assert hess.shape == (2, 2) - - def test_hessian_probs_expval_multiple_params( - self, dev_name, diff_method, hessian, grad_on_execution, 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": - 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.") - - par_0 = jax.numpy.array(0.1, dtype=jax.numpy.float64) - par_1 = jax.numpy.array(0.2, dtype=jax.numpy.float64) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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.jit(hessian(circuit, argnums=[0, 1]))(par_0, par_1) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], tuple) - assert len(hess[0]) == 2 - - for h in hess[0]: - assert isinstance(h, tuple) - for h_comp in h: - assert h_comp.shape == () - - for h in hess[1]: - assert isinstance(h, tuple) - for h_comp in h: - assert h_comp.shape == (2,) - - def test_hessian_probs_expval_multiple_param_array( - 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 == "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( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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.jit(hessian(circuit))(params) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - assert isinstance(hess[0], jax.numpy.ndarray) - assert hess[0].shape == (2, 2) - - assert isinstance(hess[1], jax.numpy.ndarray) - assert hess[1].shape == (2, 2, 2) - - def test_hessian_probs_var_multiple_params( - self, dev_name, diff_method, hessian, 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.") - elif diff_method == "hadamard": - pytest.skip("Test does not supports hadamard because of var.") - - par_0 = jax.numpy.array(0.1, dtype=jax.numpy.float64) - par_1 = jax.numpy.array(0.2, dtype=jax.numpy.float64) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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.jit(hessian(circuit, argnums=[0, 1]))(par_0, par_1) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], tuple) - assert len(hess[0]) == 2 - - for h in hess[0]: - assert isinstance(h, tuple) - for h_comp in h: - assert h_comp.shape == () - - for h in hess[1]: - assert isinstance(h, tuple) - for h_comp in h: - assert h_comp.shape == (2,) - - def test_hessian_probs_var_multiple_param_array( - 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 == "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( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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.jit(hessian(circuit))(params) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - assert isinstance(hess[0], jax.numpy.ndarray) - assert hess[0].shape == (2, 2) - - assert isinstance(hess[1], jax.numpy.ndarray) - assert hess[1].shape == (2, 2, 2) - - -@pytest.mark.parametrize("hessian", hessian_fn) -@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) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = jax.numpy.array([1.0, 2.0]) - a, b = x - - hess = jax.jit(hessian(circuit))(x) - - expected_hess = [ - [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)], - [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)], - ] - shots_tol = 0.1 - assert np.allclose(hess, expected_hess, atol=shots_tol, rtol=0) - - -@pytest.mark.parametrize("jit_inside", [True, False]) -@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 -) -class TestSubsetArgnums: - def test_single_measurement( - self, - interface, - dev_name, - diff_method, - grad_on_execution, - jacobian, - argnums, - jit_inside, - tol, - ): - """Test single measurement with different diff methods with argnums.""" - - dev = qml.device(dev_name, wires=3) - - kwargs = {} - if diff_method == "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, - cache=False, - **kwargs, - ) - 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)) - - a = jax.numpy.array(1.0) - b = jax.numpy.array(2.0) - - if jit_inside: - jac = jacobian(jax.jit(circuit), argnums=argnums)(a, b) - else: - jac = jax.jit(jacobian(circuit, argnums=argnums))(a, b) - - expected = np.array([-np.sin(a), 0]) - - if argnums == 0: - assert np.allclose(jac, expected[0], atol=tol) - elif argnums == 1: - assert np.allclose(jac, expected[1], atol=tol) - else: - assert np.allclose(jac[0], expected[0], atol=tol) - assert np.allclose(jac[1], expected[1], atol=tol) - - def test_multi_measurements( - self, - interface, - dev_name, - diff_method, - grad_on_execution, - jacobian, - argnums, - jit_inside, - tol, - ): - """Test multiple measurements with different diff methods with argnums.""" - dev = qml.device(dev_name, wires=3) - - kwargs = {} - if diff_method == "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, - **kwargs, - ) - 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)) - - a = jax.numpy.array(1.0) - b = jax.numpy.array(2.0) - - if jit_inside: - jac = jacobian(jax.jit(circuit), argnums=argnums)(a, b) - else: - jac = jax.jit(jacobian(circuit, argnums=argnums))(a, b) - - expected = np.array([[-np.sin(a), 0], [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)]]) - - if argnums == 0: - assert np.allclose(jac, expected.T[0], atol=tol) - elif argnums == 1: - assert np.allclose(jac, expected.T[1], atol=tol) - else: - assert np.allclose(jac[0], expected[0], atol=tol) - assert np.allclose(jac[1], expected[1], atol=tol) diff --git a/tests/interfaces/legacy_devices_integration/test_jax_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_legacy.py deleted file mode 100644 index ecaa78a4164..00000000000 --- a/tests/interfaces/legacy_devices_integration/test_jax_legacy.py +++ /dev/null @@ -1,876 +0,0 @@ -# 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_compute_derivatives") - - 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.target_device._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 == 1 - - # With caching, also 2 evaluations are required. One - # for the forward pass, and one for the backward pass. - dev.target_device._num_executions = 0 - jac_fn = jax.grad(cost) - jac_fn(params, cache=True) - assert dev.num_executions == 1 - - -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/legacy_devices_integration/test_jax_qnode_legacy.py b/tests/interfaces/legacy_devices_integration/test_jax_qnode_legacy.py deleted file mode 100644 index 4aaed88e12c..00000000000 --- a/tests/interfaces/legacy_devices_integration/test_jax_qnode_legacy.py +++ /dev/null @@ -1,2507 +0,0 @@ -# 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. -"""Integration tests for using the JAX-Python interface with a QNode""" -# 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.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_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 - -jax = pytest.importorskip("jax") -jax.config.update("jax_enable_x64", True) - -TOL_FOR_SPSA = 1.0 -SEED_FOR_SPSA = 32651 -H_FOR_SPSA = 0.05 - - -@pytest.mark.parametrize( - "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_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 - ) - 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) - circuit(a) - - assert circuit.interface == interface - - # 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) - assert grad.shape == () - - 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": - pytest.skip("Test only supports parameter-shift") - - 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, - ) - def circuit(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.expval(qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliY(1)])) - - grad_fn = jax.grad(circuit, argnums=[0, 1]) - res = grad_fn(a, b) - - # the tape has reported both arguments as trainable - assert circuit.qtape.trainable_params == [0, 1] - - expected = [-np.sin(a) + np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)] - assert np.allclose(res, expected, atol=tol, rtol=0) - - # make the second QNode argument a constant - grad_fn = jax.grad(circuit, argnums=0) - res = grad_fn(a, b) - - # the tape has reported only the first argument as trainable - assert circuit.qtape.trainable_params == [0] - - 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): - """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 - ) - def circuit(a, b, c): - qml.RY(a * c, wires=0) - qml.RZ(b, wires=0) - qml.RX(c + c**2 + jax.numpy.sin(a), wires=0) - return qml.expval(qml.PauliZ(0)) - - res = jax.grad(circuit, argnums=[0, 2])(a, b, c) - - if diff_method == "finite-diff": - assert circuit.qtape.trainable_params == [0, 2] - - assert len(res) == 2 - - 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 - ) - def circuit(U, a): - qml.QubitUnitary(U, wires=0) - qml.RY(a, wires=0) - return qml.expval(qml.PauliZ(0)) - - res = jax.grad(circuit, argnums=1)(U, a) - assert np.allclose(res, np.sin(a), atol=tol, rtol=0) - - 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): - """Test that operation and nested tape expansion - is differentiable""" - - 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} - tol = TOL_FOR_SPSA - - class U3(qml.U3): - def decomposition(self): - theta, phi, lam = self.data - wires = self.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]) - - @qnode(dev, **kwargs) - def circuit(a, p): - qml.RX(a, wires=0) - U3(p[0], p[1], p[2], wires=0) - return qml.expval(qml.PauliX(0)) - - res = 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 np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.grad(circuit, argnums=1)(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=tol, rtol=0) - - 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, - 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) - return qml.expval(qml.PauliZ(0)) - - jax.jacobian(circuit)(a) - - -@pytest.mark.parametrize( - "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_name, diff_method, grad_on_execution, interface, tol): - """Test jacobian calculation""" - 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 - - 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): - 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)) - - res = circuit(a, b) - - assert circuit.qtape.trainable_params == [0, 1] - assert isinstance(res, tuple) - assert len(res) == 2 - - expected = [np.cos(a), -np.cos(a) * np.sin(b)] - assert np.allclose(res[0], expected[0], atol=tol, rtol=0) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - 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 isinstance(res, tuple) - assert len(res) == 2 - - assert isinstance(res[0], tuple) - assert isinstance(res[0][0], jax.numpy.ndarray) - assert res[0][0].shape == () - assert np.allclose(res[0][0], expected[0][0], atol=tol, rtol=0) - assert isinstance(res[0][1], jax.numpy.ndarray) - assert res[0][1].shape == () - assert np.allclose(res[0][1], expected[0][1], atol=tol, rtol=0) - - assert isinstance(res[1], tuple) - assert isinstance(res[1][0], jax.numpy.ndarray) - assert res[1][0].shape == () - assert np.allclose(res[1][0], expected[1][0], atol=tol, rtol=0) - assert isinstance(res[1][1], jax.numpy.ndarray) - 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): - """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 - ) - if diff_method == "spsa": - kwargs["sampler_rng"] = np.random.default_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, **kwargs) - 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)) - - jac_fn = jax.jacobian(circuit, argnums=[0, 1]) - res = jac_fn(a, b) - - assert isinstance(res, tuple) - assert len(res) == 2 - - 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) - - # call the Jacobian with new parameters - a = jax.numpy.array(0.6) - b = jax.numpy.array(0.832) - - res = jac_fn(a, b) - - assert isinstance(res, tuple) - assert len(res) == 2 - - 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) - - 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 = 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 - - 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) - - @qnode(dev, **kwargs) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[1]) - - res = jax.jacobian(circuit, argnums=[0, 1])(x, y) - - expected = np.array( - [ - [-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2], - [np.cos(y) * np.sin(x) / 2, np.cos(x) * np.sin(y) / 2], - ] - ) - - assert isinstance(res, tuple) - assert len(res) == 2 - - assert isinstance(res[0], jax.numpy.ndarray) - assert res[0].shape == (2,) - - assert isinstance(res[1], jax.numpy.ndarray) - assert res[1].shape == (2,) - - 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): - """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["sampler_rng"] = np.random.default_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) - - @qnode(dev, **kwargs) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=[0]), qml.probs(wires=[1, 2]) - - res = circuit(x, y) - - assert isinstance(res, tuple) - assert len(res) == 2 - - expected = [ - [np.cos(x / 2) ** 2, np.sin(x / 2) ** 2], - [(1 + np.cos(x) * np.cos(y)) / 2, 0, (1 - np.cos(x) * np.cos(y)) / 2, 0], - ] - - assert isinstance(res[0], jax.numpy.ndarray) - 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,) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) - expected_0 = np.array( - [ - [-np.sin(x) / 2, np.sin(x) / 2], - [0, 0], - ] - ) - - expected_1 = np.array( - [ - [-np.cos(y) * np.sin(x) / 2, 0, np.sin(x) * np.cos(y) / 2, 0], - [-np.cos(x) * np.sin(y) / 2, 0, np.cos(x) * np.sin(y) / 2, 0], - ] - ) - - assert isinstance(jac, tuple) - assert isinstance(jac[0], tuple) - - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == (2,) - assert np.allclose(jac[0][0], expected_0[0], atol=tol, rtol=0) - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == (2,) - assert np.allclose(jac[0][1], expected_0[1], atol=tol, rtol=0) - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == (4,) - - assert np.allclose(jac[1][0], expected_1[0], atol=tol, rtol=0) - assert isinstance(jac[1][1], jax.numpy.ndarray) - 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): - """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["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) - - @qnode(dev, **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=[1]) - - res = circuit(x, y) - expected = [np.cos(x), [(1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2]] - assert isinstance(res, tuple) - assert len(res) == 2 - - assert isinstance(res[0], jax.numpy.ndarray) - 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,) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) - expected = [ - [-np.sin(x), 0], - [ - [-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2], - [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], - ], - ] - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert np.allclose(jac[0][0], expected[0][0], atol=tol, rtol=0) - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - assert np.allclose(jac[0][1], expected[0][1], atol=tol, rtol=0) - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == (2,) - assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - assert isinstance(jac[1][1], jax.numpy.ndarray) - assert jac[1][1].shape == (2,) - 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 - ): - """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 - 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) - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **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=[1]) - - jac = jax.jacobian(circuit, argnums=[0])(x, y) - - expected = [ - [-np.sin(x), 0], - [ - [-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2], - [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], - ], - ] - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 1 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert np.allclose(jac[0][0], expected[0][0], atol=tol, rtol=0) - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 1 - assert isinstance(jac[1][0], jax.numpy.ndarray) - 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): - """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": - pytest.skip("Hadamard does not support var") - 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) - - @qnode(dev, **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.probs(wires=[1]) - - res = circuit(x, y) - - expected = [ - np.sin(x) ** 2, - [(1 + np.cos(x) * np.cos(y)) / 2, (1 - np.cos(x) * np.cos(y)) / 2], - ] - - assert isinstance(res[0], jax.numpy.ndarray) - 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,) - assert np.allclose(res[1], expected[1], atol=tol, rtol=0) - - jac = jax.jacobian(circuit, argnums=[0, 1])(x, y) - expected = [ - [2 * np.cos(x) * np.sin(x), 0], - [ - [-np.sin(x) * np.cos(y) / 2, np.cos(y) * np.sin(x) / 2], - [-np.cos(x) * np.sin(y) / 2, np.cos(x) * np.sin(y) / 2], - ], - ] - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert np.allclose(jac[0][0], expected[0][0], atol=tol, rtol=0) - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - assert np.allclose(jac[0][1], expected[0][1], atol=tol, rtol=0) - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == (2,) - assert np.allclose(jac[1][0], expected[1][0], atol=tol, rtol=0) - assert isinstance(jac[1][1], jax.numpy.ndarray) - assert jac[1][1].shape == (2,) - assert np.allclose(jac[1][1], expected[1][1], atol=tol, rtol=0) - - -@pytest.mark.parametrize("interface", ["auto", "jax", "jax-python"]) -class TestShotsIntegration: - """Test that the QNode correctly changes shot value, and - 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) - - @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)), 1) - - 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) - 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)) - - # execute with device default shots (None) - with pytest.raises(qml.QuantumFunctionError): - circuit(a, b) - - # execute with shots=100 - 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) - 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)) - - res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000) - assert dev.shots == qml.measurements.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): - """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") - 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)) - - # 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 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_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_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.") - - 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) - def circuit(): - qml.Hadamard(wires=[0]) - qml.CNOT(wires=[0, 1]) - return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliX(1)) - - res = circuit() - - assert isinstance(res, tuple) - - assert isinstance(res[0], jax.Array) - assert res[0].shape == (10,) - assert isinstance(res[1], jax.Array) - assert res[1].shape == (10,) - - 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 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) - 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), - ) - - res = circuit() - - assert isinstance(res, tuple) - - assert isinstance(res[0], dict) - assert len(res[0]) == 2 - assert isinstance(res[1], dict) - assert len(res[1]) == 2 - - def test_chained_qnodes(self, 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) - - 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) - 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) - def circuit2(data, weights): - qml.templates.AngleEmbedding(jax.numpy.stack([data, 0.7]), wires=[0, 1]) - Template(weights, wires=[0, 1]) - return qml.expval(qml.PauliX(0)) - - def cost(weights): - w1, w2 = weights - c1 = circuit1(w1) - c2 = circuit2(c1, w2) - return jax.numpy.sum(c2) ** 2 - - w1 = qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=3) - w2 = qml.templates.StronglyEntanglingLayers.shape(n_wires=2, n_layers=4) - - weights = [ - jax.numpy.array(np.random.random(w1)), - jax.numpy.array(np.random.random(w2)), - ] - - grad_fn = jax.grad(cost) - res = grad_fn(weights) - - assert len(res) == 2 - - -@pytest.mark.parametrize( - "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_name, diff_method, grad_on_execution, 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, - ) - 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} - 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) - qml.RX(x[1], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = jax.numpy.array([1.0, 2.0]) - res = circuit(x) - g = jax.grad(circuit)(x) - g2 = jax.grad(lambda x: jax.numpy.sum(jax.grad(circuit)(x)))(x) - - a, b = x - - expected_res = np.cos(a) * np.cos(b) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - - expected_g2 = [ - -np.cos(a) * np.cos(b) + np.sin(a) * np.sin(b), - np.sin(a) * np.sin(b) - np.cos(a) * np.cos(b), - ] - if diff_method == "finite-diff": - assert np.allclose(g2, expected_g2, atol=10e-2, rtol=0) - else: - assert np.allclose(g2, expected_g2, atol=tol, rtol=0) - - 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": - pytest.skip("Adjoint does not support second derivative.") - elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "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, - max_diff=2, - **gradient_kwargs, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.expval(qml.PauliZ(0)) - - x = jax.numpy.array([1.0, 2.0]) - res = circuit(x) - - a, b = x - - expected_res = np.cos(a) * np.cos(b) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - grad_fn = jax.grad(circuit) - g = grad_fn(x) - - expected_g = [-np.sin(a) * np.cos(b), -np.cos(a) * np.sin(b)] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - - hess = jax.jacobian(grad_fn)(x) - - expected_hess = [ - [-np.cos(a) * np.cos(b), np.sin(a) * np.sin(b)], - [np.sin(a) * np.sin(b), -np.cos(a) * np.cos(b)], - ] - if diff_method == "finite-diff": - assert np.allclose(hess, expected_hess, atol=10e-2, rtol=0) - 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): - """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": - gradient_kwargs = { - "h": H_FOR_SPSA, - "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, - max_diff=2, - **gradient_kwargs, - ) - def circuit(x): - qml.RY(x[0], wires=0) - qml.RX(x[1], wires=0) - return qml.probs(wires=0) - - x = jax.numpy.array([1.0, 2.0]) - res = circuit(x) - - a, b = x - - expected_res = [0.5 + 0.5 * np.cos(a) * np.cos(b), 0.5 - 0.5 * np.cos(a) * np.cos(b)] - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - jac_fn = jax.jacobian(circuit) - g = jac_fn(x) - - expected_g = [ - [-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)], - ] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - - hess = jax.jacobian(jac_fn)(x) - - expected_hess = [ - [ - [-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)], - ], - ] - if diff_method == "finite-diff": - assert np.allclose(hess, expected_hess, atol=10e-2, rtol=0) - else: - 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 - ): - """Test hessian calculation of a vector valued QNode with post-processing""" - gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support second derivative.") - elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "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, - max_diff=2, - **gradient_kwargs, - ) - def circuit(x): - qml.RX(x[0], wires=0) - qml.RY(x[1], wires=0) - return qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(0)) - - def cost_fn(x): - return x @ jax.numpy.array(circuit(x)) - - x = jax.numpy.array([0.76, -0.87]) - res = cost_fn(x) - - a, b = x - - expected_res = x @ jax.numpy.array([np.cos(a) * np.cos(b), np.cos(a) * np.cos(b)]) - assert np.allclose(res, expected_res, atol=tol, rtol=0) - - grad_fn = jax.grad(cost_fn) - g = grad_fn(x) - - expected_g = [ - np.cos(b) * (np.cos(a) - (a + b) * np.sin(a)), - np.cos(a) * (np.cos(b) - (a + b) * np.sin(b)), - ] - assert np.allclose(g, expected_g, atol=tol, rtol=0) - hess = jax.jacobian(grad_fn)(x) - - expected_hess = [ - [ - -(np.cos(b) * ((a + b) * np.cos(a) + 2 * np.sin(a))), - -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b), - ], - [ - -(np.cos(b) * np.sin(a)) + (-np.cos(a) + (a + b) * np.sin(a)) * np.sin(b), - -(np.cos(a) * ((a + b) * np.cos(b) + 2 * np.sin(b))), - ], - ] - - if diff_method == "finite-diff": - assert np.allclose(hess, expected_hess, atol=10e-2, rtol=0) - else: - 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 - ): - """Test hessian calculation of a vector valued QNode that has separate input arguments""" - gradient_kwargs = {} - if diff_method == "adjoint": - pytest.skip("Adjoint does not support second derivative.") - elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "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, - max_diff=2, - **gradient_kwargs, - ) - def circuit(a, b): - qml.RY(a, wires=0) - qml.RX(b, wires=0) - return qml.probs(wires=0) - - a = jax.numpy.array(1.0) - b = jax.numpy.array(2.0) - res = 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 np.allclose(res, expected_res, atol=tol, rtol=0) - - jac_fn = jax.jacobian(circuit, argnums=[0, 1]) - g = jac_fn(a, b) - - expected_g = np.array( - [ - [-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)], - ] - ) - 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( - [ - [ - [-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.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)], - ], - ] - ) - if diff_method == "finite-diff": - assert np.allclose(hess, expected_hess, atol=10e-2, rtol=0) - else: - assert np.allclose(hess, expected_hess, atol=tol, rtol=0) - - 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 - - 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 - ) - def circuit(x, y): - qml.RX(x, wires=[0]) - qml.RY(y, wires=[1]) - qml.CNOT(wires=[0, 1]) - return qml.state() - - def cost_fn(x, y): - res = circuit(x, y) - assert res.dtype is np.dtype("complex128") - probs = jax.numpy.abs(res) ** 2 - return probs[0] + probs[2] - - res = cost_fn(x, y) - - if diff_method not in {"backprop"}: - pytest.skip("Test only supports backprop") - - res = jax.grad(cost_fn, argnums=[0, 1])(x, y) - expected = np.array([-np.sin(x) * np.cos(y) / 2, -np.cos(x) * np.sin(y) / 2]) - 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): - """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("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 - - @qnode( - dev, - diff_method=diff_method, - interface=interface, - grad_on_execution=grad_on_execution, - **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.Projector(P, wires=0) @ qml.PauliX(1)) - - 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 np.allclose(res, expected, atol=tol, rtol=0) - - res = jax.grad(circuit, argnums=[0, 1])(x, y) - expected = np.array( - [ - 0.5 * np.sin(x) * (np.cos(x / 2) ** 2 + np.cos(2 * y) * np.sin(x / 2) ** 2), - -2 * np.cos(y) * np.sin(x / 2) ** 4 * np.sin(y), - ] - ) - assert np.allclose(res, expected, atol=tol, rtol=0) - - -@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.target_device, "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 - with the JAX interface""" - - @pytest.mark.parametrize("max_diff", [1, 2]) - def test_gradient_expansion_trainable_only( - 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 - - 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, - interface=interface, - ) - def circuit(x, y): - qml.Hadamard(wires=0) - PhaseShift(x, wires=0) - PhaseShift(2 * y, wires=0) - return qml.expval(qml.PauliX(0)) - - 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 - ): - """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 diff_method == "adjoint": - pytest.skip("The adjoint method does not yet support Hamiltonians") - elif diff_method == "hadamard": - pytest.skip("Hadamard does not yet support Hamiltonians.") - elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "num_directions": 20, - "sampler_rng": np.random.default_rng(SEED_FOR_SPSA), - } - 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)] - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - max_diff=max_diff, - **gradient_kwargs, - ) - def circuit(data, weights, coeffs): - weights = weights.reshape(1, -1) - qml.templates.AngleEmbedding(data, wires=[0, 1]) - qml.templates.BasicEntanglerLayers(weights, wires=[0, 1]) - return qml.expval(qml.Hamiltonian(coeffs, obs)) - - d = jax.numpy.array([0.1, 0.2]) - w = jax.numpy.array([0.654, -0.734]) - c = jax.numpy.array([-0.6543, 0.24, 0.54]) - - # test output - 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() - - # test gradients - 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]), - ] - expected_c = [0, -np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]), np.cos(d[1] + w[1])] - assert np.allclose(grad[0], expected_w, atol=tol) - assert np.allclose(grad[1], expected_c, atol=tol) - - # TODO: Add parameter shift when the bug with trainable params and hamiltonian_grad is solved. - # test second-order derivatives - if diff_method in "backprop" and max_diff == 2: - grad2_c = jax.jacobian(jax.grad(circuit, argnums=[2]), argnums=[2])(d, w, c) - assert np.allclose(grad2_c, 0, atol=tol) - - grad2_w_c = jax.jacobian(jax.grad(circuit, argnums=[1]), argnums=[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]), - -np.sin(d[1] + w[1]), - ] - 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 - ): - """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 diff_method in ("adjoint", "backprop", "finite-diff"): - pytest.skip("The adjoint and backprop methods do not yet support sampling") - elif diff_method == "hadamard": - pytest.skip("Hadamard does not yet support Hamiltonians.") - elif diff_method == "spsa": - gradient_kwargs = { - "h": H_FOR_SPSA, - "sampler_rng": np.random.default_rng(SEED_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( - dev, - interface=interface, - diff_method=diff_method, - grad_on_execution=grad_on_execution, - max_diff=max_diff, - **gradient_kwargs, - ) - def circuit(data, weights, coeffs): - weights = weights.reshape(1, -1) - qml.templates.AngleEmbedding(data, wires=[0, 1]) - qml.templates.BasicEntanglerLayers(weights, wires=[0, 1]) - H = qml.Hamiltonian(coeffs, obs) - H.compute_grouping() - return qml.expval(H) - - d = jax.numpy.array([0.1, 0.2]) - w = jax.numpy.array([0.654, -0.734]) - c = jax.numpy.array([-0.6543, 0.24, 0.54]) - - # test output - 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) - 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]), - ] - expected_c = [0, -np.sin(d[0] + w[0]) * np.sin(d[1] + w[1]), np.cos(d[1] + w[1])] - assert np.allclose(grad[0], expected_w, atol=tol) - assert np.allclose(grad[1], expected_c, atol=tol) - - # TODO: Fix hamiltonian grad for parameter shift and jax - # # test second-order derivatives - # if diff_method == "parameter-shift" and max_diff == 2: - - # grad2_c = jax.jacobian(jax.grad(circuit, argnum=2), argnum=2)(d, w, c) - # assert np.allclose(grad2_c, 0, atol=tol) - - # grad2_w_c = jax.jacobian(jax.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]), - # -np.sin(d[1] + w[1]), - # ] - # assert np.allclose(grad2_w_c, expected, atol=tol) - - -jacobian_fn = [jax.jacobian, jax.jacrev, jax.jacfwd] - - -@pytest.mark.parametrize("shots", [None, 10000]) -@pytest.mark.parametrize( - "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_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 - ) - 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) - - grad = jax.grad(circuit)(a) - - 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 - ): - """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 - ) - 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) - - grad = jax.grad(circuit, argnums=[0, 1])(a, b) - - assert isinstance(grad, tuple) - assert len(grad) == 2 - assert grad[0].shape == () - assert grad[1].shape == () - - def test_grad_single_measurement_multiple_param_array( - 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 - ) - 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]) - - 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_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 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) - - @qnode( - 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 = jax.numpy.array(0.1) - - 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_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 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) - - @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=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) - - assert isinstance(jac, tuple) - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == (4,) - - assert isinstance(jac[1], jax.numpy.ndarray) - assert jac[1].shape == (4,) - - @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 - ): - """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 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 - ) - 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) - - 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 - ): - """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 = jax.numpy.array(0.1) - par_1 = jax.numpy.array(0.2) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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) - - assert isinstance(jac, tuple) - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == () - assert isinstance(jac[1][1], jax.numpy.ndarray) - assert jac[1][1].shape == () - - @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 - ): - """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=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.PauliX(1)), qml.expval(qml.PauliZ(0)) - - a = jax.numpy.array([0.1, 0.2]) - - jac = jacobian(circuit)(a) - - assert isinstance(jac, tuple) - assert len(jac) == 2 # measurements - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == (2,) - - assert isinstance(jac[1], jax.numpy.ndarray) - assert jac[1].shape == (2,) - - @pytest.mark.parametrize("jacobian", jacobian_fn) - def test_jacobian_var_var_multiple_params( - 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("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) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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) - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == () - assert isinstance(jac[1][1], jax.numpy.ndarray) - assert jac[1][1].shape == () - - @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 - ): - """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("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 - ) - 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) - - assert isinstance(jac, tuple) - assert len(jac) == 2 # measurements - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == (2,) - - assert isinstance(jac[1], jax.numpy.ndarray) - assert jac[1].shape == (2,) - - @pytest.mark.parametrize("jacobian", jacobian_fn) - def test_jacobian_multiple_measurement_single_param( - 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") - - 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 - ) - 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) - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == () - - assert isinstance(jac[1], jax.numpy.ndarray) - assert jac[1].shape == (4,) - - @pytest.mark.parametrize("jacobian", jacobian_fn) - def test_jacobian_multiple_measurement_multiple_param( - 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") - - 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 - ) - 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) - - assert isinstance(jac, tuple) - assert len(jac) == 2 - - assert isinstance(jac[0], tuple) - assert len(jac[0]) == 2 - assert isinstance(jac[0][0], jax.numpy.ndarray) - assert jac[0][0].shape == () - assert isinstance(jac[0][1], jax.numpy.ndarray) - assert jac[0][1].shape == () - - assert isinstance(jac[1], tuple) - assert len(jac[1]) == 2 - assert isinstance(jac[1][0], jax.numpy.ndarray) - assert jac[1][0].shape == (4,) - assert isinstance(jac[1][1], jax.numpy.ndarray) - assert jac[1][1].shape == (4,) - - @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 - ): - """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) - - @qnode( - 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 = jax.numpy.array([0.1, 0.2]) - - jac = jacobian(circuit)(a) - - assert isinstance(jac, tuple) - assert len(jac) == 2 # measurements - - assert isinstance(jac[0], jax.numpy.ndarray) - assert jac[0].shape == (2,) - - assert isinstance(jac[1], jax.numpy.ndarray) - assert jac[1].shape == (4, 2) - - def test_hessian_expval_multiple_params( - 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.") - - 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, - grad_on_execution=grad_on_execution, - ) - 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) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], tuple) - assert len(hess[0]) == 2 - assert isinstance(hess[0][0], jax.numpy.ndarray) - assert hess[0][0].shape == () - assert hess[0][1].shape == () - - assert isinstance(hess[1], tuple) - assert len(hess[1]) == 2 - assert isinstance(hess[1][0], jax.numpy.ndarray) - assert hess[1][0].shape == () - assert hess[1][1].shape == () - - def test_hessian_expval_multiple_param_array( - 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": - pytest.skip("Test does not supports adjoint because second order diff.") - 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, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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) - - 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 - ): - """The hessian of single a measurement with multiple params return a tuple of arrays.""" - if diff_method == "adjoint": - pytest.skip("Test does not supports adjoint because second order diff.") - 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) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], tuple) - assert len(hess[0]) == 2 - assert isinstance(hess[0][0], jax.numpy.ndarray) - assert hess[0][0].shape == () - assert hess[0][1].shape == () - - assert isinstance(hess[1], tuple) - assert len(hess[1]) == 2 - assert isinstance(hess[1][0], jax.numpy.ndarray) - assert hess[1][0].shape == () - assert hess[1][1].shape == () - - def test_hessian_var_multiple_param_array( - 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": - pytest.skip("Test does not supports adjoint because second order diff.") - 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) - - params = jax.numpy.array([0.1, 0.2]) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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) - - 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 - ): - """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": - pytest.skip("Test does not supports diff of non commuting obs.") - - if shots is not None and diff_method in ("backprop", "adjoint"): - pytest.skip("Test does not support finite shots and adjoint/backprop") - - 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, - grad_on_execution=grad_on_execution, - ) - 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) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], tuple) - assert len(hess[0]) == 2 - assert isinstance(hess[0][0], tuple) - assert len(hess[0][0]) == 2 - assert isinstance(hess[0][0][0], jax.numpy.ndarray) - assert hess[0][0][0].shape == () - assert isinstance(hess[0][0][1], jax.numpy.ndarray) - assert hess[0][0][1].shape == () - assert isinstance(hess[0][1], tuple) - assert len(hess[0][1]) == 2 - assert isinstance(hess[0][1][0], jax.numpy.ndarray) - assert hess[0][1][0].shape == () - assert isinstance(hess[0][1][1], jax.numpy.ndarray) - assert hess[0][1][1].shape == () - - assert isinstance(hess[1], tuple) - assert len(hess[1]) == 2 - assert isinstance(hess[1][0], tuple) - assert len(hess[1][0]) == 2 - assert isinstance(hess[1][0][0], jax.numpy.ndarray) - assert hess[1][0][0].shape == (2,) - assert isinstance(hess[1][0][1], jax.numpy.ndarray) - assert hess[1][0][1].shape == (2,) - assert isinstance(hess[1][1], tuple) - assert len(hess[1][1]) == 2 - assert isinstance(hess[1][1][0], jax.numpy.ndarray) - assert hess[1][1][0].shape == (2,) - assert isinstance(hess[1][1][1], jax.numpy.ndarray) - 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 - ): - """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.") - elif diff_method == "hadamard": - pytest.skip("Test does not supports diff of non commuting obs.") - 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, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], jax.numpy.ndarray) - assert hess[0].shape == (2, 2) - - assert isinstance(hess[1], jax.numpy.ndarray) - assert hess[1].shape == (2, 2, 2) - - def test_hessian_probs_var_multiple_params( - 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": - pytest.skip("Test does not supports adjoint because second order diff.") - 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 = qml.numpy.array(0.1) - par_1 = qml.numpy.array(0.2) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], tuple) - assert len(hess[0]) == 2 - assert isinstance(hess[0][0], tuple) - assert len(hess[0][0]) == 2 - assert isinstance(hess[0][0][0], jax.numpy.ndarray) - assert hess[0][0][0].shape == () - assert isinstance(hess[0][0][1], jax.numpy.ndarray) - assert hess[0][0][1].shape == () - assert isinstance(hess[0][1], tuple) - assert len(hess[0][1]) == 2 - assert isinstance(hess[0][1][0], jax.numpy.ndarray) - assert hess[0][1][0].shape == () - assert isinstance(hess[0][1][1], jax.numpy.ndarray) - assert hess[0][1][1].shape == () - - assert isinstance(hess[1], tuple) - assert len(hess[1]) == 2 - assert isinstance(hess[1][0], tuple) - assert len(hess[1][0]) == 2 - assert isinstance(hess[1][0][0], jax.numpy.ndarray) - assert hess[1][0][0].shape == (2,) - assert isinstance(hess[1][0][1], jax.numpy.ndarray) - assert hess[1][0][1].shape == (2,) - assert isinstance(hess[1][1], tuple) - assert len(hess[1][1]) == 2 - assert isinstance(hess[1][1][0], jax.numpy.ndarray) - assert hess[1][1][0].shape == (2,) - assert isinstance(hess[1][1][1], jax.numpy.ndarray) - 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 - ): - """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.") - 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) - - params = jax.numpy.array([0.1, 0.2]) - - @qnode( - dev, - interface=interface, - diff_method=diff_method, - max_diff=2, - grad_on_execution=grad_on_execution, - ) - 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) - - assert isinstance(hess, tuple) - assert len(hess) == 2 - - assert isinstance(hess[0], jax.numpy.ndarray) - assert hess[0].shape == (2, 2) - - assert isinstance(hess[1], jax.numpy.ndarray) - assert hess[1].shape == (2, 2, 2) - - -@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 = qml.device(dev_name, wires=1) - - @qml.qnode(dev, interface="jax") - def circuit(): - qml.Hadamard(wires=0) - return qml.state() - - res = circuit() - assert isinstance(res, jax.numpy.ndarray) 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 deleted file mode 100644 index d87db9f89da..00000000000 --- a/tests/interfaces/legacy_devices_integration/test_jax_qnode_shot_vector_legacy.py +++ /dev/null @@ -1,932 +0,0 @@ -# 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""" -from contextlib import nullcontext - -# 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.raises( - qml.QuantumFunctionError, - match="adjoint with requested circuit.", - ) - if isinstance(shots, tuple) - else nullcontext() - ): - - @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) - - 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.raises( - qml.QuantumFunctionError, - match="adjoint with requested circuit.", - ): - - @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) - - 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 index 6e9739a631f..619bd74a066 100644 --- a/tests/interfaces/legacy_devices_integration/test_set_shots_legacy.py +++ b/tests/interfaces/legacy_devices_integration/test_set_shots_legacy.py @@ -14,7 +14,7 @@ """ Tests for workflow.set_shots """ - +import pytest import pennylane as qml from pennylane.measurements import Shots @@ -24,16 +24,18 @@ 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 + dev = qml.devices.DefaultMixed(wires=1) + with pytest.warns(qml.PennyLaneDeprecationWarning): + 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) + with pytest.warns(qml.PennyLaneDeprecationWarning): + with set_shots(dev, shot_tuples): + assert dev.shots == 20 + assert dev.shot_vector == list(shot_tuples.shot_vector) assert dev.shots is None @@ -42,6 +44,7 @@ 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 + dev = qml.devices.DefaultMixed(wires=1) + with pytest.warns(qml.PennyLaneDeprecationWarning): + with set_shots(dev, False): + assert dev.shots is None diff --git a/tests/interfaces/test_autograd_qnode.py b/tests/interfaces/test_autograd_qnode.py index 3db765be87d..129ab56dfe8 100644 --- a/tests/interfaces/test_autograd_qnode.py +++ b/tests/interfaces/test_autograd_qnode.py @@ -539,17 +539,26 @@ def cost_fn(a, b): qml.CNOT(wires=[0, 1]) return qml.expval(qml.PauliY(1)) - assert cost_fn.gradient_fn == "backprop" # gets restored to default + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + 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 spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift - assert cost_fn.gradient_fn is qml.gradients.param_shift + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + 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" + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + assert cost_fn.gradient_fn == "backprop" assert spy.call_args[1]["gradient_fn"] == "backprop" diff --git a/tests/interfaces/test_execute.py b/tests/interfaces/test_execute.py index 18003f95586..1732f741d8a 100644 --- a/tests/interfaces/test_execute.py +++ b/tests/interfaces/test_execute.py @@ -13,6 +13,7 @@ # limitations under the License. """Tests for exeuction with default qubit 2 independent of any interface.""" +import numpy as np import pytest import pennylane as qml @@ -42,3 +43,15 @@ def test_caching(gradient_fn): assert tracker.totals["batches"] == 1 assert tracker.totals["executions"] == 1 assert cache[qs.hash] == -1.0 + + +def test_execute_legacy_device(): + """Test that qml.execute works when passed a legacy device class.""" + + dev = qml.devices.DefaultMixed(wires=2) + + tape = qml.tape.QuantumScript([qml.RX(0.1, 0)], [qml.expval(qml.Z(0))]) + + res = qml.execute((tape,), dev) + + assert qml.math.allclose(res[0], np.cos(0.1)) diff --git a/tests/interfaces/test_jacobian_products.py b/tests/interfaces/test_jacobian_products.py index 2fc275558e7..3d758c38642 100644 --- a/tests/interfaces/test_jacobian_products.py +++ b/tests/interfaces/test_jacobian_products.py @@ -129,21 +129,6 @@ def test_device_jacobians_initialization_new_dev(self): assert isinstance(jpc._jacs_cache, LRUCache) assert len(jpc._jacs_cache) == 0 - def test_device_jacobians_initialization_old_dev(self): - """Test the private attributes are set during initialization of a DeviceDerivatives class with the - old device interface.""" - - device = qml.devices.DefaultQubitLegacy(wires=5) - - jpc = DeviceDerivatives(device, aj_config) - - assert jpc._device is device - assert jpc._execution_config == aj_config - assert isinstance(jpc._results_cache, LRUCache) - assert len(jpc._results_cache) == 0 - assert isinstance(jpc._jacs_cache, LRUCache) - assert len(jpc._jacs_cache) == 0 - def test_device_jacobians_repr(self): """Test the repr method for device jacobians.""" device = qml.device("default.qubit") diff --git a/tests/interfaces/test_jax_jit_qnode.py b/tests/interfaces/test_jax_jit_qnode.py index 4d959de446f..cce76a83b9e 100644 --- a/tests/interfaces/test_jax_jit_qnode.py +++ b/tests/interfaces/test_jax_jit_qnode.py @@ -847,11 +847,17 @@ def cost_fn(a, b): 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 + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + 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" + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + assert cost_fn.gradient_fn == "backprop" # if we set the shots to None, backprop can now be used assert spy.call_args[1]["gradient_fn"] == "backprop" diff --git a/tests/interfaces/test_jax_qnode.py b/tests/interfaces/test_jax_qnode.py index 3ea650f96e8..d24dec3383d 100644 --- a/tests/interfaces/test_jax_qnode.py +++ b/tests/interfaces/test_jax_qnode.py @@ -779,12 +779,18 @@ def cost_fn(a, b): 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 + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + assert cost_fn.gradient_fn == 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" + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + assert cost_fn.gradient_fn == "backprop" assert spy.call_args[1]["gradient_fn"] == "backprop" diff --git a/tests/interfaces/test_tensorflow_qnode.py b/tests/interfaces/test_tensorflow_qnode.py index e3c4597ae06..c09f1632202 100644 --- a/tests/interfaces/test_tensorflow_qnode.py +++ b/tests/interfaces/test_tensorflow_qnode.py @@ -540,12 +540,18 @@ def circuit(weights): 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 + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + assert circuit.gradient_fn == 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 circuit(weights) - assert circuit.gradient_fn == "backprop" + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + assert circuit.gradient_fn == "backprop" assert spy.call_args[1]["gradient_fn"] == "backprop" diff --git a/tests/interfaces/test_torch_qnode.py b/tests/interfaces/test_torch_qnode.py index 35854f13694..82dbda669d4 100644 --- a/tests/interfaces/test_torch_qnode.py +++ b/tests/interfaces/test_torch_qnode.py @@ -642,7 +642,10 @@ def cost_fn(a, b): 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 + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + assert cost_fn.gradient_fn == 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 diff --git a/tests/logging/test_logging_autograd.py b/tests/logging/test_logging_autograd.py index 01d637d025d..8ae7a9baaa2 100644 --- a/tests/logging/test_logging_autograd.py +++ b/tests/logging/test_logging_autograd.py @@ -69,7 +69,7 @@ def circuit(params): return qml.expval(qml.PauliZ(0)) circuit(params) - assert len(caplog.records) == 8 + assert len(caplog.records) == 9 log_records_expected = [ ( "pennylane.workflow.qnode", @@ -79,6 +79,10 @@ def circuit(params): "pennylane.workflow.qnode", ["Calling " ) - @pytest.mark.autograd - def test_diff_method_none(self, tol): - """Test that diff_method=None creates a QNode with no interface, and no - device swapping.""" - dev = qml.device("default.qubit.legacy", wires=1) - - @qnode(dev, diff_method=None) - def circuit(x): - qml.RX(x, wires=0) - return qml.expval(qml.PauliZ(0)) - - assert circuit.interface is None - assert circuit.gradient_fn is None - assert circuit.device is dev - # QNode can still be executed - assert np.allclose(circuit(0.5), np.cos(0.5), atol=tol, rtol=0) + assert np.allclose(qn(0.5), np.cos(0.5), rtol=0) with pytest.warns(UserWarning, match="Attempted to differentiate a function with no"): - grad = qml.grad(circuit)(0.5) + grad = qml.grad(qn)(0.5) assert np.allclose(grad, 0) @@ -545,29 +440,6 @@ def func(x, y): assert np.allclose(res, res2, atol=tol, rtol=0) assert qn.qtape is not old_tape - def test_jacobian(self): - """Test the jacobian computation""" - dev = qml.device("default.qubit.legacy", wires=2) - - def func(x, y): - qml.RX(x, wires=0) - qml.RY(y, wires=1) - qml.CNOT(wires=[0, 1]) - return qml.probs(wires=0), qml.probs(wires=1) - - qn = QNode( - func, dev, interface="autograd", diff_method="finite-diff", h=1e-8, approx_order=2 - ) - assert qn.gradient_kwargs["h"] == 1e-8 - assert qn.gradient_kwargs["approx_order"] == 2 - - jac = qn.gradient_fn(qn)( - pnp.array(0.45, requires_grad=True), pnp.array(0.1, requires_grad=True) - ) - assert isinstance(jac, tuple) and len(jac) == 2 - assert len(jac[0]) == 2 - assert len(jac[1]) == 2 - def test_returning_non_measurements(self): """Test that an exception is raised if a non-measurement is returned from the QNode.""" @@ -941,7 +813,9 @@ def test_no_defer_measurements_if_supported(self, mocker): """Test that the defer_measurements transform is not used during QNode construction if the device supports mid-circuit measurements.""" dev = qml.device("default.qubit.legacy", wires=3) - mocker.patch.object(qml.Device, "_capabilities", {"supports_mid_measure": True}) + mocker.patch.object( + qml.devices.LegacyDevice, "_capabilities", {"supports_mid_measure": True} + ) spy = mocker.spy(qml.defer_measurements, "_transform") @qml.qnode(dev) diff --git a/tests/test_return_types_qnode.py b/tests/test_return_types_qnode.py index f1dfa2cabdd..364eb468922 100644 --- a/tests/test_return_types_qnode.py +++ b/tests/test_return_types_qnode.py @@ -762,7 +762,7 @@ def circuit(x): assert sum(res.values()) == shots -devices = ["default.qubit.jax", "default.mixed"] +devices = ["default.mixed"] @pytest.mark.jax @@ -775,7 +775,7 @@ def test_state_default(self, wires): import jax - dev = qml.device("default.qubit.jax", wires=wires) + dev = qml.device("default.qubit", wires=wires) def circuit(x): qml.Hadamard(wires=[0]) @@ -1893,7 +1893,7 @@ def circuit(x): assert t.shape == () -devices = ["default.qubit.jax", "default.mixed"] +devices = ["default.mixed"] @pytest.mark.jax @@ -2204,7 +2204,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2232,7 +2234,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2268,7 +2272,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2295,7 +2301,9 @@ def circuit(x): all_shot_copies = [ shot_tuple.shots for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) for _ in range(shot_tuple.copies) ] @@ -2326,7 +2334,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2358,7 +2368,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2395,7 +2407,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2427,7 +2441,9 @@ def circuit(x): all_shot_copies = [ shot_tuple.shots for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) for _ in range(shot_tuple.copies) ] @@ -2455,7 +2471,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2552,7 +2570,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2584,7 +2604,9 @@ def test_scalar_sample_with_obs(self, shot_vector, meas1, meas2, device): raw_shot_vector = [ shot_tuple.shots for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) for _ in range(shot_tuple.copies) ] @@ -2601,7 +2623,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2642,7 +2666,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2656,7 +2682,9 @@ def circuit(x): for m in measurement_res ) - for shot_tuple in dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector: + for shot_tuple in ( + dev.shot_vector if isinstance(dev, qml.devices.LegacyDevice) else dev.shots.shot_vector + ): for idx in range(shot_tuple.copies): for i, r in enumerate(res[idx]): if i % 2 == 0 or shot_tuple.shots == 1: @@ -2674,7 +2702,9 @@ def test_scalar_counts_with_obs(self, shot_vector, meas1, meas2, device): raw_shot_vector = [ shot_tuple.shots for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) for _ in range(shot_tuple.copies) ] @@ -2691,7 +2721,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2726,7 +2758,9 @@ def test_scalar_counts_no_obs(self, shot_vector, meas1, meas2, device): raw_shot_vector = [ shot_tuple.shots for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) for _ in range(shot_tuple.copies) ] @@ -2743,7 +2777,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2774,7 +2810,9 @@ def test_probs_sample(self, shot_vector, sample_obs, device): raw_shot_vector = [ shot_tuple.shots for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) for _ in range(shot_tuple.copies) ] @@ -2799,7 +2837,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2835,7 +2875,9 @@ def test_probs_counts(self, shot_vector, sample_obs, device): raw_shot_vector = [ shot_tuple.shots for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) for _ in range(shot_tuple.copies) ] @@ -2860,7 +2902,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2896,7 +2940,9 @@ def test_sample_counts(self, shot_vector, sample_wires, counts_wires, device): raw_shot_vector = [ shot_tuple.shots for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) for _ in range(shot_tuple.copies) ] @@ -2927,7 +2973,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) @@ -2960,7 +3008,9 @@ def test_scalar_probs_sample_counts(self, shot_vector, meas1, meas2, device): raw_shot_vector = [ shot_tuple.shots for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) for _ in range(shot_tuple.copies) ] @@ -2983,7 +3033,9 @@ def circuit(x): [ shot_tuple.copies for shot_tuple in ( - dev.shot_vector if isinstance(dev, qml.Device) else dev.shots.shot_vector + dev.shot_vector + if isinstance(dev, qml.devices.LegacyDevice) + else dev.shots.shot_vector ) ] ) diff --git a/tests/test_vqe.py b/tests/test_vqe.py index e8edc6c84cf..ca7ea9a9e38 100644 --- a/tests/test_vqe.py +++ b/tests/test_vqe.py @@ -193,33 +193,6 @@ def amp_embed_and_strong_ent_layer(params, wires=None): (amp_embed_and_strong_ent_layer, (EMBED_PARAMS, LAYER_PARAMS)), ] -##################################################### -# Device - - -@pytest.fixture(scope="function", name="mock_device") -def mock_device_fixture(monkeypatch): - with monkeypatch.context() as m: - m.setattr(qml.Device, "__abstractmethods__", frozenset()) - m.setattr( - qml.Device, "_capabilities", {"supports_tensor_observables": True, "model": "qubit"} - ) - m.setattr(qml.Device, "operations", ["RX", "RY", "Rot", "CNOT", "Hadamard", "StatePrep"]) - m.setattr( - qml.Device, "observables", ["PauliX", "PauliY", "PauliZ", "Hadamard", "Hermitian"] - ) - m.setattr(qml.Device, "short_name", "MockDevice") - m.setattr(qml.Device, "expval", lambda self, x, y, z: 1) - m.setattr(qml.Device, "var", lambda self, x, y, z: 2) - m.setattr(qml.Device, "sample", lambda self, x, y, z: 3) - m.setattr(qml.Device, "apply", lambda self, x, y, z: None) - - def get_device(wires=1): - return qml.Device(wires=wires) # pylint:disable=abstract-class-instantiated - - yield get_device - - ##################################################### # Queues diff --git a/tests/transforms/test_diagonalize_measurements.py b/tests/transforms/test_diagonalize_measurements.py index a6774662d0b..9c6c1fa7929 100644 --- a/tests/transforms/test_diagonalize_measurements.py +++ b/tests/transforms/test_diagonalize_measurements.py @@ -35,6 +35,8 @@ null_postprocessing, ) +# pylint: disable=protected-access + class TestDiagonalizeObservable: """Tests for the _diagonalize_observable method""" @@ -310,10 +312,15 @@ def test_diagonalize_all_measurements(self, to_eigvals): new_tape = tapes[0] if to_eigvals: - assert new_tape.measurements == [ - ExpectationMP(eigvals=[1.0, -1.0], wires=[0]), - VarianceMP(eigvals=[2.0, 0.0, 0.0, -2.0], wires=[1, 2]), - ] + assert len(new_tape.measurements) == 2 + assert isinstance(new_tape.measurements[0], ExpectationMP) + assert isinstance(new_tape.measurements[1], VarianceMP) + assert new_tape.measurements[0].wires == qml.wires.Wires([0]) + assert new_tape.measurements[1].wires == qml.wires.Wires([1, 2]) + assert qml.math.allclose(sorted(new_tape.measurements[0]._eigvals), [-1.0, 1.0]) + assert qml.math.allclose( + sorted(new_tape.measurements[1]._eigvals), [-2.0, 0.0, 0.0, 2.0] + ) else: assert new_tape.measurements == [qml.expval(Z(0)), qml.var(Z(1) + Z(2))] assert new_tape.operations == diagonalize_qwc_pauli_words([X(0), X(1), Y(2)])[0] @@ -439,11 +446,16 @@ def test_with_duplicate_measurements(self, to_eigvals, supported_base_obs): new_tape = tapes[0] if to_eigvals: - assert new_tape.measurements == [ - ExpectationMP(eigvals=[1.0, -1], wires=[0]), - VarianceMP(eigvals=[2.0, 0.0, 0.0, -2.0], wires=[1, 2]), - SampleMP(eigvals=[1.0, -1.0, -1.0, 1.0], wires=[0, 2]), - ] + assert len(new_tape.measurements) == 3 + assert isinstance(new_tape.measurements[0], ExpectationMP) + assert isinstance(new_tape.measurements[1], VarianceMP) + assert isinstance(new_tape.measurements[2], SampleMP) + assert new_tape.measurements[0].wires == qml.wires.Wires([0]) + assert new_tape.measurements[1].wires == qml.wires.Wires([1, 2]) + assert new_tape.measurements[2].wires == qml.wires.Wires([0, 2]) + assert np.allclose(sorted(new_tape.measurements[0]._eigvals), [-1.0, 1]) + assert np.allclose(sorted(new_tape.measurements[1]._eigvals), [-2.0, 0, 0, 2.0]) + assert np.allclose(sorted(new_tape.measurements[2]._eigvals), [-1.0, -1.0, 1.0, 1.0]) else: assert new_tape.measurements == [ qml.expval(Z(0)), diff --git a/tests/transforms/test_qcut.py b/tests/transforms/test_qcut.py index 1a78c1ab4e2..96b0d1e3af8 100644 --- a/tests/transforms/test_qcut.py +++ b/tests/transforms/test_qcut.py @@ -4681,7 +4681,7 @@ def test_init_raises(self, devices, imbalance_tolerance, num_fragments_probed): """Test if ill-initialized instances throw errors.""" if ( - isinstance(devices, (qml.Device, qml.devices.Device)) + isinstance(devices, qml.devices.Device) and imbalance_tolerance is None and num_fragments_probed is None ):