From 41bab702f55d78ce20e991f3fb3dee40084580cd Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 5 Mar 2024 16:25:07 -0500 Subject: [PATCH 01/19] Add `_cache_transform()` function to replace `cache_execute()` --- pennylane/workflow/execution.py | 68 ++++++++++++++++++++++++++------- 1 file changed, 54 insertions(+), 14 deletions(-) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index 650187e2fc6..b8b3cac25a6 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -27,13 +27,14 @@ import inspect import warnings from functools import wraps, partial -from typing import Callable, Sequence, Optional, Union, Tuple +from typing import Callable, MutableMapping, Sequence, Optional, Union, Tuple import logging from cachetools import LRUCache, Cache import pennylane as qml from pennylane.tape import QuantumTape +from pennylane.transforms import transform from pennylane.typing import ResultBatch from .set_shots import set_shots @@ -259,13 +260,17 @@ def _make_inner_execute( if isinstance(device, qml.Device): device_execution = set_shots(device, override_shots)(device.batch_execute) - else: device_execution = partial(device.execute, execution_config=execution_config) - cached_device_execution = qml.workflow.cache_execute( - device_execution, cache, return_tuple=False - ) + def cached_device_execution(tapes): + tapes, post_processing = _cache_transform(tapes, cache=cache) + return post_processing(device_execution(tapes)) + + if cache in [None, False]: + inner_device_execution = device_execution + else: + inner_device_execution = cached_device_execution def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: """Execution that occurs within a machine learning framework boundary. @@ -273,14 +278,14 @@ def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: Closure Variables: expand_fn (Callable[[QuantumTape], QuantumTape]): A device preprocessing step numpy_only (bool): whether or not to convert the data to numpy or leave as is - cached_device_execution (Callable[[Sequence[QuantumTape]], ResultBatch]) + inner_device_execution (Callable[[Sequence[QuantumTape]], ResultBatch]) """ if expand_fn: tapes = tuple(expand_fn(t) for t in tapes) if numpy_only: tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes) - return cached_device_execution(tapes) + return inner_device_execution(tapes) return inner_execute @@ -322,6 +327,8 @@ def cache_execute(fn: Callable, cache, pass_kwargs=False, return_tuple=True, exp function: a wrapped version of the execution function ``fn`` with caching support """ + # TODO: Add deprecation warning. + # This function has been replaced by ``_cache_transform()``. if logger.isEnabledFor(logging.DEBUG): logger.debug( "Entry with args=(fn=%s, cache=%s, pass_kwargs=%s, return_tuple=%s, expand_fn=%s) called by=%s", @@ -431,6 +438,36 @@ def wrapper(tapes: Sequence[QuantumTape], **kwargs): return wrapper +@transform +def _cache_transform(tape: QuantumTape, cache: MutableMapping): + """Caches the result of ``tape`` using the provided ``cache``.""" + + def cache_hit_postprocessing(_results: Tuple[Tuple]) -> Tuple: + result = cache[tape.hash] + if result is not None: + return result + + raise RuntimeError( + "Result for tape is missing from the execution cache. " + "This is likely the result of a race condition." + ) + + if tape.hash in cache: + return [], cache_hit_postprocessing + + def cache_miss_postprocessing(results: Tuple[Tuple]) -> Tuple: + result = results[0] + cache[tape.hash] = result + return result + + # Adding a ``None`` entry to the cache indicates that a result will eventually be available for + # the tape. This assumes that post-processing functions are called in the same order in which + # the transforms are invoked. Otherwise, ``cache_hit_postprocessing()`` may be called before the + # result of the corresponding tape is placed in the cache by ``cache_miss_postprocessing()``. + cache[tape.hash] = None + return [tape], cache_miss_postprocessing + + def execute( tapes: Sequence[QuantumTape], device: device_type, @@ -797,18 +834,21 @@ def inner_execute_with_empty_jac(tapes, **_): # replace the backward gradient computation gradient_fn_with_shots = set_shots(device, override_shots)(device.gradients) - cached_gradient_fn = qml.workflow.cache_execute( - gradient_fn_with_shots, - cache, - pass_kwargs=True, - return_tuple=False, - ) + # TODO: Use cache_transform(). + # cached_gradient_fn = qml.workflow.cache_execute( + # gradient_fn_with_shots, + # cache, + # pass_kwargs=True, + # return_tuple=False, + # ) def device_gradient_fn(inner_tapes, **gradient_kwargs): numpy_tapes = tuple( qml.transforms.convert_to_numpy_parameters(t) for t in inner_tapes ) - return cached_gradient_fn(numpy_tapes, **gradient_kwargs) + # TODO: Use the result from cache_transform() above. + # return cached_gradient_fn(numpy_tapes, **gradient_kwargs) + return gradient_fn_with_shots(numpy_tapes, **gradient_kwargs) gradient_fn = device_gradient_fn From dc13c4c07af02646911e00c3fb0db02531945d8e Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 5 Mar 2024 16:25:32 -0500 Subject: [PATCH 02/19] Update docstring to mention `_cache_transform()` --- pennylane/workflow/jacobian_products.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/workflow/jacobian_products.py b/pennylane/workflow/jacobian_products.py index 06648993fe5..3d2c92f6703 100644 --- a/pennylane/workflow/jacobian_products.py +++ b/pennylane/workflow/jacobian_products.py @@ -223,7 +223,7 @@ class TransformJacobianProducts(JacobianProductCalculator): instead of treating each call as independent. This keyword argument is used to patch problematic autograd behavior when caching is turned off. In this case, caching will be based on the identity of the batch, rather than the potentially expensive :attr:`~.QuantumScript.hash` that is used - by :func:`~.cache_execute`. + by :func:`~._cache_transform`. >>> inner_execute = qml.device('default.qubit').execute >>> gradient_transform = qml.gradients.param_shift From 1b0c129ab472dda2c5b0a168b529211385525b88 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 5 Mar 2024 16:26:05 -0500 Subject: [PATCH 03/19] Update tests to spy on `_cache_transform()` --- tests/interfaces/test_autograd.py | 10 +++++----- tests/interfaces/test_jax.py | 12 ++++++------ tests/interfaces/test_jax_jit.py | 12 ++++++------ tests/interfaces/test_tensorflow.py | 8 ++++---- tests/interfaces/test_torch.py | 8 ++++---- 5 files changed, 25 insertions(+), 25 deletions(-) diff --git a/tests/interfaces/test_autograd.py b/tests/interfaces/test_autograd.py index f914c902a52..ef1385ea3be 100644 --- a/tests/interfaces/test_autograd.py +++ b/tests/interfaces/test_autograd.py @@ -228,7 +228,7 @@ class TestCaching: 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, "cache_execute") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cachesize): with qml.queuing.AnnotatedQueue() as q: @@ -241,7 +241,7 @@ def cost(a, cachesize): params = np.array([0.1, 0.2]) qml.jacobian(cost)(params, cachesize=2) - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache.maxsize == 2 assert cache.currsize == 2 @@ -250,7 +250,7 @@ def cost(a, cachesize): def test_custom_cache(self, mocker): """Test the use of a custom cache object""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow, "cache_execute") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cache): with qml.queuing.AnnotatedQueue() as q: @@ -265,7 +265,7 @@ def cost(a, cache): params = np.array([0.1, 0.2]) qml.jacobian(cost)(params, cache=custom_cache) - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache is custom_cache def test_caching_param_shift(self, tol): @@ -398,7 +398,7 @@ def cost(a, cache): )[0] ) - # no cache_execute caching, but jac for each batch still stored. + # no caching, but jac for each batch still stored. qml.jacobian(cost)(params, cache=None) assert dev.num_executions == 2 diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py index 44801ce4c90..7a91b19aa04 100644 --- a/tests/interfaces/test_jax.py +++ b/tests/interfaces/test_jax.py @@ -195,7 +195,7 @@ class TestCaching: 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, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") def cost(a, cachesize): with qml.queuing.AnnotatedQueue() as q: @@ -213,7 +213,7 @@ def cost(a, cachesize): params = jax.numpy.array([0.1, 0.2]) jax.grad(cost)(params, cachesize=2) - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache.maxsize == 2 assert cache.currsize == 2 @@ -222,7 +222,7 @@ def cost(a, cachesize): def test_custom_cache(self, mocker): """Test the use of a custom cache object""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") def cost(a, cache): with qml.queuing.AnnotatedQueue() as q: @@ -242,13 +242,13 @@ def cost(a, cache): params = jax.numpy.array([0.1, 0.2]) jax.grad(cost)(params, cache=custom_cache) - cache = spy.call_args[0][1] + 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, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) @@ -277,7 +277,7 @@ def cost(a, b, cache): custom_cache = {} jax.grad(cost)(a, b, cache=custom_cache) - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache is custom_cache def test_caching_param_shift(self, tol): diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py index 509b537f15b..908f577de56 100644 --- a/tests/interfaces/test_jax_jit.py +++ b/tests/interfaces/test_jax_jit.py @@ -184,7 +184,7 @@ class TestCaching: 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, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") def cost(a, cachesize): with qml.queuing.AnnotatedQueue() as q: @@ -203,7 +203,7 @@ def cost(a, cachesize): params = jax.numpy.array([0.1, 0.2]) jax.jit(jax.grad(cost), static_argnums=1)(params, cachesize=2) - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache.maxsize == 2 assert cache.currsize == 2 @@ -212,7 +212,7 @@ def cost(a, cachesize): def test_custom_cache(self, mocker): """Test the use of a custom cache object""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") def cost(a, cache): with qml.queuing.AnnotatedQueue() as q: @@ -233,13 +233,13 @@ def cost(a, cache): params = jax.numpy.array([0.1, 0.2]) jax.grad(cost)(params, cache=custom_cache) - cache = spy.call_args[0][1] + 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, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) @@ -270,7 +270,7 @@ def cost(a, b, cache): custom_cache = {} jax.grad(cost)(a, b, cache=custom_cache) - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache is custom_cache def test_caching_param_shift(self, tol): diff --git a/tests/interfaces/test_tensorflow.py b/tests/interfaces/test_tensorflow.py index 05840c3cbc9..85159ff926c 100644 --- a/tests/interfaces/test_tensorflow.py +++ b/tests/interfaces/test_tensorflow.py @@ -137,7 +137,7 @@ class TestCaching: 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, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as t: @@ -150,7 +150,7 @@ def test_cache_maxsize(self, mocker): res = execute([tape], dev, gradient_fn=param_shift, cachesize=2, interface="tf")[0] t.jacobian(res, a) - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache.maxsize == 2 assert cache.currsize == 2 @@ -159,7 +159,7 @@ def test_cache_maxsize(self, mocker): 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, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") a = tf.Variable([0.1, 0.2]) custom_cache = {} @@ -176,7 +176,7 @@ def test_custom_cache(self, mocker): t.jacobian(res, a) - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache is custom_cache unwrapped_tape = qml.transforms.convert_to_numpy_parameters(tape) diff --git a/tests/interfaces/test_torch.py b/tests/interfaces/test_torch.py index 71e2145d054..2c15f06bbd7 100644 --- a/tests/interfaces/test_torch.py +++ b/tests/interfaces/test_torch.py @@ -171,7 +171,7 @@ class TestCaching: 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, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") def cost(a, cachesize): with qml.queuing.AnnotatedQueue() as q: @@ -188,7 +188,7 @@ def cost(a, cachesize): params = torch.tensor([0.1, 0.2], requires_grad=True) res = cost(params, cachesize=2) res.backward() - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache.maxsize == 2 assert cache.currsize == 2 @@ -197,7 +197,7 @@ def cost(a, cachesize): def test_custom_cache(self, mocker): """Test the use of a custom cache object""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow, "cache_execute") + spy = mocker.spy(qml.workflow.execution, "_cache_transform") def cost(a, cache): with qml.queuing.AnnotatedQueue() as q: @@ -216,7 +216,7 @@ def cost(a, cache): res = cost(params, cache=custom_cache) res.backward() - cache = spy.call_args[0][1] + cache = spy.call_args.kwargs["cache"] assert cache is custom_cache def test_caching_param_shift(self): From e8463fa64be5006d99d0d124dd68ced821cab827 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 5 Mar 2024 16:26:22 -0500 Subject: [PATCH 04/19] Add explicit tests for `_cache_transform()` --- tests/workflow/test_cache_transform.py | 90 ++++++++++++++++++++++++++ 1 file changed, 90 insertions(+) create mode 100644 tests/workflow/test_cache_transform.py diff --git a/tests/workflow/test_cache_transform.py b/tests/workflow/test_cache_transform.py new file mode 100644 index 00000000000..83e6887f1c2 --- /dev/null +++ b/tests/workflow/test_cache_transform.py @@ -0,0 +1,90 @@ +# Copyright 2018-2024 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the :func:``_cache_transform()`` function.""" +from typing import MutableMapping +import pytest + +import pennylane as qml +from pennylane.tape import QuantumScript +from pennylane.workflow.execution import _cache_transform + + +@pytest.fixture +def tape() -> QuantumScript: + """Returns a ``QuantumScript`` object.""" + return QuantumScript([], [qml.expval(qml.Z(0))]) + + +@pytest.fixture +def cache() -> MutableMapping: + """Returns an object which can be used as a cache.""" + return {} + + +def test_cache_miss_before_cache_hit(tape, cache): + """Tests that the "miss" post-processing function updates the cache so that + calling the "hit" post-processing function afterwards returns the cached + result for the tape. + """ + miss_tapes, miss_fn = _cache_transform(tape, cache=cache) + hit_tapes, hit_fn = _cache_transform(tape, cache=cache) + + assert miss_tapes + assert not hit_tapes + + result = (1.23,) + + assert miss_fn((result,)) == result + assert hit_fn(((),)) == result + + +def test_cache_hit_before_cache_miss(tape, cache): + """Tests that a RuntimeError is raised if the "hit" post-processing function + is called before the "miss" post-processing function for the same tape. + """ + _cache_transform(tape, cache=cache) + _, hit_fn = _cache_transform(tape, cache=cache) + + match = ( + r"Result for tape is missing from the execution cache\. " + r"This is likely the result of a race condition\." + ) + with pytest.raises(RuntimeError, match=match): + hit_fn(((1.23,),)) + + +def test_batch_of_different_tapes(cache): + """Tests that the results of different tapes are not cached under the same key.""" + tape_1 = QuantumScript([], [qml.expval(qml.X(0))]) + tape_2 = QuantumScript([], [qml.expval(qml.Y(0))]) + tape_3 = QuantumScript([], [qml.expval(qml.Z(0))]) + + batch_tapes, batch_fns = _cache_transform([tape_1, tape_2, tape_3], cache=cache) + assert len(batch_tapes) == 3 + + results = ((1.0,), (2.0,), (3.0,)) + assert batch_fns(results) == results + + +def test_batch_of_identical_tapes(cache): + """Tests that the result of identical tapes are cached under the same key.""" + tape_1 = QuantumScript([], [qml.expval(qml.Z(0))]) + tape_2 = QuantumScript([], [qml.expval(qml.Z(0))]) + tape_3 = QuantumScript([], [qml.expval(qml.Z(0))]) + + batch_tapes, batch_fns = _cache_transform([tape_1, tape_2, tape_3], cache=cache) + assert len(batch_tapes) == 1 + + result = (1.23,) + assert batch_fns((result,)) == (result, result, result) From eb3f1dcc1a4906fc80cd4dfba481ae696f396954 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 5 Mar 2024 17:06:46 -0500 Subject: [PATCH 05/19] Extract `_apply_cache_transform()` for better reuse --- pennylane/workflow/execution.py | 64 +++++++++++++++++++-------------- 1 file changed, 38 insertions(+), 26 deletions(-) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index b8b3cac25a6..334343c50b6 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -263,14 +263,7 @@ def _make_inner_execute( else: device_execution = partial(device.execute, execution_config=execution_config) - def cached_device_execution(tapes): - tapes, post_processing = _cache_transform(tapes, cache=cache) - return post_processing(device_execution(tapes)) - - if cache in [None, False]: - inner_device_execution = device_execution - else: - inner_device_execution = cached_device_execution + cached_device_execution = _apply_cache_transform(fn=device_execution, cache=cache) def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: """Execution that occurs within a machine learning framework boundary. @@ -278,14 +271,14 @@ def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: Closure Variables: expand_fn (Callable[[QuantumTape], QuantumTape]): A device preprocessing step numpy_only (bool): whether or not to convert the data to numpy or leave as is - inner_device_execution (Callable[[Sequence[QuantumTape]], ResultBatch]) + cached_device_execution (Callable[[Sequence[QuantumTape]], ResultBatch]) """ if expand_fn: tapes = tuple(expand_fn(t) for t in tapes) if numpy_only: tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes) - return inner_device_execution(tapes) + return cached_device_execution(tapes) return inner_execute @@ -440,7 +433,12 @@ def wrapper(tapes: Sequence[QuantumTape], **kwargs): @transform def _cache_transform(tape: QuantumTape, cache: MutableMapping): - """Caches the result of ``tape`` using the provided ``cache``.""" + """Caches the result of ``tape`` using the provided ``cache``. + + .. note:: + + This function makes use of :attr:`.QuantumTape.hash` to identify unique tapes. + """ def cache_hit_postprocessing(_results: Tuple[Tuple]) -> Tuple: result = cache[tape.hash] @@ -468,6 +466,25 @@ def cache_miss_postprocessing(results: Tuple[Tuple]) -> Tuple: return [tape], cache_miss_postprocessing +def _apply_cache_transform(fn: Callable, cache: Optional[MutableMapping]) -> Callable: + """Wraps the given execution function with ``_cache_transform()`` using the provided cache. + + Args: + fn (Callable): The execution function to be augmented with caching. This function should + have the signature ``fn(tapes, **kwargs)`` and return ``list[tensor_like]`` with the + same length as the input ``tapes``. + cache (None | MutableMapping): The cache to use. If ``None``, caching will not occur. + """ + if cache is None: + return fn + + def execution_function_with_caching(tapes): + tapes, post_processing_fn = _cache_transform(tapes, cache=cache) + return post_processing_fn(fn(tapes)) + + return execution_function_with_caching + + def execute( tapes: Sequence[QuantumTape], device: device_type, @@ -477,7 +494,7 @@ def execute( config=None, grad_on_execution="best", gradient_kwargs=None, - cache: Union[bool, dict, Cache] = True, + cache: Union[None, bool, dict, Cache] = True, cachesize=10000, max_diff=1, override_shots: int = False, @@ -508,7 +525,7 @@ def execute( pass. The 'best' option chooses automatically between the two options and is default. gradient_kwargs (dict): dictionary of keyword arguments to pass when determining the gradients of tapes - cache (bool, dict, Cache): Whether to cache evaluations. This can result in + cache (None, bool, dict, Cache): Whether to cache evaluations. This can result in a significant reduction in quantum evaluations during gradient computations. cachesize (int): the size of the cache max_diff (int): If ``gradient_fn`` is a gradient transform, this option specifies @@ -661,10 +678,13 @@ def cost_fn(params, x): else: transform_program = qml.transforms.core.TransformProgram() - if isinstance(cache, bool) and cache: - # cache=True: create a LRUCache object + # If caching is desired but an explicit cache is not provided, use an ``LRUCache``. + if cache is True: cache = LRUCache(maxsize=cachesize) - setattr(cache, "_persistent_cache", False) + + # Ensure that ``cache`` is not a Boolean to simplify downstream code. + elif cache is False: + cache = None expand_fn = _preprocess_expand_fn(expand_fn, device, max_expansion) @@ -834,21 +854,13 @@ def inner_execute_with_empty_jac(tapes, **_): # replace the backward gradient computation gradient_fn_with_shots = set_shots(device, override_shots)(device.gradients) - # TODO: Use cache_transform(). - # cached_gradient_fn = qml.workflow.cache_execute( - # gradient_fn_with_shots, - # cache, - # pass_kwargs=True, - # return_tuple=False, - # ) + cached_gradient_fn = _apply_cache_transform(fn=gradient_fn_with_shots, cache=cache) def device_gradient_fn(inner_tapes, **gradient_kwargs): numpy_tapes = tuple( qml.transforms.convert_to_numpy_parameters(t) for t in inner_tapes ) - # TODO: Use the result from cache_transform() above. - # return cached_gradient_fn(numpy_tapes, **gradient_kwargs) - return gradient_fn_with_shots(numpy_tapes, **gradient_kwargs) + return cached_gradient_fn(numpy_tapes, **gradient_kwargs) gradient_fn = device_gradient_fn From dbbd873add9c4bb6a30370f678bccbffffb4421e Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 5 Mar 2024 17:14:08 -0500 Subject: [PATCH 06/19] Add note to changelog under 'Improvements' section --- doc/releases/changelog-dev.md | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 03f2ec08c26..1cce92612bf 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -11,14 +11,17 @@ * The `molecular_hamiltonian` function calls `PySCF` directly when `method='pyscf'` is selected. [(#5118)](https://github.com/PennyLaneAI/pennylane/pull/5118) - -* All generators in the source code (except those in the `qchem` module) no longer return + +* All generators in the source code (except those in the `qchem` module) no longer return `Hamiltonian` or `Tensor` instances. Wherever possible, these return `Sum`, `SProd`, and `Prod` instances. [(#5253)](https://github.com/PennyLaneAI/pennylane/pull/5253) * Upgraded `null.qubit` to the new device API. Also, added support for all measurements and various modes of differentiation. [(#5211)](https://github.com/PennyLaneAI/pennylane/pull/5211) +* Replaced `cache_execute` with an alternate implementation based on `@transform`. + [(#5318)](https://github.com/PennyLaneAI/pennylane/pull/5318) +

