From 5f4169a297605ab6c785eef11828aecb749067d5 Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Tue, 10 Sep 2024 16:01:34 -0400 Subject: [PATCH 01/15] Remove `Operator.expand` (#6227) `Operator.expand` provided a slower version of `Operator.decomposition`, and was essentially unused, duplicate information. Therefore it was deprecated last cycle and is ready for removal. [sc-73119] --------- Co-authored-by: Astral Cai --- doc/development/deprecations.rst | 5 +++++ doc/releases/changelog-dev.md | 3 +++ pennylane/operation.py | 27 --------------------------- tests/test_operation.py | 15 --------------- 4 files changed, 8 insertions(+), 42 deletions(-) diff --git a/doc/development/deprecations.rst b/doc/development/deprecations.rst index 3954c633d00..e9d23e790db 100644 --- a/doc/development/deprecations.rst +++ b/doc/development/deprecations.rst @@ -82,6 +82,11 @@ Other deprecations Completed deprecation cycles ---------------------------- +* `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/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 7ac1ce219bb..c46feffc318 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -77,6 +77,9 @@ 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 👋

Documentation 📝

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/tests/test_operation.py b/tests/test_operation.py index 3062803df0c..a9938b66817 100644 --- a/tests/test_operation.py +++ b/tests/test_operation.py @@ -321,21 +321,6 @@ class DummyOp(Operator): assert op._ndim_params == (ndim_params,) assert op.ndim_params == (0,) - def test_expand_deprecated(self): - - class MyOp(qml.operation.Operation): - num_wires = 1 - has_decomposition = True - - @staticmethod - def compute_decomposition(*params, wires=None, **hyperparameters): - return [qml.Hadamard(wires=wires)] - - op = MyOp(wires=0) - - with pytest.warns(qml.PennyLaneDeprecationWarning, match="'Operator.expand' is deprecated"): - op.expand() - class TestPytreeMethods: def test_pytree_defaults(self): From 2ec6442fcee69379998caf632855224cf08bafda Mon Sep 17 00:00:00 2001 From: ringo-but-quantum Date: Wed, 11 Sep 2024 09:51:44 +0000 Subject: [PATCH 02/15] [no ci] bump nightly version --- pennylane/_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/_version.py b/pennylane/_version.py index 86f7c12846d..7a90dd82366 100644 --- a/pennylane/_version.py +++ b/pennylane/_version.py @@ -16,4 +16,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "0.39.0-dev11" +__version__ = "0.39.0-dev12" From df613f60e97b9583de2f38b6561d02e2f79db84d Mon Sep 17 00:00:00 2001 From: ringo-but-quantum Date: Thu, 12 Sep 2024 09:51:40 +0000 Subject: [PATCH 03/15] [no ci] bump nightly version --- pennylane/_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/_version.py b/pennylane/_version.py index 7a90dd82366..0c24fd96c36 100644 --- a/pennylane/_version.py +++ b/pennylane/_version.py @@ -16,4 +16,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "0.39.0-dev12" +__version__ = "0.39.0-dev13" From 6f3c27687078ad77e21dcb39e3813476dbab9089 Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Thu, 12 Sep 2024 10:06:42 -0400 Subject: [PATCH 04/15] Remove default qubit jax (#6209) [sc-72785] Yet another device removal. Code cov failures are in default.qubit.legacy and QubitDevice, so can be ignored. --- doc/development/guide/logging.rst | 2 +- doc/introduction/interfaces/jax.rst | 8 +- doc/releases/changelog-dev.md | 3 +- pennylane/devices/__init__.py | 1 - pennylane/devices/default_qubit_jax.py | 371 --- pennylane/devices/default_qubit_legacy.py | 1 - setup.py | 1 - tests/devices/test_default_qubit_autograd.py | 1 - tests/devices/test_default_qubit_jax.py | 1336 -------- tests/devices/test_default_qubit_legacy.py | 41 - .../core/test_adjoint_metric_tensor.py | 20 +- .../gradients/core/test_hadamard_gradient.py | 44 +- tests/gradients/core/test_jvp.py | 36 +- tests/gradients/core/test_pulse_gradient.py | 133 +- tests/gradients/core/test_pulse_odegen.py | 75 +- tests/gradients/core/test_vjp.py | 31 +- .../finite_diff/test_finite_difference.py | 57 +- .../test_finite_difference_shot_vec.py | 51 +- .../finite_diff/test_spsa_gradient.py | 48 +- .../test_spsa_gradient_shot_vec.py | 51 +- .../parameter_shift/test_parameter_shift.py | 89 +- .../test_parameter_shift_shot_vec.py | 60 +- .../test_jax_jit_legacy.py | 899 ----- .../test_jax_jit_qnode_legacy.py | 2955 ----------------- .../test_jax_legacy.py | 876 ----- .../test_jax_qnode_legacy.py | 2507 -------------- .../test_jax_qnode_shot_vector_legacy.py | 932 ------ tests/pulse/test_parametrized_evolution.py | 30 +- .../test_subroutines/test_qubitization.py | 8 +- tests/test_qubit_device.py | 1 - tests/test_return_types_qnode.py | 6 +- 31 files changed, 281 insertions(+), 10393 deletions(-) delete mode 100644 pennylane/devices/default_qubit_jax.py delete mode 100644 tests/devices/test_default_qubit_jax.py delete mode 100644 tests/interfaces/legacy_devices_integration/test_jax_jit_legacy.py delete mode 100644 tests/interfaces/legacy_devices_integration/test_jax_jit_qnode_legacy.py delete mode 100644 tests/interfaces/legacy_devices_integration/test_jax_legacy.py delete mode 100644 tests/interfaces/legacy_devices_integration/test_jax_qnode_legacy.py delete mode 100644 tests/interfaces/legacy_devices_integration/test_jax_qnode_shot_vector_legacy.py 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/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-dev.md b/doc/releases/changelog-dev.md index c46feffc318..2f0eeee194f 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -59,9 +59,10 @@ * 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`, and `DefaultQubitJax` 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) * `expand_fn`, `max_expansion`, `override_shots`, and `device_batch_transform` are removed from the signature of `qml.execute`. diff --git a/pennylane/devices/__init__.py b/pennylane/devices/__init__.py index 87dc6c3f6e6..823f2b7f41d 100644 --- a/pennylane/devices/__init__.py +++ b/pennylane/devices/__init__.py @@ -27,7 +27,6 @@ default_qubit default_qubit_legacy - default_qubit_jax default_qubit_autograd default_gaussian default_mixed 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..bdcfd1ad1da 100644 --- a/pennylane/devices/default_qubit_legacy.py +++ b/pennylane/devices/default_qubit_legacy.py @@ -715,7 +715,6 @@ def capabilities(cls): returns_state=True, passthru_devices={ "autograd": "default.qubit.autograd", - "jax": "default.qubit.jax", }, ) return capabilities diff --git a/setup.py b/setup.py index f1f77907b6a..ec97eea8503 100644 --- a/setup.py +++ b/setup.py @@ -52,7 +52,6 @@ "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/devices/test_default_qubit_autograd.py b/tests/devices/test_default_qubit_autograd.py index 2041e2f402a..bb7c3cda0b4 100644 --- a/tests/devices/test_default_qubit_autograd.py +++ b/tests/devices/test_default_qubit_autograd.py @@ -57,7 +57,6 @@ def test_defines_correct_capabilities(self): "supports_broadcasting": True, "passthru_devices": { "autograd": "default.qubit.autograd", - "jax": "default.qubit.jax", }, } assert cap == capabilities 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..3ab933495fe 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( @@ -1009,7 +1008,6 @@ def test_defines_correct_capabilities(self): "supports_broadcasting": True, "passthru_devices": { "autograd": "default.qubit.autograd", - "jax": "default.qubit.jax", }, } assert cap == capabilities @@ -2085,34 +2083,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""" @@ -2432,17 +2402,6 @@ def test_trainable_autograd(self, is_state_batched): 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/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 1dfc7ff9777..e8ad4e1d614 100644 --- a/tests/gradients/core/test_jvp.py +++ b/tests/gradients/core/test_jvp.py @@ -654,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) @@ -672,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) @@ -686,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: @@ -709,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) @@ -724,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) @@ -746,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) @@ -762,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) @@ -785,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) 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..237178b538d 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) @@ -3598,14 +3586,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 +3606,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 +3615,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 +3633,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 +3642,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 +3659,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 +3669,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..d7a0b893ce0 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) @@ -874,18 +866,12 @@ 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, 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 +894,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 +955,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 +963,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 +982,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) @@ -2268,15 +2252,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 +2280,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 +2290,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 +2317,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 +2326,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 +2348,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 +2358,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_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/pulse/test_parametrized_evolution.py b/tests/pulse/test_parametrized_evolution.py index 802f02c64e4..051c4d72182 100644 --- a/tests/pulse/test_parametrized_evolution.py +++ b/tests/pulse/test_parametrized_evolution.py @@ -21,7 +21,7 @@ import pytest import pennylane as qml -from pennylane.devices import DefaultQubit, DefaultQubitLegacy +from pennylane.devices import DefaultQubit from pennylane.operation import AnyWires from pennylane.ops import QubitUnitary from pennylane.pulse import ParametrizedEvolution, ParametrizedHamiltonian @@ -544,29 +544,20 @@ def test_return_intermediate_and_complementary(self, comp, len_t): class TestIntegration: """Integration tests for the ParametrizedEvolution class.""" - @pytest.mark.parametrize("device_class", ["DefaultQubit", "DefaultQubitJax"]) @pytest.mark.parametrize("time", [0.3, 1, [0, 2], [0.4, 2], (3, 3.1)]) @pytest.mark.parametrize("time_interface", ["python", "numpy", "jax"]) @pytest.mark.parametrize("use_jit", [False, True]) - def test_time_input_formats(self, device_class, time, time_interface, use_jit): + def test_time_input_formats(self, time, time_interface, use_jit): import jax import jax.numpy as jnp - from pennylane.devices.default_qubit_jax import DefaultQubitJax - if time_interface == "jax": time = jnp.array(time) elif time_interface == "numpy": time = np.array(time) H = qml.pulse.ParametrizedHamiltonian([2], [qml.PauliX(0)]) - # This weird-looking code is a temporary solution to be able - # to access both DefaultQubit and DefaultQubitJax without - # having to the break the parameterization of the test. - # Once DefaultQubitJax is removed, the 'device_class' - # parameter would be redundant and dev would always be - # default qubit. - dev = {**globals(), **locals()}[device_class](wires=1) + dev = DefaultQubit(wires=1) @qml.qnode(dev, interface="jax") def circuit(t): @@ -580,16 +571,15 @@ def circuit(t): duration = time if qml.math.ndim(time) == 0 else time[1] - time[0] assert qml.math.isclose(res, qml.math.cos(4 * duration)) - @pytest.mark.parametrize("device_class", [DefaultQubit, DefaultQubitLegacy]) # pylint: disable=unused-argument - def test_time_independent_hamiltonian(self, device_class): + def test_time_independent_hamiltonian(self): """Test the execution of a time independent hamiltonian.""" import jax import jax.numpy as jnp H = time_independent_hamiltonian() - dev = device_class(wires=2) + dev = DefaultQubit(wires=2) t = 4 @@ -621,9 +611,8 @@ def true_circuit(params): jax.grad(jitted_circuit)(params), jax.grad(true_circuit)(params), atol=1e-3 ) - @pytest.mark.parametrize("device_class", [DefaultQubit, DefaultQubitLegacy]) @pytest.mark.slow - def test_time_dependent_hamiltonian(self, device_class): + def test_time_dependent_hamiltonian(self): """Test the execution of a time dependent hamiltonian. This test approximates the time-ordered exponential with a product of exponentials using small time steps. For more information, see https://en.wikipedia.org/wiki/Ordered_exponential.""" @@ -632,7 +621,7 @@ def test_time_dependent_hamiltonian(self, device_class): H = time_dependent_hamiltonian() - dev = device_class(wires=2) + dev = DefaultQubit(wires=2) t = 0.1 def generator(params): @@ -717,8 +706,7 @@ def true_circuit(params): jax.grad(jitted_circuit)(params), jax.grad(true_circuit)(params), atol=1e-3 ) - @pytest.mark.parametrize("device_class", [DefaultQubit, DefaultQubitLegacy]) - def test_two_commuting_parametrized_hamiltonians(self, device_class): + def test_two_commuting_parametrized_hamiltonians(self): """Test that the evolution of two parametrized hamiltonians that commute with each other is equal to evolve the two hamiltonians simultaneously. This test uses 8 wires for the device to test the case where 2 * n < N (the matrix is evolved instead of the state).""" @@ -742,7 +730,7 @@ def f3(p, t): ops = [qml.PauliX(0), qml.PauliX(2)] H2_ = qml.dot(coeffs, ops) - dev = device_class(wires=8) + dev = DefaultQubit(wires=8) @jax.jit @qml.qnode(dev, interface="jax") diff --git a/tests/templates/test_subroutines/test_qubitization.py b/tests/templates/test_subroutines/test_qubitization.py index 3381b47c30e..6545af9bbd5 100644 --- a/tests/templates/test_subroutines/test_qubitization.py +++ b/tests/templates/test_subroutines/test_qubitization.py @@ -169,8 +169,7 @@ def test_qnode_autograd(self): @pytest.mark.jax @pytest.mark.parametrize("use_jit", (False, True)) @pytest.mark.parametrize("shots", (None, 50000)) - @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) - def test_qnode_jax(self, shots, use_jit, device): + def test_qnode_jax(self, shots, use_jit): """ "Test that the QNode executes and is differentiable with JAX. The shots argument controls whether autodiff or parameter-shift gradients are used.""" import jax @@ -185,10 +184,7 @@ def test_qnode_jax(self, shots, use_jit, device): jax.config.update("jax_enable_x64", True) - if device == "default.qubit": - dev = qml.device("default.qubit", shots=shots, seed=10) - else: - dev = qml.device("default.qubit.legacy", shots=shots, wires=5) + dev = qml.device("default.qubit", shots=shots, seed=10) diff_method = "backprop" if shots is None else "parameter-shift" qnode = qml.QNode(self.circuit, dev, interface="jax", diff_method=diff_method) diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py index 1f02976527a..58fc05370ec 100644 --- a/tests/test_qubit_device.py +++ b/tests/test_qubit_device.py @@ -1511,7 +1511,6 @@ class TestResourcesTracker: devices = ( "default.qubit.legacy", "default.qubit.autograd", - "default.qubit.jax", ) @pytest.mark.all_interfaces diff --git a/tests/test_return_types_qnode.py b/tests/test_return_types_qnode.py index f1dfa2cabdd..6f74f84bb14 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 From 502cd0322b0cf1416463a35db09def76bbfdb8dc Mon Sep 17 00:00:00 2001 From: Astral Cai Date: Thu, 12 Sep 2024 13:31:45 -0400 Subject: [PATCH 05/15] Fix failing `test_diagonalize_all_measurements` for legacy opmath (#6252) `qml.eigvals` computes eigenvalues numerically from the matrix using `np.linalg.eigvalsh`, which always returns eigenvalues in ascending order, different from how `CompositeOp` computes it. This is fine, because there is no intrinsic order of eigenvalues anyway. I have tested that the correctness of things like `expval` of legacy `Hamiltonians` is not affected. --- tests/devices/test_default_qutrit_mixed.py | 2 +- .../transforms/test_diagonalize_measurements.py | 17 +++++++++++++---- 2 files changed, 14 insertions(+), 5 deletions(-) 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/transforms/test_diagonalize_measurements.py b/tests/transforms/test_diagonalize_measurements.py index a6774662d0b..abaeef466f8 100644 --- a/tests/transforms/test_diagonalize_measurements.py +++ b/tests/transforms/test_diagonalize_measurements.py @@ -310,10 +310,19 @@ 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), # pylint: disable=protected-access + [-1.0, 1.0], + ) + assert qml.math.allclose( + sorted(new_tape.measurements[1]._eigvals), # pylint: disable=protected-access + [-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] From 62b1f06a4996f6061ee7f7efb4a7cf2ee93761dd Mon Sep 17 00:00:00 2001 From: Mudit Pandey Date: Thu, 12 Sep 2024 15:25:04 -0400 Subject: [PATCH 06/15] RC sync for `v0.38.1` to `master` (#6265) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Astral Cai Co-authored-by: Christina Lee Co-authored-by: Utkarsh Co-authored-by: Pietropaolo Frisoni Co-authored-by: Guillermo Alonso-Linaje <65235481+KetpuntoG@users.noreply.github.com> Co-authored-by: Justin Pickering <79890410+justinpickering@users.noreply.github.com> Co-authored-by: lillian542 <38584660+lillian542@users.noreply.github.com> Co-authored-by: Jack Brown Co-authored-by: Diksha Dhawan <40900030+ddhawan11@users.noreply.github.com> Co-authored-by: soranjh <40344468+soranjh@users.noreply.github.com> Co-authored-by: soranjh Co-authored-by: Cristian Emiliano Godinez Ramirez <57567043+EmilianoG-byte@users.noreply.github.com> Co-authored-by: Alex Preciado Co-authored-by: Jorge J. Martínez de Lejarza <61199780+gmlejarza@users.noreply.github.com> Co-authored-by: Isaac De Vlugt <34751083+isaacdevlugt@users.noreply.github.com> Co-authored-by: Josh Izaac Co-authored-by: Austin Huang <65315367+austingmhuang@users.noreply.github.com> Co-authored-by: ringo-but-quantum Co-authored-by: Matthew Silverman Co-authored-by: David Wierichs Co-authored-by: Will Co-authored-by: Lee James O'Riordan Co-authored-by: anthayes92 <34694788+anthayes92@users.noreply.github.com> --- doc/development/release_notes.md | 2 ++ doc/releases/changelog-0.38.0.md | 2 +- doc/releases/changelog-0.38.1.md | 14 ++++++++++++++ pennylane/devices/qubit/apply_operation.py | 4 ++-- pennylane/measurements/expval.py | 2 +- pennylane/ops/op_math/controlled_decompositions.py | 2 +- pennylane/ops/qubit/special_unitary.py | 4 ++-- 7 files changed, 23 insertions(+), 7 deletions(-) create mode 100644 doc/releases/changelog-0.38.1.md 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/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/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/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/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 From 30331fd2d033a6ad2750f763433fa1116f34ec6b Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Thu, 12 Sep 2024 17:48:56 -0400 Subject: [PATCH 07/15] Deprecate`QNode.gradient_fn` (#6244) **Context:** The existence of `QNode.gradient_fn` ties us in to a bit more of an object oriented framework with a lot of in-place mutation of the qnode. By freeing ourselves from this property, we can have a bit more of a functional structure with less coupling and side effects. It will also free us up to start making other logical changes and improvements. `QNode.gradient_fn` is also not really defined, so it's hard to tell what it should actually be and reflect. Things have changed enough recently with more dynamic gradient validation, that it no longer really carries the same information it did when it was added. There isn't really a good analog yet of `QNode.gradient_fn`, since it's kinda a "processed diff method". We do have stories for next quarter to start adding helper transforms for things like this, but it won't be immediate. **Description of the Change:** **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:** [sc-71844] --------- Co-authored-by: Pietropaolo Frisoni --- doc/development/deprecations.rst | 6 + doc/releases/changelog-dev.md | 3 + pennylane/resource/specs.py | 14 ++- pennylane/workflow/construct_batch.py | 35 ++++-- pennylane/workflow/qnode.py | 115 +++++++++-------- .../test_default_qubit_preprocessing.py | 1 - tests/devices/test_default_mixed_autograd.py | 13 -- tests/devices/test_default_mixed_jax.py | 13 -- tests/devices/test_default_mixed_tf.py | 13 -- tests/devices/test_default_mixed_torch.py | 13 -- tests/devices/test_default_qubit_autograd.py | 16 --- tests/devices/test_default_qutrit.py | 3 - .../test_autograd_qnode_legacy.py | 15 ++- tests/interfaces/test_autograd_qnode.py | 15 ++- tests/interfaces/test_jax_jit_qnode.py | 10 +- tests/interfaces/test_jax_qnode.py | 10 +- tests/interfaces/test_tensorflow_qnode.py | 10 +- tests/interfaces/test_torch_qnode.py | 5 +- tests/logging/test_logging_autograd.py | 12 +- tests/test_qnode.py | 107 +++++++++++++--- tests/test_qnode_legacy.py | 117 +----------------- 21 files changed, 258 insertions(+), 288 deletions(-) diff --git a/doc/development/deprecations.rst b/doc/development/deprecations.rst index e9d23e790db..ffd490f09ba 100644 --- a/doc/development/deprecations.rst +++ b/doc/development/deprecations.rst @@ -9,6 +9,12 @@ deprecations are listed below. Pending deprecations -------------------- +* `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.39 + - Will be removed in v0.40 + * 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. diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 2f0eeee194f..73c266f2833 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -83,6 +83,9 @@

Deprecations 👋

+* `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 🐛

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/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/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/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 index bb7c3cda0b4..87d402cab03 100644 --- a/tests/devices/test_default_qubit_autograd.py +++ b/tests/devices/test_default_qubit_autograd.py @@ -82,8 +82,6 @@ def circuit(x): 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): @@ -100,7 +98,6 @@ def circuit(x): 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): @@ -278,7 +275,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) @@ -314,7 +310,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) @@ -419,9 +414,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) @@ -617,14 +609,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_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/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py index eb65a655877..97f9ff5a58f 100644 --- a/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py +++ b/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py @@ -589,16 +589,25 @@ def cost_fn(a, b): 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 + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + 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" + with pytest.warns( + qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" + ): + assert cost_fn.gradient_fn == "backprop" cost_fn(a, b) - 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 assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift 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_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 +455,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.""" From 87c1bc383983727ebfc563c3f9a5ffa7fab70148 Mon Sep 17 00:00:00 2001 From: ringo-but-quantum Date: Fri, 13 Sep 2024 09:51:50 +0000 Subject: [PATCH 08/15] [no ci] bump nightly version --- pennylane/_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/_version.py b/pennylane/_version.py index 0c24fd96c36..a6ead820881 100644 --- a/pennylane/_version.py +++ b/pennylane/_version.py @@ -16,4 +16,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "0.39.0-dev13" +__version__ = "0.39.0-dev14" From e484ba21c7a46b50b020d2026ba824a96b418d83 Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Fri, 13 Sep 2024 10:36:10 -0400 Subject: [PATCH 09/15] deprecate set_shots (#6250) [sc-71546] We no longer interact with the legacy device interface during out workflow. Therefore, we should always set shots via the new methods and not use the `set_shots` context manager. --- pennylane/devices/_qubit_device.py | 10 +++++++- pennylane/measurements/classical_shadow.py | 9 +++++++- pennylane/workflow/set_shots.py | 21 +++++++++++++++++ .../test_set_shots_legacy.py | 23 +++++++++++-------- 4 files changed, 51 insertions(+), 12 deletions(-) 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/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/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/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 From 060bb9a7479dcf6735671d809da19d500c4b7a89 Mon Sep 17 00:00:00 2001 From: Will Date: Fri, 13 Sep 2024 11:53:28 -0400 Subject: [PATCH 10/15] Fix FABLE template to return the correct result in JIT mode (#6263) This PR fixes bug #6262 --- doc/releases/changelog-dev.md | 3 +++ pennylane/templates/subroutines/fable.py | 20 ++++++++++--------- .../templates/test_subroutines/test_fable.py | 11 ++++++---- 3 files changed, 21 insertions(+), 13 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 73c266f2833..90dbe86a773 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -105,6 +105,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/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/tests/templates/test_subroutines/test_fable.py b/tests/templates/test_subroutines/test_fable.py index 8649fe71748..d2ba5f2496a 100644 --- a/tests/templates/test_subroutines/test_fable.py +++ b/tests/templates/test_subroutines/test_fable.py @@ -235,7 +235,7 @@ def circuit_jax(input_matrix): assert np.allclose(gradient_numeric, gradient_jax[0, 0], rtol=0.001) @pytest.mark.jax - def test_fable_grad_jax_jit(self, input_matrix): + def test_fable_jax_jit(self, input_matrix): """Test that FABLE is differentiable when using jax.""" import jax import jax.numpy as jnp @@ -272,18 +272,21 @@ def test_fable_grad_jax_jit(self, input_matrix): input_jax_negative_delta = jnp.array(input_negative_delta) input_matrix_jax = jnp.array(input_matrix) - @jax.jit @qml.qnode(dev, diff_method="backprop") def circuit_jax(input_matrix): qml.FABLE(input_matrix, wires=range(5), tol=0) return qml.expval(qml.PauliZ(wires=0)) - grad_fn = jax.grad(circuit_jax) + jitted_fn = jax.jit(circuit_jax) + + grad_fn = jax.grad(jitted_fn) gradient_numeric = ( circuit_jax(input_jax_positive_delta) - circuit_jax(input_jax_negative_delta) ) / (2 * delta) gradient_jax = grad_fn(input_matrix_jax) - assert np.allclose(gradient_numeric, gradient_jax[0, 0], rtol=0.001) + + assert qml.math.allclose(gradient_numeric, gradient_jax[0, 0], rtol=0.001) + assert qml.math.allclose(jitted_fn(input_matrix), circuit_jax(input_matrix)) @pytest.mark.jax def test_fable_grad_jax_jit_error(self, input_matrix): From 1cbaeeccf554b0b8e5a3f5e9746f2a48b6ab6758 Mon Sep 17 00:00:00 2001 From: Astral Cai Date: Fri, 13 Sep 2024 16:20:14 -0400 Subject: [PATCH 11/15] Fix failing test in legacy opmath (#6272) A follow up of https://github.com/PennyLaneAI/pennylane/pull/6252, does the same but with a different test. --- .../test_diagonalize_measurements.py | 25 +++++++++++-------- 1 file changed, 14 insertions(+), 11 deletions(-) diff --git a/tests/transforms/test_diagonalize_measurements.py b/tests/transforms/test_diagonalize_measurements.py index abaeef466f8..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""" @@ -315,13 +317,9 @@ def test_diagonalize_all_measurements(self, to_eigvals): 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[0]._eigvals), # pylint: disable=protected-access - [-1.0, 1.0], - ) - assert qml.math.allclose( - sorted(new_tape.measurements[1]._eigvals), # pylint: disable=protected-access - [-2.0, 0.0, 0.0, 2.0], + 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))] @@ -448,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)), From 08961ddf5c6a42c62b427c2ccbfcf6f0ca7d7520 Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Fri, 13 Sep 2024 16:54:07 -0400 Subject: [PATCH 12/15] remove default.qubit.autograd (#6210) [sc-72795] --------- Co-authored-by: Pietropaolo Frisoni --- Makefile | 2 - doc/development/deprecations.rst | 24 +- doc/releases/changelog-dev.md | 5 +- pennylane/devices/__init__.py | 4 - pennylane/devices/default_qubit_autograd.py | 141 - pennylane/devices/default_qubit_legacy.py | 4 +- pennylane/devices/legacy_facade.py | 94 +- pennylane/devices/tests/conftest.py | 1 - pennylane/tape/qscript.py | 2 +- pennylane/workflow/execution.py | 2 +- setup.py | 1 - tests/devices/test_default_qubit_autograd.py | 786 ------ tests/devices/test_default_qubit_legacy.py | 17 +- .../test_default_qubit_legacy_broadcasting.py | 2024 -------------- tests/devices/test_legacy_facade.py | 74 - .../test_parameter_shift_shot_vec.py | 1 - .../test_autograd_legacy.py | 1367 ---------- .../test_autograd_qnode_legacy.py | 2365 ----------------- .../test_autograd_qnode_shot_vector_legacy.py | 658 ----- tests/interfaces/test_execute.py | 13 + tests/test_qnode_legacy.py | 15 - tests/test_qubit_device.py | 5 +- 22 files changed, 39 insertions(+), 7566 deletions(-) delete mode 100644 pennylane/devices/default_qubit_autograd.py delete mode 100644 tests/devices/test_default_qubit_autograd.py delete mode 100644 tests/devices/test_default_qubit_legacy_broadcasting.py delete mode 100644 tests/interfaces/legacy_devices_integration/test_autograd_legacy.py delete mode 100644 tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py delete mode 100644 tests/interfaces/legacy_devices_integration/test_autograd_qnode_shot_vector_legacy.py 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 ffd490f09ba..967417ad71b 100644 --- a/doc/development/deprecations.rst +++ b/doc/development/deprecations.rst @@ -15,18 +15,6 @@ Pending deprecations - Deprecated in v0.39 - Will be removed in v0.40 -* All of the legacy devices (any with the name ``default.qubit.{autograd,torch,tf,jax,legacy}``) are deprecated. Use ``default.qubit`` instead, - as it supports backpropagation for the many backends the legacy devices support. - - - Deprecated in v0.38 - - Will be removed in v0.39 - -* The 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``. - - - Deprecated in v0.38 - - Will be removed in v0.39 - * The ``decomp_depth`` argument in ``qml.device`` is deprecated. - Deprecated in v0.38 @@ -88,6 +76,18 @@ 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 diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 90dbe86a773..7b9416d62c3 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -59,10 +59,12 @@ * Remove support for Python 3.9. [(#6223)](https://github.com/PennyLaneAI/pennylane/pull/6223) -* `DefaultQubitTF`, `DefaultQubitTorch`, and `DefaultQubitJax` 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`. @@ -81,6 +83,7 @@ * `Operator.expand` is now removed. Use `qml.tape.QuantumScript(op.deocomposition())` instead. [(#6227)](https://github.com/PennyLaneAI/pennylane/pull/6227) +

Deprecations 👋

* `QNode.gradient_fn` is deprecated. Please use `QNode.diff_method` and `QNode.get_gradient_fn` instead. diff --git a/pennylane/devices/__init__.py b/pennylane/devices/__init__.py index 823f2b7f41d..a542ba7df1d 100644 --- a/pennylane/devices/__init__.py +++ b/pennylane/devices/__init__.py @@ -27,7 +27,6 @@ default_qubit default_qubit_legacy - default_qubit_autograd default_gaussian default_mixed default_qutrit @@ -153,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/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_legacy.py b/pennylane/devices/default_qubit_legacy.py index bdcfd1ad1da..868c426d47f 100644 --- a/pennylane/devices/default_qubit_legacy.py +++ b/pennylane/devices/default_qubit_legacy.py @@ -713,9 +713,7 @@ def capabilities(cls): supports_analytic_computation=True, supports_broadcasting=True, returns_state=True, - passthru_devices={ - "autograd": "default.qubit.autograd", - }, + 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/tests/conftest.py b/pennylane/devices/tests/conftest.py index 8e5e9740da1..5ea86da0aae 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", } 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/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/setup.py b/setup.py index ec97eea8503..e13673fb1fa 100644 --- a/setup.py +++ b/setup.py @@ -51,7 +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.mixed = pennylane.devices.default_mixed:DefaultMixed", "null.qubit = pennylane.devices.null_qubit:NullQubit", "default.qutrit = pennylane.devices.default_qutrit:DefaultQutrit", diff --git a/tests/devices/test_default_qubit_autograd.py b/tests/devices/test_default_qubit_autograd.py deleted file mode 100644 index 87d402cab03..00000000000 --- a/tests/devices/test_default_qubit_autograd.py +++ /dev/null @@ -1,786 +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", - }, - } - 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 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 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)) - - 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)) - - 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) - - 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) - - 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_legacy.py b/tests/devices/test_default_qubit_legacy.py index 3ab933495fe..f3c47f90e23 100644 --- a/tests/devices/test_default_qubit_legacy.py +++ b/tests/devices/test_default_qubit_legacy.py @@ -1006,9 +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", - }, + "passthru_devices": {}, } assert cap == capabilities @@ -2389,19 +2387,6 @@ def test_super_expval_not_called(self, is_state_batched, mocker): 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)) - 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_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/gradients/parameter_shift/test_parameter_shift_shot_vec.py b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py index d7a0b893ce0..51ffd43753b 100644 --- a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py +++ b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py @@ -864,7 +864,6 @@ 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)]) def test_fallback_probs( self, RX, RY, argnum, mocker, broadcast 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 97f9ff5a58f..00000000000 --- a/tests/interfaces/legacy_devices_integration/test_autograd_qnode_legacy.py +++ /dev/null @@ -1,2365 +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 - with pytest.warns( - qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" - ): - 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" - with pytest.warns( - qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" - ): - assert cost_fn.gradient_fn == "backprop" - - cost_fn(a, b) - with pytest.warns( - qml.PennyLaneDeprecationWarning, match=r"QNode.gradient_fn is deprecated" - ): - 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/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/test_qnode_legacy.py b/tests/test_qnode_legacy.py index 8d433da51fe..f8e56cc6fb4 100644 --- a/tests/test_qnode_legacy.py +++ b/tests/test_qnode_legacy.py @@ -76,21 +76,6 @@ def capabilities(self): return capabilities -def test_backprop_switching_deprecation(): - """Test that a PennyLaneDeprecationWarning is raised when a device is subtituted - for a different backprop device. - """ - - with pytest.warns(qml.PennyLaneDeprecationWarning): - - @qml.qnode(DummyDevice(shots=None), interface="autograd") - def circ(x): - qml.RX(x, 0) - return qml.expval(qml.Z(0)) - - circ(pnp.array(3)) - - # pylint: disable=too-many-public-methods class TestValidation: """Tests for QNode creation and validation""" diff --git a/tests/test_qubit_device.py b/tests/test_qubit_device.py index 58fc05370ec..c42ca5a6d28 100644 --- a/tests/test_qubit_device.py +++ b/tests/test_qubit_device.py @@ -1508,10 +1508,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", - ) + devices = ("default.qubit.legacy",) @pytest.mark.all_interfaces @pytest.mark.parametrize("dev_name", devices) From 43cac7225d8b1c5f29391c5c4e542c006eb6ff12 Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Fri, 13 Sep 2024 17:37:12 -0400 Subject: [PATCH 13/15] Deprecate top level access to legacy device base classes (#6238) **Context:** `qml.Device`, `qml.QubitDevice`, and `qml.QutritDevice` all reflect the legacy device interface, which is no longer the recommended way of creating devices. **Description of the Change:** Deprecate top level access to the `Device`, `QubitDevice`, and `QutritDevice`. **Benefits:** Further isolation of the legacy device interface. **Possible Drawbacks:** All deprecations propagate through the ecosystem and can cause issues. **Related GitHub Issues:** [sc-71519] --------- Co-authored-by: Mudit Pandey --- .github/workflows/interface-unit-tests.yml | 3 +- doc/development/deprecations.rst | 7 ++ doc/releases/changelog-dev.md | 5 + pennylane/__init__.py | 19 +++- pennylane/devices/tests/conftest.py | 12 -- .../tests/test_compare_default_qubit.py | 5 +- pennylane/devices/tests/test_gates.py | 22 ++-- .../devices/tests/test_gates_with_expval.py | 2 +- pennylane/devices/tests/test_measurements.py | 56 +++++----- pennylane/devices/tests/test_templates.py | 2 +- pennylane/devices/tests/test_tracker.py | 8 +- pennylane/qcut/cutstrategy.py | 2 +- pennylane/workflow/interfaces/jax_jit.py | 15 +-- pennylane/workflow/jacobian_products.py | 6 +- tests/conftest.py | 32 +++--- tests/devices/test_default_qubit_legacy.py | 4 +- .../{test_device.py => test_legacy_device.py} | 10 +- tests/{ => devices}/test_qubit_device.py | 12 +- tests/{ => devices}/test_qutrit_device.py | 8 +- .../parameter_shift/test_parameter_shift.py | 3 +- .../test_parameter_shift_shot_vec.py | 3 +- tests/interfaces/test_jacobian_products.py | 15 --- tests/test_debugging.py | 10 +- tests/test_qnode_legacy.py | 4 +- tests/test_return_types_qnode.py | 104 +++++++++++++----- tests/test_vqe.py | 27 ----- tests/transforms/test_qcut.py | 2 +- 27 files changed, 218 insertions(+), 180 deletions(-) rename tests/devices/{test_device.py => test_legacy_device.py} (99%) rename tests/{ => devices}/test_qubit_device.py (99%) rename tests/{ => devices}/test_qutrit_device.py (99%) 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/doc/development/deprecations.rst b/doc/development/deprecations.rst index 967417ad71b..c91979963a9 100644 --- a/doc/development/deprecations.rst +++ b/doc/development/deprecations.rst @@ -9,6 +9,13 @@ deprecations are listed below. Pending deprecations -------------------- +* ``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 top level access in v0.39 + - Top level access removed in v0.40 + * `QNode.gradient_fn` is deprecated. Please use `QNode.diff_method` instead. `QNode.get_gradient_fn` can also be used to process the diff method. diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 7b9416d62c3..488fd8d37e8 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -86,6 +86,11 @@

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) 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/devices/tests/conftest.py b/pennylane/devices/tests/conftest.py index 5ea86da0aae..f462138c4d5 100755 --- a/pennylane/devices/tests/conftest.py +++ b/pennylane/devices/tests/conftest.py @@ -110,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") @@ -141,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/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/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 d5ede1227a5..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 @@ -334,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: @@ -399,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/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/test_default_qubit_legacy.py b/tests/devices/test_default_qubit_legacy.py index f3c47f90e23..11ca082441c 100644 --- a/tests/devices/test_default_qubit_legacy.py +++ b/tests/devices/test_default_qubit_legacy.py @@ -2344,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) @@ -2382,7 +2382,7 @@ 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() 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/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 c42ca5a6d28..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() @@ -1655,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" @@ -1676,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/gradients/parameter_shift/test_parameter_shift.py b/tests/gradients/parameter_shift/test_parameter_shift.py index 237178b538d..e5a00b2cdf0 100644 --- a/tests/gradients/parameter_shift/test_parameter_shift.py +++ b/tests/gradients/parameter_shift/test_parameter_shift.py @@ -3563,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 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 51ffd43753b..6d9212dceb9 100644 --- a/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py +++ b/tests/gradients/parameter_shift/test_parameter_shift_shot_vec.py @@ -2228,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): 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/test_debugging.py b/tests/test_debugging.py index bdd10a4cfd5..4e049fd3e72 100644 --- a/tests/test_debugging.py +++ b/tests/test_debugging.py @@ -176,7 +176,7 @@ def circuit(): if "mixed" in dev.name: qml.Snapshot(measurement=qml.density_matrix(wires=[0, 1])) - if isinstance(dev, qml.QutritDevice): + if isinstance(dev, qml.devices.QutritDevice): return qml.expval(qml.GellMann(0, 1)) return qml.expval(qml.PauliZ(0)) @@ -191,7 +191,7 @@ def circuit(): @pytest.mark.parametrize("diff_method", [None, "parameter-shift"]) def test_all_state_measurement_snapshot_pure_qubit_dev(self, dev, diff_method): """Test that the correct measurement snapshots are returned for different measurement types.""" - if isinstance(dev, (qml.devices.default_mixed.DefaultMixed, qml.QutritDevice)): + if isinstance(dev, (qml.devices.default_mixed.DefaultMixed, qml.devices.QutritDevice)): pytest.skip() @qml.qnode(dev, diff_method=diff_method) @@ -230,7 +230,7 @@ def test_empty_snapshots(self, dev): @qml.qnode(dev) def circuit(): - if isinstance(dev, qml.QutritDevice): + if isinstance(dev, qml.devices.QutritDevice): qml.THadamard(wires=0) return qml.expval(qml.GellMann(0, index=6)) @@ -238,7 +238,7 @@ def circuit(): return qml.expval(qml.PauliX(0)) result = qml.snapshots(circuit)() - if isinstance(dev, qml.QutritDevice): + if isinstance(dev, qml.devices.QutritDevice): expected = {"execution_results": np.array(0.66666667)} else: expected = {"execution_results": np.array(1.0)} @@ -700,7 +700,7 @@ def circuit(): assert np.allclose(result["execution_results"], expected["execution_results"]) - del result["execution_results"] + del result["execution_results"] # pylint: disable=unsupported-delete-operation del expected["execution_results"] _compare_numpy_dicts(result, expected) diff --git a/tests/test_qnode_legacy.py b/tests/test_qnode_legacy.py index f8e56cc6fb4..3ee36d99bdb 100644 --- a/tests/test_qnode_legacy.py +++ b/tests/test_qnode_legacy.py @@ -813,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 6f74f84bb14..364eb468922 100644 --- a/tests/test_return_types_qnode.py +++ b/tests/test_return_types_qnode.py @@ -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_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 ): From d0344b02309fee6b2d01c7a9a433e89af4ad5a78 Mon Sep 17 00:00:00 2001 From: David Wierichs Date: Sun, 15 Sep 2024 21:47:02 +0200 Subject: [PATCH 14/15] [Program Capture] Add pytree support to captured `qml.grad` and `qml.jacobian` (#6134) **Context:** #6120 and #6127 add support to capture `qml.grad` and `qml.jacobian` in plxpr. Once captured, they dispatch to `jax.grad` and `jax.jacobian`. **Description of the Change:** This PR adds support for pytree inputs and outputs of the differentiated functions, similar to #6081. For this, it extends the internal class `FlatFn` by the extra functionality to turn the wrapper into a `*flat_args -> *flat_outputs` function, instead of a `*pytree_args -> *flat_outputs` function. **Benefits:** Pytree support :deciduous_tree: **Possible Drawbacks:** **Related GitHub Issues:** [sc-70930] [sc-71862] --------- Co-authored-by: Christina Lee Co-authored-by: Mudit Pandey --- doc/releases/changelog-dev.md | 3 +- pennylane/_grad.py | 50 ++++++-- pennylane/capture/__init__.py | 3 + pennylane/capture/capture_diff.py | 2 +- pennylane/capture/capture_qnode.py | 10 +- pennylane/capture/explanations.md | 7 +- pennylane/capture/flatfn.py | 31 ++++- tests/capture/test_capture_diff.py | 188 ++++++++++++++++++++++++++--- 8 files changed, 260 insertions(+), 34 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 488fd8d37e8..0a93727855d 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -25,9 +25,10 @@ * Differentiation of hybrid programs via `qml.grad` and `qml.jacobian` can now be captured into plxpr. When evaluating a captured `qml.grad` (`qml.jacobian`) instruction, it will dispatch to `jax.grad` (`jax.jacobian`), which differs from the Autograd implementation - without capture. + without capture. Pytree inputs and outputs are supported. [(#6120)](https://github.com/PennyLaneAI/pennylane/pull/6120) [(#6127)](https://github.com/PennyLaneAI/pennylane/pull/6127) + [(#6134)](https://github.com/PennyLaneAI/pennylane/pull/6134) * Improve unit testing for capturing of nested control flows. [(#6111)](https://github.com/PennyLaneAI/pennylane/pull/6111) diff --git a/pennylane/_grad.py b/pennylane/_grad.py index 859ae5d9fbb..d7cb52e0d52 100644 --- a/pennylane/_grad.py +++ b/pennylane/_grad.py @@ -25,6 +25,7 @@ from pennylane.capture import enabled from pennylane.capture.capture_diff import _get_grad_prim, _get_jacobian_prim +from pennylane.capture.flatfn import FlatFn from pennylane.compiler import compiler from pennylane.compiler.compiler import CompileError @@ -33,18 +34,53 @@ def _capture_diff(func, argnum=None, diff_prim=None, method=None, h=None): """Capture-compatible gradient computation.""" - import jax # pylint: disable=import-outside-toplevel + # pylint: disable=import-outside-toplevel + import jax + from jax.tree_util import tree_flatten, tree_leaves, tree_unflatten, treedef_tuple - if isinstance(argnum, int): - argnum = [argnum] if argnum is None: - argnum = [0] + argnum = 0 + if argnum_is_int := isinstance(argnum, int): + argnum = [argnum] @wraps(func) def new_func(*args, **kwargs): - jaxpr = jax.make_jaxpr(partial(func, **kwargs))(*args) - prim_kwargs = {"argnum": argnum, "jaxpr": jaxpr.jaxpr, "n_consts": len(jaxpr.consts)} - return diff_prim.bind(*jaxpr.consts, *args, **prim_kwargs, method=method, h=h) + flat_args, in_trees = zip(*(tree_flatten(arg) for arg in args)) + full_in_tree = treedef_tuple(in_trees) + + # Create a new input tree that only takes inputs marked by argnum into account + trainable_in_trees = (in_tree for i, in_tree in enumerate(in_trees) if i in argnum) + # If an integer was provided as argnum, unpack the arguments axis of the derivatives + if argnum_is_int: + trainable_in_tree = list(trainable_in_trees)[0] + else: + trainable_in_tree = treedef_tuple(trainable_in_trees) + + # Create argnum for the flat list of input arrays. For each flattened argument, + # add a list of flat argnums if the argument is trainable and an empty list otherwise. + start = 0 + flat_argnum_gen = ( + ( + list(range(start, (start := start + len(flat_arg)))) + if i in argnum + else list(range((start := start + len(flat_arg)), start)) + ) + for i, flat_arg in enumerate(flat_args) + ) + flat_argnum = sum(flat_argnum_gen, start=[]) + + # Create fully flattened function (flat inputs & outputs) + flat_fn = FlatFn(partial(func, **kwargs) if kwargs else func, full_in_tree) + flat_args = sum(flat_args, start=[]) + jaxpr = jax.make_jaxpr(flat_fn)(*flat_args) + prim_kwargs = {"argnum": flat_argnum, "jaxpr": jaxpr.jaxpr, "n_consts": len(jaxpr.consts)} + out_flat = diff_prim.bind(*jaxpr.consts, *flat_args, **prim_kwargs, method=method, h=h) + # flatten once more to go from 2D derivative structure (outputs, args) to flat structure + out_flat = tree_leaves(out_flat) + assert flat_fn.out_tree is not None, "out_tree should be set after executing flat_fn" + # The derivative output tree is the composition of output tree and trainable input trees + combined_tree = flat_fn.out_tree.compose(trainable_in_tree) + return tree_unflatten(combined_tree, out_flat) return new_func diff --git a/pennylane/capture/__init__.py b/pennylane/capture/__init__.py index 6deeef29682..2e43a246e2c 100644 --- a/pennylane/capture/__init__.py +++ b/pennylane/capture/__init__.py @@ -34,6 +34,7 @@ ~create_measurement_wires_primitive ~create_measurement_mcm_primitive ~qnode_call + ~FlatFn The ``primitives`` submodule offers easy access to objects with jax dependencies such as @@ -154,6 +155,7 @@ def _(*args, **kwargs): create_measurement_mcm_primitive, ) from .capture_qnode import qnode_call +from .flatfn import FlatFn # by defining this here, we avoid # E0611: No name 'AbstractOperator' in module 'pennylane.capture' (no-name-in-module) @@ -196,4 +198,5 @@ def __getattr__(key): "AbstractOperator", "AbstractMeasurement", "qnode_prim", + "FlatFn", ) diff --git a/pennylane/capture/capture_diff.py b/pennylane/capture/capture_diff.py index 92dde5a2956..829a9516af1 100644 --- a/pennylane/capture/capture_diff.py +++ b/pennylane/capture/capture_diff.py @@ -103,7 +103,7 @@ def _(*args, argnum, jaxpr, n_consts, method, h): def func(*inner_args): return jax.core.eval_jaxpr(jaxpr, consts, *inner_args) - return jax.jacobian(func, argnums=argnum)(*args) + return jax.tree_util.tree_leaves(jax.jacobian(func, argnums=argnum)(*args)) # pylint: disable=unused-argument @jacobian_prim.def_abstract_eval diff --git a/pennylane/capture/capture_qnode.py b/pennylane/capture/capture_qnode.py index f7f06451230..491b9f3f6a4 100644 --- a/pennylane/capture/capture_qnode.py +++ b/pennylane/capture/capture_qnode.py @@ -82,8 +82,12 @@ def _(*args, qnode, shots, device, qnode_kwargs, qfunc_jaxpr, n_consts): mps = qfunc_jaxpr.outvars return _get_shapes_for(*mps, shots=shots, num_device_wires=len(device.wires)) - def _qnode_jvp(*args_and_tangents, **impl_kwargs): - return jax.jvp(partial(qnode_prim.impl, **impl_kwargs), *args_and_tangents) + def make_zero(tan, arg): + return jax.lax.zeros_like_array(arg) if isinstance(tan, ad.Zero) else tan + + def _qnode_jvp(args, tangents, **impl_kwargs): + tangents = tuple(map(make_zero, tangents, args)) + return jax.jvp(partial(qnode_prim.impl, **impl_kwargs), args, tangents) ad.primitive_jvps[qnode_prim] = _qnode_jvp @@ -174,7 +178,7 @@ def f(x): qnode_kwargs = {"diff_method": qnode.diff_method, **execute_kwargs, **mcm_config} qnode_prim = _get_qnode_prim() - flat_args, _ = jax.tree_util.tree_flatten(args) + flat_args = jax.tree_util.tree_leaves(args) res = qnode_prim.bind( *qfunc_jaxpr.consts, *flat_args, diff --git a/pennylane/capture/explanations.md b/pennylane/capture/explanations.md index 61b88b94030..84feef9786f 100644 --- a/pennylane/capture/explanations.md +++ b/pennylane/capture/explanations.md @@ -159,12 +159,11 @@ You can also see the const variable `a` as argument `e:i32[]` to the inner neste ### Pytree handling Evaluating a jaxpr requires accepting and returning a flat list of tensor-like inputs and outputs. -list of tensor-like outputs. These long lists can be hard to manage and are very -restrictive on the allowed functions, but we can take advantage of pytrees to allow handling -arbitrary functions. +These long lists can be hard to manage and are very restrictive on the allowed functions, but we +can take advantage of pytrees to allow handling arbitrary functions. To start, we import the `FlatFn` helper. This class converts a function to one that caches -the resulting result pytree into `flat_fn.out_tree` when executed. This can be used to repack the +the result pytree into `flat_fn.out_tree` when executed. This can be used to repack the results into the correct shape. It also returns flattened results. This does not particularly matter for program capture, as we will only be producing jaxpr from the function, not calling it directly. diff --git a/pennylane/capture/flatfn.py b/pennylane/capture/flatfn.py index 4ba6005f41e..59e1c52b948 100644 --- a/pennylane/capture/flatfn.py +++ b/pennylane/capture/flatfn.py @@ -29,23 +29,48 @@ class FlatFn: property, so that the results can be repacked later. It also returns flattened results instead of the original result object. + If an ``in_tree`` is provided, the function accepts flattened inputs instead of the + original inputs with tree structure given by ``in_tree``. + + **Example** + + >>> import jax + >>> from pennylane.capture.flatfn import FlatFn >>> def f(x): ... return {"y": 2+x["x"]} >>> flat_f = FlatFn(f) - >>> res = flat_f({"x": 0}) + >>> arg = {"x": 0.5} + >>> res = flat_f(arg) + >>> res + [2.5] + >>> jax.tree_util.tree_unflatten(flat_f.out_tree, res) + {'y': 2.5} + + If we want to use a fully flattened function that also takes flat inputs instead of + the original inputs with tree structure, we can provide the treedef for this input + structure: + + >>> flat_args, in_tree = jax.tree_util.tree_flatten((arg,)) + >>> flat_f = FlatFn(f, in_tree) + >>> res = flat_f(*flat_args) >>> res - [2] + [2.5] >>> jax.tree_util.tree_unflatten(flat_f.out_tree, res) {'y': 2.5} + Note that the ``in_tree`` has to be created by flattening a tuple of all input + arguments, even if there is only a single argument. """ - def __init__(self, f): + def __init__(self, f, in_tree=None): self.f = f + self.in_tree = in_tree self.out_tree = None update_wrapper(self, f) def __call__(self, *args): + if self.in_tree is not None: + args = jax.tree_util.tree_unflatten(self.in_tree, args) out = self.f(*args) out_flat, out_tree = jax.tree_util.tree_flatten(out) self.out_tree = out_tree diff --git a/tests/capture/test_capture_diff.py b/tests/capture/test_capture_diff.py index cf7834aafb8..07596e01b47 100644 --- a/tests/capture/test_capture_diff.py +++ b/tests/capture/test_capture_diff.py @@ -99,11 +99,11 @@ def func_jax(x): assert qml.math.allclose(func_qml(x), jax_out) # Check overall jaxpr properties - if isinstance(argnum, int): - argnum = [argnum] jaxpr = jax.make_jaxpr(func_qml)(x) assert jaxpr.in_avals == [jax.core.ShapedArray((), fdtype, weak_type=True)] assert len(jaxpr.eqns) == 3 + if isinstance(argnum, int): + argnum = [argnum] assert jaxpr.out_avals == [jax.core.ShapedArray((), fdtype, weak_type=True)] * len(argnum) grad_eqn = jaxpr.eqns[2] @@ -260,6 +260,106 @@ def circuit(x): jax.config.update("jax_enable_x64", initial_mode) + @pytest.mark.parametrize("argnum", ([0, 1], [0], [1])) + def test_grad_pytree_input(self, x64_mode, argnum): + """Test that the qml.grad primitive can be captured with pytree inputs.""" + + initial_mode = jax.config.jax_enable_x64 + jax.config.update("jax_enable_x64", x64_mode) + fdtype = jax.numpy.float64 if x64_mode else jax.numpy.float32 + + def inner_func(x, y): + return jnp.prod(jnp.sin(x["a"]) * jnp.cos(y[0]["b"][1]) ** 2) + + def func_qml(x): + return qml.grad(inner_func, argnum=argnum)( + {"a": x}, ({"b": [None, 0.4 * jnp.sqrt(x)]},) + ) + + def func_jax(x): + return jax.grad(inner_func, argnums=argnum)( + {"a": x}, ({"b": [None, 0.4 * jnp.sqrt(x)]},) + ) + + x = 0.7 + jax_out = func_jax(x) + jax_out_flat, jax_out_tree = jax.tree_util.tree_flatten(jax_out) + qml_out_flat, qml_out_tree = jax.tree_util.tree_flatten(func_qml(x)) + assert jax_out_tree == qml_out_tree + assert qml.math.allclose(jax_out_flat, qml_out_flat) + + # Check overall jaxpr properties + jaxpr = jax.make_jaxpr(func_qml)(x) + assert jaxpr.in_avals == [jax.core.ShapedArray((), fdtype, weak_type=True)] + assert len(jaxpr.eqns) == 3 + argnum = [argnum] if isinstance(argnum, int) else argnum + assert jaxpr.out_avals == [jax.core.ShapedArray((), fdtype, weak_type=True)] * len(argnum) + + grad_eqn = jaxpr.eqns[2] + diff_eqn_assertions(grad_eqn, grad_prim, argnum=argnum) + assert [var.aval for var in grad_eqn.outvars] == jaxpr.out_avals + assert len(grad_eqn.params["jaxpr"].eqns) == 6 # 5 numeric eqns, 1 conversion eqn + + manual_out = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, x) + manual_out_flat, manual_out_tree = jax.tree_util.tree_flatten(manual_out) + # Assert that the output from the manual evaluation is flat + assert manual_out_tree == jax.tree_util.tree_flatten(manual_out_flat)[1] + assert qml.math.allclose(jax_out_flat, manual_out_flat) + + jax.config.update("jax_enable_x64", initial_mode) + + @pytest.mark.parametrize("argnum", ([0, 1, 2], [0, 2], [1], 0)) + def test_grad_qnode_with_pytrees(self, argnum, x64_mode): + """Test capturing the gradient of a qnode that uses Pytrees.""" + # pylint: disable=protected-access + initial_mode = jax.config.jax_enable_x64 + jax.config.update("jax_enable_x64", x64_mode) + fdtype = jax.numpy.float64 if x64_mode else jax.numpy.float32 + + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y, z): + qml.RX(x["a"], wires=0) + qml.RY(y, wires=0) + qml.RZ(z[1][0], wires=0) + return qml.expval(qml.X(0)) + + dcircuit = qml.grad(circuit, argnum=argnum) + x = {"a": 0.6, "b": 0.9} + y = 0.6 + z = ({"c": 0.5}, [0.2, 0.3]) + qml_out = dcircuit(x, y, z) + qml_out_flat, qml_out_tree = jax.tree_util.tree_flatten(qml_out) + jax_out = jax.grad(circuit, argnums=argnum)(x, y, z) + jax_out_flat, jax_out_tree = jax.tree_util.tree_flatten(jax_out) + assert jax_out_tree == qml_out_tree + assert qml.math.allclose(jax_out_flat, qml_out_flat) + + jaxpr = jax.make_jaxpr(dcircuit)(x, y, z) + + assert len(jaxpr.eqns) == 1 # grad equation + assert jaxpr.in_avals == [jax.core.ShapedArray((), fdtype, weak_type=True)] * 6 + argnum = [argnum] if isinstance(argnum, int) else argnum + num_out_avals = 2 * (0 in argnum) + (1 in argnum) + 3 * (2 in argnum) + assert jaxpr.out_avals == [jax.core.ShapedArray((), fdtype, weak_type=True)] * num_out_avals + + grad_eqn = jaxpr.eqns[0] + assert all(invar.aval == in_aval for invar, in_aval in zip(grad_eqn.invars, jaxpr.in_avals)) + flat_argnum = [0, 1] * (0 in argnum) + [2] * (1 in argnum) + [3, 4, 5] * (2 in argnum) + diff_eqn_assertions(grad_eqn, grad_prim, argnum=flat_argnum) + grad_jaxpr = grad_eqn.params["jaxpr"] + assert len(grad_jaxpr.eqns) == 1 # qnode equation + + flat_args = jax.tree_util.tree_leaves((x, y, z)) + manual_out = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, *flat_args) + manual_out_flat, manual_out_tree = jax.tree_util.tree_flatten(manual_out) + # Assert that the output from the manual evaluation is flat + assert manual_out_tree == jax.tree_util.tree_flatten(manual_out_flat)[1] + assert qml.math.allclose(jax_out_flat, manual_out_flat) + + jax.config.update("jax_enable_x64", initial_mode) + def _jac_allclose(jac1, jac2, num_axes, atol=1e-8): """Test that two Jacobians, given as nested sequences of arrays, are equal.""" @@ -279,10 +379,6 @@ class TestJacobian: @pytest.mark.parametrize("argnum", ([0, 1], [0], [1], 0, 1)) def test_classical_jacobian(self, x64_mode, argnum): """Test that the qml.jacobian primitive can be captured with classical nodes.""" - if isinstance(argnum, list) and len(argnum) > 1: - # These cases will only be unlocked with Pytree support - pytest.xfail() - initial_mode = jax.config.jax_enable_x64 jax.config.update("jax_enable_x64", x64_mode) fdtype = jnp.float64 if x64_mode else jnp.float32 @@ -307,14 +403,14 @@ def inner_func(x, y): func_jax = jax.jacobian(inner_func, argnums=argnum) jax_out = func_jax(x, y) - num_axes = 1 if isinstance(argnum, int) else 2 - assert _jac_allclose(func_qml(x, y), jax_out, num_axes) + qml_out = func_qml(x, y) + num_axes = 1 if (int_argnum := isinstance(argnum, int)) else 2 + assert _jac_allclose(qml_out, jax_out, num_axes) # Check overall jaxpr properties - jaxpr = jax.make_jaxpr(func_jax)(x, y) jaxpr = jax.make_jaxpr(func_qml)(x, y) - if isinstance(argnum, int): + if int_argnum: argnum = [argnum] exp_in_avals = [shaped_array(shape) for shape in [(4,), (2, 3)]] @@ -331,6 +427,9 @@ def inner_func(x, y): diff_eqn_assertions(jac_eqn, jacobian_prim, argnum=argnum) manual_eval = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, x, y) + # Evaluating jaxpr gives flat list results. Need to adapt the JAX output to that + if not int_argnum: + jax_out = sum(jax_out, start=()) assert _jac_allclose(manual_eval, jax_out, num_axes) jax.config.update("jax_enable_x64", initial_mode) @@ -354,7 +453,6 @@ def func(x): return jnp.prod(x) * jnp.sin(x), jnp.sum(x**2) x = jnp.array([0.7, -0.9, 0.6, 0.3]) - x = x[:1] dim = len(x) eye = jnp.eye(dim) @@ -382,15 +480,20 @@ def func(x): # 2nd order qml_func_2 = qml.jacobian(qml_func_1) + hyperdiag = qml.numpy.zeros((4, 4, 4)) + for i in range(4): + hyperdiag[i, i, i] = 1 expected_2 = ( prod_sin[:, None, None] / x[None, :, None] / x[None, None, :] + - jnp.tensordot(prod_sin, eye / x**2, axes=0) # Correct diagonal entries + prod_cos_e_i[:, :, None] / x[None, None, :] + prod_cos_e_i[:, None, :] / x[None, :, None] - - jnp.tensordot(prod_sin, eye + eye / x**2, axes=0), - jnp.tensordot(jnp.ones(dim), eye * 2, axes=0), + - prod_sin * hyperdiag, + eye * 2, ) # Output only has one tuple axis - assert _jac_allclose(qml_func_2(x), expected_2, 1) + atol = 1e-8 if x64_mode else 2e-7 + assert _jac_allclose(qml_func_2(x), expected_2, 1, atol=atol) jaxpr_2 = jax.make_jaxpr(qml_func_2)(x) assert jaxpr_2.in_avals == [jax.core.ShapedArray((dim,), fdtype)] @@ -406,7 +509,7 @@ def func(x): assert jac_eqn.params["jaxpr"].eqns[0].primitive == jacobian_prim manual_eval_2 = jax.core.eval_jaxpr(jaxpr_2.jaxpr, jaxpr_2.consts, x) - assert _jac_allclose(manual_eval_2, expected_2, 1) + assert _jac_allclose(manual_eval_2, expected_2, 1, atol=atol) jax.config.update("jax_enable_x64", initial_mode) @@ -473,3 +576,58 @@ def circuit(x): assert _jac_allclose(manual_res, expected_res, 1) jax.config.update("jax_enable_x64", initial_mode) + + @pytest.mark.parametrize("argnum", ([0, 1], [0], [1])) + def test_jacobian_pytrees(self, x64_mode, argnum): + """Test that the qml.jacobian primitive can be captured with + pytree inputs and outputs.""" + + initial_mode = jax.config.jax_enable_x64 + jax.config.update("jax_enable_x64", x64_mode) + fdtype = jax.numpy.float64 if x64_mode else jax.numpy.float32 + + def inner_func(x, y): + return { + "prod_cos": jnp.prod(jnp.sin(x["a"]) * jnp.cos(y[0]["b"][1]) ** 2), + "sum_sin": jnp.sum(jnp.sin(x["a"]) * jnp.sin(y[1]["c"]) ** 2), + } + + def func_qml(x): + return qml.jacobian(inner_func, argnum=argnum)( + {"a": x}, ({"b": [None, 0.4 * jnp.sqrt(x)]}, {"c": 0.5}) + ) + + def func_jax(x): + return jax.jacobian(inner_func, argnums=argnum)( + {"a": x}, ({"b": [None, 0.4 * jnp.sqrt(x)]}, {"c": 0.5}) + ) + + x = 0.7 + jax_out = func_jax(x) + jax_out_flat, jax_out_tree = jax.tree_util.tree_flatten(jax_out) + qml_out_flat, qml_out_tree = jax.tree_util.tree_flatten(func_qml(x)) + assert jax_out_tree == qml_out_tree + assert qml.math.allclose(jax_out_flat, qml_out_flat) + + # Check overall jaxpr properties + jaxpr = jax.make_jaxpr(func_qml)(x) + assert jaxpr.in_avals == [jax.core.ShapedArray((), fdtype, weak_type=True)] + assert len(jaxpr.eqns) == 3 + + argnum = [argnum] if isinstance(argnum, int) else argnum + # Compute the flat argnum in order to determine the expected number of out tracers + flat_argnum = [0] * (0 in argnum) + [1, 2] * (1 in argnum) + assert jaxpr.out_avals == [jax.core.ShapedArray((), fdtype)] * (2 * len(flat_argnum)) + + jac_eqn = jaxpr.eqns[2] + + diff_eqn_assertions(jac_eqn, jacobian_prim, argnum=flat_argnum) + assert [var.aval for var in jac_eqn.outvars] == jaxpr.out_avals + + manual_out = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, x) + manual_out_flat, manual_out_tree = jax.tree_util.tree_flatten(manual_out) + # Assert that the output from the manual evaluation is flat + assert manual_out_tree == jax.tree_util.tree_flatten(manual_out_flat)[1] + assert qml.math.allclose(jax_out_flat, manual_out_flat) + + jax.config.update("jax_enable_x64", initial_mode) From 79fc6d3dadb4a1dcd435f28e0d6e26aceb72ae3f Mon Sep 17 00:00:00 2001 From: ringo-but-quantum Date: Mon, 16 Sep 2024 09:51:46 +0000 Subject: [PATCH 15/15] [no ci] bump nightly version --- pennylane/_version.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/_version.py b/pennylane/_version.py index a6ead820881..77639685bc6 100644 --- a/pennylane/_version.py +++ b/pennylane/_version.py @@ -16,4 +16,4 @@ Version number (major.minor.patch[-label]) """ -__version__ = "0.39.0-dev14" +__version__ = "0.39.0-dev15"