Breaking changes 💔

Deprecations 👋

From f116252611747ec1d683cc2cee29ff09ea4d3427 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 5 Mar 2024 17:24:56 -0500 Subject: [PATCH 07/19] Disable pylint fixture warnings --- tests/workflow/test_cache_transform.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/workflow/test_cache_transform.py b/tests/workflow/test_cache_transform.py index 83e6887f1c2..6b687487882 100644 --- a/tests/workflow/test_cache_transform.py +++ b/tests/workflow/test_cache_transform.py @@ -11,7 +11,10 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Tests for the :func:``_cache_transform()`` function.""" +""" +Unit tests for the :func:``_cache_transform()`` function. +""" +# pylint: disable=redefined-outer-name from typing import MutableMapping import pytest From 94a8e743b5eb08da19127da9c9c6dc9a2a7585f9 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Fri, 8 Mar 2024 17:49:45 -0500 Subject: [PATCH 08/19] Restore tape with finite shots and persistent cache warning --- pennylane/workflow/execution.py | 24 +++++++++++++++--------- tests/workflow/test_cache_transform.py | 13 +++++++++++++ 2 files changed, 28 insertions(+), 9 deletions(-) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index 334343c50b6..e99763c8b10 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -86,6 +86,17 @@ """list[str]: allowed interface strings""" +_CACHED_EXECUTION_WITH_FINITE_SHOTS_WARNINGS = ( + "Cached execution with finite shots detected!\n" + "Note that samples as well as all noisy quantities computed via sampling " + "will be identical across executions. This situation arises where tapes " + "are executed with identical operations, measurements, and parameters.\n" + "To avoid this behaviour, provide 'cache=False' to the QNode or execution " + "function." +) +"""str: warning message to display when cached execution is used with finite shots""" + + def _adjoint_jacobian_expansion( tapes: Sequence[QuantumTape], grad_on_execution: bool, interface: str, max_expansion: int ): @@ -383,15 +394,7 @@ def wrapper(tapes: Sequence[QuantumTape], **kwargs): # Tape exists within the cache, store the cached result cached_results[i] = cache[hashes[i]] if tape.shots and getattr(cache, "_persistent_cache", True): - warnings.warn( - "Cached execution with finite shots detected!\n" - "Note that samples as well as all noisy quantities computed via sampling " - "will be identical across executions. This situation arises where tapes " - "are executed with identical operations, measurements, and parameters.\n" - "To avoid this behavior, provide 'cache=False' to the QNode or execution " - "function.", - UserWarning, - ) + warnings.warn(_CACHED_EXECUTION_WITH_FINITE_SHOTS_WARNINGS, UserWarning) else: # Tape does not exist within the cache, store the tape # for execution via the execution function. @@ -443,6 +446,8 @@ def _cache_transform(tape: QuantumTape, cache: MutableMapping): def cache_hit_postprocessing(_results: Tuple[Tuple]) -> Tuple: result = cache[tape.hash] if result is not None: + if tape.shots and getattr(cache, "_persistent_cache", True): + warnings.warn(_CACHED_EXECUTION_WITH_FINITE_SHOTS_WARNINGS, UserWarning) return result raise RuntimeError( @@ -681,6 +686,7 @@ def cost_fn(params, x): # If caching is desired but an explicit cache is not provided, use an ``LRUCache``. if cache is True: cache = LRUCache(maxsize=cachesize) + setattr(cache, "_persistent_cache", False) # Ensure that ``cache`` is not a Boolean to simplify downstream code. elif cache is False: diff --git a/tests/workflow/test_cache_transform.py b/tests/workflow/test_cache_transform.py index 6b687487882..c451873f7c2 100644 --- a/tests/workflow/test_cache_transform.py +++ b/tests/workflow/test_cache_transform.py @@ -91,3 +91,16 @@ def test_batch_of_identical_tapes(cache): result = (1.23,) assert batch_fns((result,)) == (result, result, result) + + +def test_finite_shots_with_persistent_cache_warning(cache): + """Tests that a UserWarning is emitted if a cache hit occurs for a tape with + finite shots that uses a persistent cache. + """ + tape = QuantumScript([], [qml.expval(qml.Z(0))], shots=1) + + batch_tapes, batch_fns = _cache_transform([tape, tape], cache=cache) + assert len(batch_tapes) == 1 + + with pytest.warns(UserWarning, match=r"Cached execution with finite shots detected!"): + batch_fns(((1.23,),)) From 34b9ea9463d6476cb5044894506acda6dacdfb61 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Fri, 8 Mar 2024 17:50:12 -0500 Subject: [PATCH 09/19] Implement `_make_inner_execute()` using TransformProgram --- pennylane/workflow/execution.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index e99763c8b10..cf90612b365 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -274,22 +274,26 @@ def _make_inner_execute( else: device_execution = partial(device.execute, execution_config=execution_config) - cached_device_execution = _apply_cache_transform(fn=device_execution, cache=cache) - def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: """Execution that occurs within a machine learning framework boundary. Closure Variables: expand_fn (Callable[[QuantumTape], QuantumTape]): A device preprocessing step numpy_only (bool): whether or not to convert the data to numpy or leave as is - cached_device_execution (Callable[[Sequence[QuantumTape]], ResultBatch]) - + device_execution (Callable[[Sequence[QuantumTape]], ResultBatch]) + cache (None | MutableMapping): The cache to use. If ``None``, caching will not occur. """ + transform_program = qml.transforms.core.TransformProgram() + transform_program.add_transform(_cache_transform, cache=cache) + + # TODO: Apply expand_fn() and convert_to_numpy_parameters() as transforms. if expand_fn: tapes = tuple(expand_fn(t) for t in tapes) if numpy_only: tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes) - return cached_device_execution(tapes) + + transformed_tapes, transform_post_processing = transform_program(tapes) + return transform_post_processing(device_execution(transformed_tapes)) return inner_execute From 664f3ee8063672bd6bf390d3c7aeb057ee08ca67 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Mon, 11 Mar 2024 11:21:44 -0400 Subject: [PATCH 10/19] Add check for `None` cache --- pennylane/workflow/execution.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index cf90612b365..909f0c5cbdf 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -284,7 +284,9 @@ def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: cache (None | MutableMapping): The cache to use. If ``None``, caching will not occur. """ transform_program = qml.transforms.core.TransformProgram() - transform_program.add_transform(_cache_transform, cache=cache) + + if cache is not None: + transform_program.add_transform(_cache_transform, cache=cache) # TODO: Apply expand_fn() and convert_to_numpy_parameters() as transforms. if expand_fn: From d7b5dba5cb73ded8a441e69fda7685f51538ac52 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 12 Mar 2024 13:31:19 -0400 Subject: [PATCH 11/19] Update remaining broken `mocker.spy` references Co-authored-by: Christina Lee --- tests/interfaces/test_jax.py | 6 +++--- tests/interfaces/test_jax_jit.py | 6 +++--- tests/interfaces/test_tensorflow.py | 4 ++-- tests/interfaces/test_torch.py | 4 ++-- 4 files changed, 10 insertions(+), 10 deletions(-) diff --git a/tests/interfaces/test_jax.py b/tests/interfaces/test_jax.py index 7a91b19aa04..b37c67e6c45 100644 --- a/tests/interfaces/test_jax.py +++ b/tests/interfaces/test_jax.py @@ -195,7 +195,7 @@ class TestCaching: 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") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cachesize): with qml.queuing.AnnotatedQueue() as q: @@ -222,7 +222,7 @@ def cost(a, cachesize): def test_custom_cache(self, mocker): """Test the use of a custom cache object""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution, "_cache_transform") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cache): with qml.queuing.AnnotatedQueue() as q: @@ -248,7 +248,7 @@ def cost(a, cache): def test_custom_cache_multiple(self, mocker): """Test the use of a custom cache object with multiple tapes""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution, "_cache_transform") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) diff --git a/tests/interfaces/test_jax_jit.py b/tests/interfaces/test_jax_jit.py index 908f577de56..b9a937837da 100644 --- a/tests/interfaces/test_jax_jit.py +++ b/tests/interfaces/test_jax_jit.py @@ -184,7 +184,7 @@ class TestCaching: 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") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cachesize): with qml.queuing.AnnotatedQueue() as q: @@ -212,7 +212,7 @@ def cost(a, cachesize): def test_custom_cache(self, mocker): """Test the use of a custom cache object""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution, "_cache_transform") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cache): with qml.queuing.AnnotatedQueue() as q: @@ -239,7 +239,7 @@ def cost(a, cache): def test_custom_cache_multiple(self, mocker): """Test the use of a custom cache object with multiple tapes""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution, "_cache_transform") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") a = jax.numpy.array(0.1) b = jax.numpy.array(0.2) diff --git a/tests/interfaces/test_tensorflow.py b/tests/interfaces/test_tensorflow.py index 85159ff926c..fef89a65e19 100644 --- a/tests/interfaces/test_tensorflow.py +++ b/tests/interfaces/test_tensorflow.py @@ -137,7 +137,7 @@ class TestCaching: 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") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") a = tf.Variable([0.1, 0.2]) with tf.GradientTape() as t: @@ -159,7 +159,7 @@ def test_cache_maxsize(self, mocker): 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") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") a = tf.Variable([0.1, 0.2]) custom_cache = {} diff --git a/tests/interfaces/test_torch.py b/tests/interfaces/test_torch.py index 2c15f06bbd7..bdcd7939160 100644 --- a/tests/interfaces/test_torch.py +++ b/tests/interfaces/test_torch.py @@ -171,7 +171,7 @@ class TestCaching: 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") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cachesize): with qml.queuing.AnnotatedQueue() as q: @@ -197,7 +197,7 @@ def cost(a, cachesize): def test_custom_cache(self, mocker): """Test the use of a custom cache object""" dev = qml.device("default.qubit.legacy", wires=1) - spy = mocker.spy(qml.workflow.execution, "_cache_transform") + spy = mocker.spy(qml.workflow.execution._cache_transform, "_transform") def cost(a, cache): with qml.queuing.AnnotatedQueue() as q: From 531beff581e6301a5caa3aad155e2af769940e95 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 12 Mar 2024 13:33:17 -0400 Subject: [PATCH 12/19] Avoid passing empty tape sequence to device execution Co-authored-by: Christina Lee --- pennylane/workflow/execution.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index 909f0c5cbdf..37f7fb854dc 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -295,6 +295,8 @@ def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes) transformed_tapes, transform_post_processing = transform_program(tapes) + if len(transformed_tapes) == 0: + return transform_post_processing(tuple()) return transform_post_processing(device_execution(transformed_tapes)) return inner_execute From 07e57ee7dad1c6fc792ad5448f20f25544b79344 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 12 Mar 2024 13:36:13 -0400 Subject: [PATCH 13/19] Avoid comparing length of sequence to 0 --- pennylane/workflow/execution.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index 37f7fb854dc..5f436788632 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -295,9 +295,13 @@ def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: tapes = tuple(qml.transforms.convert_to_numpy_parameters(t) for t in tapes) transformed_tapes, transform_post_processing = transform_program(tapes) - if len(transformed_tapes) == 0: - return transform_post_processing(tuple()) - return transform_post_processing(device_execution(transformed_tapes)) + + if transformed_tapes: + results = device_execution(transformed_tapes) + else: + results = () + + return transform_post_processing(results) return inner_execute From 521cb00fd7abba991d4bcbf3291497c725af332d Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 12 Mar 2024 14:31:50 -0400 Subject: [PATCH 14/19] Update number of expected logs in tests --- tests/logging/test_logging_autograd.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/logging/test_logging_autograd.py b/tests/logging/test_logging_autograd.py index ec968cb842a..c6f491432b6 100644 --- a/tests/logging/test_logging_autograd.py +++ b/tests/logging/test_logging_autograd.py @@ -78,7 +78,7 @@ def circuit(params): circuit(params) - assert len(caplog.records) == 4 + assert len(caplog.records) == 3 log_records_expected = [ ( @@ -99,7 +99,7 @@ def circuit(params): assert all(msg in actual.getMessage() for msg in expected[1]) @pytest.mark.parametrize( - "diff_method,num_records", [("parameter-shift", 9), ("backprop", 4), ("adjoint", 7)] + "diff_method,num_records", [("parameter-shift", 8), ("backprop", 3), ("adjoint", 6)] ) def test_dq_qnode_execution_grad(self, caplog, diff_method, num_records): "Test logging of QNode with parameterised gradients" From 234d5643af8d262a3a4419372217303c358ca0f6 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 12 Mar 2024 16:33:45 -0400 Subject: [PATCH 15/19] Delete `cache_execute()` and any mentions of it --- pennylane/workflow/__init__.py | 3 +- pennylane/workflow/execution.py | 147 +--------------------- pennylane/workflow/interfaces/autograd.py | 4 +- 3 files changed, 5 insertions(+), 149 deletions(-) diff --git a/pennylane/workflow/__init__.py b/pennylane/workflow/__init__.py index b859a2e5ba0..e8c2d739785 100644 --- a/pennylane/workflow/__init__.py +++ b/pennylane/workflow/__init__.py @@ -23,7 +23,6 @@ :toctree: api ~execute - ~workflow.cache_execute ~workflow.set_shots ~workflow.construct_batch ~workflow.get_transform_program @@ -55,6 +54,6 @@ """ from .set_shots import set_shots -from .execution import execute, cache_execute, SUPPORTED_INTERFACES, INTERFACE_MAP +from .execution import execute, SUPPORTED_INTERFACES, INTERFACE_MAP from .qnode import QNode, qnode from .construct_batch import construct_batch, get_transform_program diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index 5f436788632..26e98832388 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -12,11 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -Contains the cache_execute decoratator, for adding caching to a function -that executes multiple tapes on a device. - -Also contains the general execute function, for exectuting tapes on -devices with autodifferentiation support. +Contains the general execute function, for executing tapes on devices with auto- +differentiation support. """ # pylint: disable=import-outside-toplevel,too-many-branches,not-callable,unexpected-keyword-arg @@ -306,146 +303,6 @@ def inner_execute(tapes: Sequence[QuantumTape], **_) -> ResultBatch: return inner_execute -def cache_execute(fn: Callable, cache, pass_kwargs=False, return_tuple=True, expand_fn=None): - """Decorator that adds caching to a function that executes - multiple tapes on a device. - - This decorator makes use of :attr:`.QuantumTape.hash` to identify - unique tapes. - - - If a tape does not match a hash in the cache, then the tape - has not been previously executed. It is executed, and the result - added to the cache. - - - If a tape matches a hash in the cache, then the tape has been previously - executed. The corresponding cached result is - extracted, and the tape is not passed to the execution function. - - - Finally, there might be the case where one or more tapes in the current - set of tapes to be executed are identical and thus share a hash. If this is the case, - duplicates are removed, to avoid redundant evaluations. - - Args: - fn (callable): The execution function to add caching to. - This function should have the signature ``fn(tapes, **kwargs)``, - and it should return ``list[tensor_like]``, with the - same length as the input ``tapes``. - cache (None or dict or Cache or bool): The cache to use. If ``None``, - caching will not occur. - pass_kwargs (bool): If ``True``, keyword arguments passed to the - wrapped function will be passed directly to ``fn``. If ``False``, - they will be ignored. - return_tuple (bool): If ``True``, the output of ``fn`` is returned - as a tuple ``(fn_ouput, [])``, to match the output of execution functions - that also return gradients. - - Returns: - function: a wrapped version of the execution function ``fn`` with caching - support - """ - # TODO: Add deprecation warning. - # This function has been replaced by ``_cache_transform()``. - if logger.isEnabledFor(logging.DEBUG): - logger.debug( - "Entry with args=(fn=%s, cache=%s, pass_kwargs=%s, return_tuple=%s, expand_fn=%s) called by=%s", - ( - fn - if not (logger.isEnabledFor(qml.logging.TRACE) and inspect.isfunction(fn)) - else "\n" + inspect.getsource(fn) - ), - cache, - pass_kwargs, - return_tuple, - ( - expand_fn - if not (logger.isEnabledFor(qml.logging.TRACE) and inspect.isfunction(expand_fn)) - else "\n" + inspect.getsource(expand_fn) + "\n" - ), - "::L".join(str(i) for i in inspect.getouterframes(inspect.currentframe(), 2)[1][1:3]), - ) - - if expand_fn is not None: - original_fn = fn - - def fn(tapes: Sequence[QuantumTape], **kwargs): # pylint: disable=function-redefined - tapes = [expand_fn(tape) for tape in tapes] - return original_fn(tapes, **kwargs) - - @wraps(fn) - def wrapper(tapes: Sequence[QuantumTape], **kwargs): - if not pass_kwargs: - kwargs = {} - - if cache is None or (isinstance(cache, bool) and not cache): - # No caching. Simply execute the execution function - # and return the results. - - # must convert to list as new device interface returns tuples - res = list(fn(tapes, **kwargs)) - return (res, []) if return_tuple else res - - execution_tapes = {} - cached_results = {} - hashes = {} - repeated = {} - - for i, tape in enumerate(tapes): - h = tape.hash - - if h in hashes.values(): - # Tape already exists within ``tapes``. Determine the - # index of the first occurrence of the tape, store this, - # and continue to the next iteration. - idx = list(hashes.keys())[list(hashes.values()).index(h)] - repeated[i] = idx - continue - - hashes[i] = h - - if hashes[i] in cache: - # Tape exists within the cache, store the cached result - cached_results[i] = cache[hashes[i]] - if tape.shots and getattr(cache, "_persistent_cache", True): - warnings.warn(_CACHED_EXECUTION_WITH_FINITE_SHOTS_WARNINGS, UserWarning) - else: - # Tape does not exist within the cache, store the tape - # for execution via the execution function. - execution_tapes[i] = tape - - # if there are no execution tapes, simply return! - if not execution_tapes: - if not repeated: - res = list(cached_results.values()) - return (res, []) if return_tuple else res - - else: - # execute all unique tapes that do not exist in the cache - # convert to list as new device interface returns a tuple - res = list(fn(tuple(execution_tapes.values()), **kwargs)) - - final_res = [] - - for i, tape in enumerate(tapes): - if i in cached_results: - # insert cached results into the results vector - final_res.append(cached_results[i]) - - elif i in repeated: - # insert repeated results into the results vector - final_res.append(final_res[repeated[i]]) - - else: - # insert evaluated results into the results vector - r = res.pop(0) - final_res.append(r) - cache[hashes[i]] = r - - return (final_res, []) if return_tuple else final_res - - wrapper.fn = fn - return wrapper - - @transform def _cache_transform(tape: QuantumTape, cache: MutableMapping): """Caches the result of ``tape`` using the provided ``cache``. diff --git a/pennylane/workflow/interfaces/autograd.py b/pennylane/workflow/interfaces/autograd.py index a1d4810d69a..86670cd112f 100644 --- a/pennylane/workflow/interfaces/autograd.py +++ b/pennylane/workflow/interfaces/autograd.py @@ -70,8 +70,8 @@ def grad_fn(dy): request indepedent vjps for each entry in the output, even though the internal circuits will be exactly the same. -When normal caching provided by :func:`~.cache_execute` is present, the expensive part (re-executing -identical circuits) is avoided, but when normal caching is turned off, the above can lead to an explosion +When caching is enabled, the expensive part (re-executing identical circuits) is +avoided, but when normal caching is turned off, the above can lead to an explosion in the number of required circuit executions. To avoid this explosion in the number of executed circuits when caching is turned off, we will instead internally From 89da01d9846d1cdee2c39111ca5a1ed8fe946522 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Tue, 12 Mar 2024 16:34:51 -0400 Subject: [PATCH 16/19] Delete unused `wraps` import --- pennylane/workflow/execution.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/workflow/execution.py b/pennylane/workflow/execution.py index 26e98832388..33f8c1e3b68 100644 --- a/pennylane/workflow/execution.py +++ b/pennylane/workflow/execution.py @@ -23,7 +23,7 @@ import inspect import warnings -from functools import wraps, partial +from functools import partial from typing import Callable, MutableMapping, Sequence, Optional, Union, Tuple import logging From b43eb18ea6c6e6e40da04b816b7fe9a5d4654c47 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Wed, 13 Mar 2024 12:00:14 -0400 Subject: [PATCH 17/19] Add unit tests for `_apply_cache_transform()` --- tests/workflow/test_cache_transform.py | 29 +++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/tests/workflow/test_cache_transform.py b/tests/workflow/test_cache_transform.py index c451873f7c2..d563f78b265 100644 --- a/tests/workflow/test_cache_transform.py +++ b/tests/workflow/test_cache_transform.py @@ -12,15 +12,16 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -Unit tests for the :func:``_cache_transform()`` function. +Unit tests for the :func:``_cache_transform()`` and :func:``_apply__cache_transform()`` functions. """ -# pylint: disable=redefined-outer-name +# pylint: disable=protected-access,redefined-outer-name from typing import MutableMapping +from unittest.mock import MagicMock import pytest import pennylane as qml from pennylane.tape import QuantumScript -from pennylane.workflow.execution import _cache_transform +from pennylane.workflow.execution import _apply_cache_transform, _cache_transform @pytest.fixture @@ -35,6 +36,12 @@ def cache() -> MutableMapping: return {} +@pytest.fixture +def transform_spy(mocker) -> MagicMock: + """Returns a spy on the underlying ``_cache_transform()`` function.""" + return mocker.spy(qml.workflow.execution._cache_transform, "_transform") + + def test_cache_miss_before_cache_hit(tape, cache): """Tests that the "miss" post-processing function updates the cache so that calling the "hit" post-processing function afterwards returns the cached @@ -104,3 +111,19 @@ def test_finite_shots_with_persistent_cache_warning(cache): with pytest.warns(UserWarning, match=r"Cached execution with finite shots detected!"): batch_fns(((1.23,),)) + + +def test_apply_cache_transform_with_cache(transform_spy, tape, cache): + """Tests that calling ``_apply_cache_transform()`` with a cache returns a + function that applies the cache transform. + """ + _apply_cache_transform(MagicMock(return_value=[1.23]), cache=cache)([tape]) + transform_spy.assert_called_once_with(tape, cache=cache) + + +def test_apply_cache_transform_without_cache(transform_spy, tape): + """Tests that calling ``_apply_cache_transform()`` without a cache returns a + function that does not apply the cache transform. + """ + _apply_cache_transform(MagicMock(return_value=[1.23]), cache=None)([tape]) + transform_spy.assert_not_called() From de4fbe1a2cc34df01da0365f0b7431a5b2b5154b Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Wed, 13 Mar 2024 12:02:28 -0400 Subject: [PATCH 18/19] Fix reST formatting of test module docstring --- tests/workflow/test_cache_transform.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/workflow/test_cache_transform.py b/tests/workflow/test_cache_transform.py index d563f78b265..eae734a8c57 100644 --- a/tests/workflow/test_cache_transform.py +++ b/tests/workflow/test_cache_transform.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -Unit tests for the :func:``_cache_transform()`` and :func:``_apply__cache_transform()`` functions. +Unit tests for the :func:`_cache_transform` and :func:`_apply_cache_transform` functions. """ # pylint: disable=protected-access,redefined-outer-name from typing import MutableMapping From 3844a61db697547cc882618e715ea70d08c9ad74 Mon Sep 17 00:00:00 2001 From: Mikhail Andrenkov Date: Thu, 14 Mar 2024 09:28:21 -0400 Subject: [PATCH 19/19] Avoid referencing `_cache_transform()` in documentation Co-authored-by: Thomas R. Bromley <49409390+trbromley@users.noreply.github.com> --- pennylane/workflow/jacobian_products.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/workflow/jacobian_products.py b/pennylane/workflow/jacobian_products.py index 3d2c92f6703..12557608fc1 100644 --- a/pennylane/workflow/jacobian_products.py +++ b/pennylane/workflow/jacobian_products.py @@ -223,7 +223,7 @@ class TransformJacobianProducts(JacobianProductCalculator): instead of treating each call as independent. This keyword argument is used to patch problematic autograd behavior when caching is turned off. In this case, caching will be based on the identity of the batch, rather than the potentially expensive :attr:`~.QuantumScript.hash` that is used - by :func:`~._cache_transform`. + by the cache. >>> inner_execute = qml.device('default.qubit').execute >>> gradient_transform = qml.gradients.param_shift