From c65d301acae07edf2523df6dfbed8f20502e1c4e Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Fri, 13 May 2022 15:09:59 -0400 Subject: [PATCH 01/26] Add batch_partial implementation --- pennylane/__init__.py | 1 + pennylane/transforms/__init__.py | 2 + pennylane/transforms/batch_partial.py | 99 ++++++++++ tests/transforms/test_batch_partial.py | 249 +++++++++++++++++++++++++ 4 files changed, 351 insertions(+) create mode 100644 pennylane/transforms/batch_partial.py create mode 100644 tests/transforms/test_batch_partial.py diff --git a/pennylane/__init__.py b/pennylane/__init__.py index fcb3f0b48c3..d0b45857723 100644 --- a/pennylane/__init__.py +++ b/pennylane/__init__.py @@ -60,6 +60,7 @@ batch_params, batch_input, batch_transform, + batch_partial, cut_circuit, cut_circuit_mc, ControlledOperation, diff --git a/pennylane/transforms/__init__.py b/pennylane/transforms/__init__.py index 9b69aeec7f5..5dcce23bb3d 100644 --- a/pennylane/transforms/__init__.py +++ b/pennylane/transforms/__init__.py @@ -32,6 +32,7 @@ ~transforms.classical_jacobian ~batch_params ~batch_input + ~batch_partial ~metric_tensor ~adjoint_metric_tensor ~specs @@ -179,6 +180,7 @@ from .adjoint import adjoint from .batch_params import batch_params from .batch_input import batch_input +from .batch_partial import batch_partial from .classical_jacobian import classical_jacobian from .condition import cond, Conditional from .compile import compile diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py new file mode 100644 index 00000000000..93ce44cf984 --- /dev/null +++ b/pennylane/transforms/batch_partial.py @@ -0,0 +1,99 @@ +# 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. +""" +Contains the batch dimension transform. +""" +import copy +import functools +import inspect + +import pennylane as qml + + +def _convert_to_args(func, args, kwargs): + """ + Given a function, convert the positional and + keyword arguments to purely positional arguments. + """ + sig = inspect.signature(func).parameters + + new_args = [] + for i, param in enumerate(sig): + if param in kwargs: + # first check if the name is provided in kwargs + new_args.append(kwargs[param]) + elif i < len(sig): + # next check if the argnum is provided + new_args.append(args[i]) + else: + raise ValueError(f"Argument {param} must be provided") + + return tuple(new_args) + + +def batch_partial(qnode, **partial_kwargs): + qnode = qml.batch_params(qnode) + + # store whether this decorator is being used as a pure + # analog of functools.partial, or whether it is used to + # wrap a QNode in a more complex lambda statement + is_partial = False + if not any([callable(val) for val in partial_kwargs.values()]): + # none of the kwargs passed in are callable + is_partial = True + + sig = inspect.signature(qnode).parameters + if is_partial: + # the partially evaluated function must have at least one more + # parameter, otherwise batching doesn't make sense + if len(sig) <= len(partial_kwargs): + raise ValueError("Partial evaluation must leave at least one unevaluated parameter") + else: + # if used to wrap a QNode in a lambda statement, then check that + # all arguments are provided + if len(sig) > len(partial_kwargs): + raise ValueError("Callable argument requires all other arguments to QNode be provided") + + @functools.wraps(qnode) + def wrapper(*args, **kwargs): + + # raise an error if keyword arguments are passed, since the + # arguments are passed to the lambda statement instead of the QNode + if not is_partial and kwargs: + raise ValueError( + "Arguments must not be passed as keyword arguments to " + "callable within partial function" + ) + + # get the batch dimension (we don't have to check if all arguments + # have the same batch dim since that's done in qml.batch_params) + if args: + batch_dim = qml.math.shape(args[0])[0] + else: + batch_dim = qml.math.shape(list(kwargs.values())[0])[0] + + for key, val in partial_kwargs.items(): + if callable(val): + val = qml.math.stack([val(*a) for a in zip(*args)]) + kwargs[key] = val + else: + kwargs[key] = qml.math.stack([val] * batch_dim) + + if is_partial: + return qnode(*_convert_to_args(qnode, args, kwargs)) + else: + # don't pass the arguments to the lambda itself into the QNode + return qnode(*_convert_to_args(qnode, (), kwargs)) + + return wrapper diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py new file mode 100644 index 00000000000..03ba8b42387 --- /dev/null +++ b/tests/transforms/test_batch_partial.py @@ -0,0 +1,249 @@ +# 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. +""" +Unit tests for the batch partial transform. +""" +import pytest + +import pennylane as qml +from pennylane import numpy as np + + +def test_partial_evaluation(): + """Test partial evaluation matches individual full evaluations""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial argument to construct a new circuit with + y = np.random.uniform(size=2) + + # the batched argument to the new partial circuit + x = np.random.uniform(size=batch_size) + + batched_partial_circuit = qml.batch_partial(circuit, y=y) + res = batched_partial_circuit(x) + + # check the results against individually executed circuits + indiv_res = [] + for x_indiv in x: + indiv_res.append(circuit(x_indiv, y)) + + assert np.allclose(res, indiv_res) + + +def test_partial_evaluation_grad(): + """Test gradient of partial evaluation matches gradients of + individual full evaluations""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial argument to construct a new circuit with + y = np.random.uniform(size=2) + + # the batched argument to the new partial circuit + x = np.random.uniform(size=batch_size, requires_grad=True) + + batched_partial_circuit = qml.batch_partial(circuit, y=y) + + # we could also sum over the batch dimension and use the regular + # gradient instead of the jacobian, but either works + grad = qml.math.diagonal(qml.jacobian(batched_partial_circuit)(x)) + + # check the results against individually executed circuits + indiv_grad = [] + for x_indiv in x: + indiv_grad.append(qml.grad(circuit, argnum=0)(x_indiv, y)) + + assert np.allclose(grad, indiv_grad) + + +def test_lambda_evaluation(): + """Test lambda argument replacement matches individual full evaluations""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the first partial argument + x = np.random.uniform(size=()) + + # the base value of the second partial argument + y = np.random.uniform(size=2) + + # the second partial argument as a function of the inputs + fn = lambda y0: y + y0 * np.ones(2) + + # values for the second argument + y0 = np.random.uniform(size=batch_size) + + batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + res = batched_partial_circuit(y0) + + # check the results against individually executed circuits + indiv_res = [] + for y0_indiv in y0: + indiv_res.append(circuit(x, y + y0_indiv * np.ones(2))) + + assert np.allclose(res, indiv_res) + + +def test_lambda_evaluation_grad(): + """Test gradient of lambda argument replacement matches + gradients of individual full evaluations""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the first partial argument + x = np.random.uniform(size=()) + + # the base value of the second partial argument + y = np.random.uniform(size=2) + + # the second partial argument as a function of the inputs + fn = lambda y0: y + y0 * np.ones(2) + + # values for the second argument + y0 = np.random.uniform(size=batch_size, requires_grad=True) + + batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + + # we could also sum over the batch dimension and use the regular + # gradient instead of the jacobian, but either works + grad = qml.math.diagonal(qml.jacobian(batched_partial_circuit)(y0)) + + # check the results against individually executed circuits + indiv_grad = [] + for y0_indiv in y0: + grad_wrt_second_arg = qml.grad(circuit, argnum=1)(x, y + y0_indiv * np.ones(2)) + grad_wrt_y0 = qml.math.sum(grad_wrt_second_arg) + indiv_grad.append(grad_wrt_y0) + + assert np.allclose(grad, indiv_grad) + + +def test_full_evaluation_error(): + """Test that an error is raised when all arguments to QNode + are provided to a partial evaluation.""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial arguments + x = np.random.uniform(size=()) + y = np.random.uniform(size=2) + + with pytest.raises( + ValueError, match="Partial evaluation must leave at least one unevaluated parameter" + ): + batched_partial_circuit = qml.batch_partial(circuit, x=x, y=y) + + +def test_incomplete_evaluation_error(): + """Test that an error is raised when not all arguments to QNode + are provided to a callable wrapper""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the second partial argument as a function of the inputs + y = np.random.uniform(size=2) + fn = lambda y0: y + y0 * np.ones(2) + + # values for the second argument + y0 = np.random.uniform(size=batch_size) + + with pytest.raises( + ValueError, match="Callable argument requires all other arguments to QNode be provided" + ): + batched_partial_circuit = qml.batch_partial(circuit, y=fn) + + +def test_kwargs_callable_error(): + """Test that an error is raised when keyword arguments + are provided to a callable wrapper""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial arguments + x = np.random.uniform(size=()) + + y = np.random.uniform(size=2) + fn = lambda y0: y + y0 * np.ones(2) + y0 = np.random.uniform(size=batch_size) + + batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + + with pytest.raises( + ValueError, + match="Arguments must not be passed as keyword arguments to callable within partial function", + ): + res = batched_partial_circuit(y=y0) + + with pytest.raises( + ValueError, + match="Arguments must not be passed as keyword arguments to callable within partial function", + ): + res = batched_partial_circuit(y0=y0) From 8c80a51faab3a1c3734e798cdfd055c2ec9562a3 Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Tue, 17 May 2022 10:55:26 -0400 Subject: [PATCH 02/26] fix codefactor errors --- pennylane/transforms/batch_partial.py | 14 ++++++------ tests/transforms/test_batch_partial.py | 31 ++++++++++++++++++++++++++ 2 files changed, 38 insertions(+), 7 deletions(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index 93ce44cf984..e80b38508d1 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -14,7 +14,6 @@ """ Contains the batch dimension transform. """ -import copy import functools import inspect @@ -36,20 +35,21 @@ def _convert_to_args(func, args, kwargs): elif i < len(sig): # next check if the argnum is provided new_args.append(args[i]) - else: - raise ValueError(f"Argument {param} must be provided") return tuple(new_args) def batch_partial(qnode, **partial_kwargs): + """ + TODO: docs + """ qnode = qml.batch_params(qnode) # store whether this decorator is being used as a pure # analog of functools.partial, or whether it is used to # wrap a QNode in a more complex lambda statement is_partial = False - if not any([callable(val) for val in partial_kwargs.values()]): + if not any(callable(val) for val in partial_kwargs.values()): # none of the kwargs passed in are callable is_partial = True @@ -92,8 +92,8 @@ def wrapper(*args, **kwargs): if is_partial: return qnode(*_convert_to_args(qnode, args, kwargs)) - else: - # don't pass the arguments to the lambda itself into the QNode - return qnode(*_convert_to_args(qnode, (), kwargs)) + + # don't pass the arguments to the lambda itself into the QNode + return qnode(*_convert_to_args(qnode, (), kwargs)) return wrapper diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index 03ba8b42387..158b399ef76 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -50,6 +50,37 @@ def circuit(x, y): assert np.allclose(res, indiv_res) +def test_partial_evaluation_kwargs(): + """Test partial evaluation matches individual full evaluations + when kwargs are used""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial argument to construct a new circuit with + y = np.random.uniform(size=2) + + # the batched argument to the new partial circuit + x = np.random.uniform(size=batch_size) + + batched_partial_circuit = qml.batch_partial(circuit, y=y) + res = batched_partial_circuit(x=x) + + # check the results against individually executed circuits + indiv_res = [] + for x_indiv in x: + indiv_res.append(circuit(x_indiv, y)) + + assert np.allclose(res, indiv_res) + + def test_partial_evaluation_grad(): """Test gradient of partial evaluation matches gradients of individual full evaluations""" From 454c8e8099128d85759eb0f8baca1d1b7e49f338 Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Tue, 17 May 2022 16:02:06 -0400 Subject: [PATCH 03/26] Add tests for all interfaces --- pennylane/transforms/batch_partial.py | 4 +- tests/transforms/test_batch_partial.py | 142 ++++++++++++++++++++++++- 2 files changed, 140 insertions(+), 6 deletions(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index e80b38508d1..0098e6def41 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -39,11 +39,11 @@ def _convert_to_args(func, args, kwargs): return tuple(new_args) -def batch_partial(qnode, **partial_kwargs): +def batch_partial(qnode, all_operations=False, **partial_kwargs): """ TODO: docs """ - qnode = qml.batch_params(qnode) + qnode = qml.batch_params(qnode, all_operations=all_operations) # store whether this decorator is being used as a pure # analog of functools.partial, or whether it is used to diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index 158b399ef76..8e1eb844fd1 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -81,12 +81,14 @@ def circuit(x, y): assert np.allclose(res, indiv_res) -def test_partial_evaluation_grad(): +@pytest.mark.autograd +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +def test_partial_evaluation_autograd(diff_method): """Test gradient of partial evaluation matches gradients of - individual full evaluations""" + individual full evaluations using autograd""" dev = qml.device("default.qubit", wires=2) - @qml.qnode(dev) + @qml.qnode(dev, diff_method=diff_method) def circuit(x, y): qml.RX(x, wires=0) qml.RY(y[..., 0], wires=0) @@ -105,7 +107,7 @@ def circuit(x, y): # we could also sum over the batch dimension and use the regular # gradient instead of the jacobian, but either works - grad = qml.math.diagonal(qml.jacobian(batched_partial_circuit)(x)) + grad = np.diagonal(qml.jacobian(batched_partial_circuit)(x)) # check the results against individually executed circuits indiv_grad = [] @@ -115,6 +117,138 @@ def circuit(x, y): assert np.allclose(grad, indiv_grad) +@pytest.mark.jax +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +def test_partial_evaluation_jax(diff_method): + """Test gradient of partial evaluation matches gradients of + individual full evaluations using jax""" + import jax + import jax.numpy as jnp + + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev, interface="jax", diff_method=diff_method) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial argument to construct a new circuit with + y = jnp.array([0.3, 0.4]) + + # the batched argument to the new partial circuit + x = jnp.linspace(0.1, 0.5, batch_size) + + batched_partial_circuit = qml.batch_partial(circuit, all_operations=True, y=y) + + # we could also sum over the batch dimension and use the regular + # gradient instead of the jacobian, but either works + grad = jnp.diagonal(jax.jacrev(batched_partial_circuit)(x)) + + # check the results against individually executed circuits + indiv_grad = [] + for x_indiv in x: + indiv_grad.append(jax.grad(circuit, argnums=0)(x_indiv, y)) + + assert np.allclose(grad, indiv_grad) + + +@pytest.mark.tf +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +def test_partial_evaluation_tf(diff_method): + """Test gradient of partial evaluation matches gradients of + individual full evaluations using TF""" + import tensorflow as tf + + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev, interface="tf", diff_method=diff_method) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial argument to construct a new circuit with + y = tf.Variable(np.random.uniform(size=2), trainable=True) + + # the batched argument to the new partial circuit + x = tf.Variable(np.random.uniform(size=batch_size), trainable=True) + + batched_partial_circuit = qml.batch_partial(circuit, y=y) + + with tf.GradientTape() as tape: + out = batched_partial_circuit(x) + + # we could also sum over the batch dimension and use the regular + # gradient instead of the jacobian, but either works + grad = tf.linalg.tensor_diag_part(tape.jacobian(out, x)) + + # check the results against individually executed circuits + indiv_grad = [] + for x_indiv in x: + # create a new variable since tensors created by + # indexing a trainable variable aren't themselves trainable + x_indiv = tf.Variable(x_indiv, trainable=True) + + with tf.GradientTape() as tape: + out_indiv = circuit(x_indiv, y) + indiv_grad.append(tape.gradient(out_indiv, x_indiv)) + + assert np.allclose(grad, indiv_grad) + + +@pytest.mark.torch +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +def test_partial_evaluation_torch(diff_method): + """Test gradient of partial evaluation matches gradients of + individual full evaluations using PyTorch""" + import torch + import torch.autograd.functional as F + + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev, interface="torch", diff_method=diff_method) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial argument to construct a new circuit with + y = torch.tensor(np.random.uniform(size=2), requires_grad=True) + + # the batched argument to the new partial circuit + x = torch.tensor(np.random.uniform(size=batch_size), requires_grad=True) + + batched_partial_circuit = qml.batch_partial(circuit, y=y) + + # we could also sum over the batch dimension and use the regular + # gradient instead of the jacobian, but either works + grad = torch.diagonal(F.jacobian(batched_partial_circuit, x)) + + # check the results against individually executed circuits + indiv_grad = [] + for x_indiv in x: + # create a new variable since tensors created by + # indexing a trainable variable aren't themselves trainable + x_indiv = x_indiv.clone().detach().requires_grad_(True) + + out_indiv = circuit(x_indiv, y) + out_indiv.backward() + + indiv_grad.append(x_indiv.grad) + + assert np.allclose(grad, indiv_grad) + + def test_lambda_evaluation(): """Test lambda argument replacement matches individual full evaluations""" dev = qml.device("default.qubit", wires=2) From d953e7c228307eb55a3890418765e2602d0cee00 Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Wed, 18 May 2022 09:24:53 -0400 Subject: [PATCH 04/26] Fix tf test --- tests/transforms/test_batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index 8e1eb844fd1..0cd242061f1 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -191,7 +191,7 @@ def circuit(x, y): # check the results against individually executed circuits indiv_grad = [] - for x_indiv in x: + for x_indiv in tf.unstack(x): # create a new variable since tensors created by # indexing a trainable variable aren't themselves trainable x_indiv = tf.Variable(x_indiv, trainable=True) From 7117fcc467412e9a8fc404e66c7fb23a09ce9994 Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Wed, 18 May 2022 14:11:00 -0400 Subject: [PATCH 05/26] Add docs and a couple more tests --- pennylane/transforms/batch_partial.py | 67 +++++++- tests/transforms/test_batch_partial.py | 220 ++++++++++++++++++++++++- 2 files changed, 277 insertions(+), 10 deletions(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index 0098e6def41..4f0cf59936e 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -41,7 +41,61 @@ def _convert_to_args(func, args, kwargs): def batch_partial(qnode, all_operations=False, **partial_kwargs): """ - TODO: docs + Create a wrapper function around the QNode with partially + evaluated parameters, which supports an initial batch dimension + for other unevaluated parameters. + + Args: + qnode (pennylane.QNode): QNode to partially evaluate + all_operations (bool): If ``True``, a batch dimension will be added to *all* operations + in the QNode, rather than just trainable QNode parameters. + partial_kwargs (dict): partially-evaluated parameters to pass to the QNode + + Returns: + func: Function which accepts the same arguments as the QNode minus the + partially evaluated arguments provided, and behaves the same as the QNode + called with both the partially evaluated arguments and the extra arguments. + However, the first dimension of each argument of the returned function + will be treated as a batch dimension. The function output will also contain + an initial batch dimension. + + **Example** + + Consider the following circuit: + + .. code-block:: python + + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + The ``qml.batch_partial`` decorator allows us to create a partially evaluated + function that wraps the QNode. For example, + + >>> y = np.array([0.2, 0.3]) + >>> batched_partial_circuit = qml.batch_partial(circuit, y=y) + + The unevaluated arguments of the resulting function must now have a batch + dimension, and the output of the function also has a batch dimension: + + >>> batch_size = 4 + >>> x = np.linspace(0.1, 0.5, batch_size) + >>> batched_partial_circuit(x) + tensor([0.9316158 , 0.91092081, 0.87405565, 0.82167473], requires_grad=True) + + Gradients can be computed for the arguments of the wrapper function, but + not for any partially evaluated arguments passed to ``qml.batch_partial``: + + >>> qml.jacobian(batched_partial_circuit)(x) + array([[-0.09347337, 0. , 0. , 0. ], + [ 0. , -0.21649144, 0. , 0. ], + [ 0. , 0. , -0.33566648, 0. ], + [ 0. , 0. , 0. , -0.44888295]]) """ qnode = qml.batch_params(qnode, all_operations=all_operations) @@ -78,10 +132,13 @@ def wrapper(*args, **kwargs): # get the batch dimension (we don't have to check if all arguments # have the same batch dim since that's done in qml.batch_params) - if args: - batch_dim = qml.math.shape(args[0])[0] - else: - batch_dim = qml.math.shape(list(kwargs.values())[0])[0] + try: + if args: + batch_dim = qml.math.shape(args[0])[0] + else: + batch_dim = qml.math.shape(list(kwargs.values())[0])[0] + except IndexError: + raise ValueError("Batch dimension must be provided") from None for key, val in partial_kwargs.items(): if callable(val): diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index 0cd242061f1..d32413bc6c4 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -14,6 +14,7 @@ """ Unit tests for the batch partial transform. """ +import re import pytest import pennylane as qml @@ -137,10 +138,10 @@ def circuit(x, y): batch_size = 4 # the partial argument to construct a new circuit with - y = jnp.array([0.3, 0.4]) + y = jnp.asarray(np.random.uniform(size=2)) # the batched argument to the new partial circuit - x = jnp.linspace(0.1, 0.5, batch_size) + x = jnp.asarray(np.random.uniform(size=batch_size)) batched_partial_circuit = qml.batch_partial(circuit, all_operations=True, y=y) @@ -285,12 +286,14 @@ def circuit(x, y): assert np.allclose(res, indiv_res) -def test_lambda_evaluation_grad(): +@pytest.mark.autograd +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +def test_lambda_evaluation_autograd(diff_method): """Test gradient of lambda argument replacement matches - gradients of individual full evaluations""" + gradients of individual full evaluations using autograd""" dev = qml.device("default.qubit", wires=2) - @qml.qnode(dev) + @qml.qnode(dev, diff_method=diff_method) def circuit(x, y): qml.RX(x, wires=0) qml.RY(y[..., 0], wires=0) @@ -327,6 +330,159 @@ def circuit(x, y): assert np.allclose(grad, indiv_grad) +@pytest.mark.jax +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +def test_lambda_evaluation_jax(diff_method): + """Test gradient of lambda argument replacement matches + gradients of individual full evaluations using JAX""" + import jax + import jax.numpy as jnp + + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev, interface="jax", diff_method=diff_method) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the first partial argument + x = jnp.asarray(np.random.uniform(size=())) + + # the base value of the second partial argument + y = jnp.asarray(np.random.uniform(size=2)) + + # the second partial argument as a function of the inputs + fn = lambda y0: y + y0 * jnp.ones(2) + + # values for the second argument + y0 = jnp.asarray(np.random.uniform(size=batch_size)) + + batched_partial_circuit = qml.batch_partial(circuit, all_operations=True, x=x, y=fn) + + # we could also sum over the batch dimension and use the regular + # gradient instead of the jacobian, but either works + grad = jnp.diagonal(jax.jacrev(batched_partial_circuit)(y0)) + + # check the results against individually executed circuits + indiv_grad = [] + for y0_indiv in y0: + grad_wrt_second_arg = jax.grad(circuit, argnums=1)(x, y + y0_indiv * np.ones(2)) + grad_wrt_y0 = jnp.sum(grad_wrt_second_arg) + indiv_grad.append(grad_wrt_y0) + + assert np.allclose(grad, indiv_grad) + + +@pytest.mark.tf +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +def test_lambda_evaluation_tf(diff_method): + """Test gradient of lambda argument replacement matches + gradients of individual full evaluations using TF""" + import tensorflow as tf + + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev, interface="tf", diff_method=diff_method) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the first partial argument + x = tf.Variable(np.random.uniform(size=())) + + # the base value of the second partial argument + y = tf.Variable(np.random.uniform(size=2)) + + # the second partial argument as a function of the inputs + fn = lambda y0: y + y0 * tf.ones(2, dtype=tf.float64) + + # values for the second argument + y0 = tf.Variable(np.random.uniform(size=batch_size), trainable=True) + + batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + + with tf.GradientTape() as tape: + out = batched_partial_circuit(y0) + + # we could also sum over the batch dimension and use the regular + # gradient instead of the jacobian, but either works + grad = tf.linalg.tensor_diag_part(tape.jacobian(out, y0)) + + # check the results against individually executed circuits + indiv_grad = [] + for y0_indiv in tf.unstack(y0): + # create a new variable since tensors created by + # indexing a trainable variable aren't themselves trainable + y0_indiv = tf.Variable(y0_indiv, trainable=True) + + with tf.GradientTape() as tape: + out_indiv = circuit(x, y + y0_indiv * tf.ones(2, dtype=tf.float64)) + + indiv_grad.append(tape.gradient(out_indiv, y0_indiv)) + + assert np.allclose(grad, indiv_grad) + + +@pytest.mark.torch +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +def test_lambda_evaluation_torch(diff_method): + """Test gradient of lambda argument replacement matches + gradients of individual full evaluations using PyTorch""" + import torch + import torch.autograd.functional as F + + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev, interface="torch", diff_method=diff_method) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the first partial argument + x = torch.tensor(np.random.uniform(size=()), requires_grad=True) + + # the base value of the second partial argument + y = torch.tensor(np.random.uniform(size=2), requires_grad=True) + + # the second partial argument as a function of the inputs + fn = lambda y0: y + y0 * torch.ones(2) + + # values for the second argument + y0 = torch.tensor(np.random.uniform(size=batch_size), requires_grad=True) + + batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + + # we could also sum over the batch dimension and use the regular + # gradient instead of the jacobian, but either works + grad = torch.diagonal(F.jacobian(batched_partial_circuit, y0)) + + # check the results against individually executed circuits + indiv_grad = [] + for y0_indiv in y0: + # create a new variable since tensors created by + # indexing a trainable variable aren't themselves trainable + y0_indiv = y0_indiv.clone().detach().requires_grad_(True) + + out_indiv = circuit(x, y + y0_indiv * torch.ones(2)) + out_indiv.backward() + + indiv_grad.append(y0_indiv.grad) + + assert np.allclose(grad, indiv_grad) + + def test_full_evaluation_error(): """Test that an error is raised when all arguments to QNode are provided to a partial evaluation.""" @@ -412,3 +568,57 @@ def circuit(x, y): match="Arguments must not be passed as keyword arguments to callable within partial function", ): res = batched_partial_circuit(y0=y0) + + +def test_no_batchdim_error(): + """Test that an error is raised when no batch + dimension is given to the decorated QNode""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + # the second partial argument + y = np.random.uniform(size=2) + + # the incorrectly batched argument to the new partial circuit + x = np.random.uniform(size=()) + + batched_partial_circuit = qml.batch_partial(circuit, y=y) + + with pytest.raises(ValueError, match="Batch dimension must be provided"): + out = batched_partial_circuit(x=x) + + +def test_different_batchdim_error(): + """Test that an error is raised when different batch + dimensions are given to the decorated QNode""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y, z): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + qml.RX(z, wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size1, batch_size2 = 5, 4 + + # the second partial argument + y = np.random.uniform(size=2) + + # the batched arguments to the new partial circuit + x = np.random.uniform(size=batch_size1) + z = np.random.uniform(size=batch_size2) + + batched_partial_circuit = qml.batch_partial(circuit, y=y) + + msg = "has incorrect batch dimension. Expecting first dimension of length 5." + msg = re.escape(msg) + with pytest.raises(ValueError, match=msg): + out = batched_partial_circuit(x=x, z=z) From 90303a6131852022c1e89a9a614388848ba35601 Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Wed, 18 May 2022 15:28:43 -0400 Subject: [PATCH 06/26] Improve docs and fix broken tests --- pennylane/transforms/batch_partial.py | 29 ++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index 4f0cf59936e..c9a25521c44 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -42,8 +42,7 @@ def _convert_to_args(func, args, kwargs): def batch_partial(qnode, all_operations=False, **partial_kwargs): """ Create a wrapper function around the QNode with partially - evaluated parameters, which supports an initial batch dimension - for other unevaluated parameters. + evaluated parameters, which supports an initial batch dimension. Args: qnode (pennylane.QNode): QNode to partially evaluate @@ -96,6 +95,29 @@ def circuit(x, y): [ 0. , -0.21649144, 0. , 0. ], [ 0. , 0. , -0.33566648, 0. ], [ 0. , 0. , 0. , -0.44888295]]) + + The same ``qml.batch_partial`` decorator can also be used to replace arguments + of a QNode with functions, and calling the wrapper would evaluate + those functions and pass the results into the QNode. For example, + + >>> x = np.array(0.1) + >>> y_fn = lambda y0: y0 * np.array([0.2, 0.3]) + >>> batched_lambda_circuit = qml.batch_partial(circuit, x=x, y=y_fn) + + The wrapped function ``batched_lambda_circuit`` also expects arguments to + have an initial batch dimension: + + >>> batch_size = 4 + >>> y0 = np.linspace(0.5, 2, batch_size) + >>> batched_lambda_circuit(y0) + tensor([0.97891628, 0.9316158 , 0.85593241, 0.75638669], requires_grad=True) + + Gradients can be computed in this scenario as well: + >>> qml.jacobian(batched_lambda_circuit)(y0) + array([[-0.06402847, 0. , 0. , 0. ], + [ 0. , -0.12422434, 0. , 0. ], + [ 0. , 0. , -0.17699293, 0. ], + [ 0. , 0. , 0. , -0.21920062]]) """ qnode = qml.batch_params(qnode, all_operations=all_operations) @@ -142,7 +164,8 @@ def wrapper(*args, **kwargs): for key, val in partial_kwargs.items(): if callable(val): - val = qml.math.stack([val(*a) for a in zip(*args)]) + unstacked_args = (qml.math.unstack(arg) for arg in args) + val = qml.math.stack([val(*a) for a in zip(*unstacked_args)]) kwargs[key] = val else: kwargs[key] = qml.math.stack([val] * batch_dim) From abb5dabb895c41a1a05a6e8fc5f29fd5ae417fd6 Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Wed, 18 May 2022 17:04:55 -0400 Subject: [PATCH 07/26] Fix indent error --- pennylane/transforms/batch_partial.py | 1 + 1 file changed, 1 insertion(+) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index c9a25521c44..b5b4c7855d0 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -113,6 +113,7 @@ def circuit(x, y): tensor([0.97891628, 0.9316158 , 0.85593241, 0.75638669], requires_grad=True) Gradients can be computed in this scenario as well: + >>> qml.jacobian(batched_lambda_circuit)(y0) array([[-0.06402847, 0. , 0. , 0. ], [ 0. , -0.12422434, 0. , 0. ], From b669d45aec86402c8f8d186cc1c8b2cfaf51ee62 Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Thu, 19 May 2022 10:12:24 -0400 Subject: [PATCH 08/26] Add changelog entry --- doc/releases/changelog-dev.md | 25 +++++++++++++++++++++++-- 1 file changed, 23 insertions(+), 2 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index c7f656c6cd8..b5be04930ee 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -7,7 +7,7 @@ * Boolean mask indexing of the parameter-shift Hessian [(#2538)](https://github.com/PennyLaneAI/pennylane/pull/2538) - The `argnum` keyword argument for `param_shift_hessian` + The `argnum` keyword argument for `param_shift_hessian` is now allowed to be a twodimensional Boolean `array_like`. Only the indicated entries of the Hessian will then be computed. A particularly useful example is the computation of the diagonal @@ -37,6 +37,27 @@ The code that checks for qubit wise commuting (QWC) got a performance boost that is noticable when many commuting paulis of the same type are measured. +* Added new transform `qml.batch_partial` which behaves similarly to `functools.partial` but supports batching in the unevaluated parameters. + [(#2585)](https://github.com/PennyLaneAI/pennylane/pull/2585) + + This is useful for batching circuit executions with some identical parameters but not others: + + ```python + dev = qml.device("default.qubit", wires=1) + + @qml.qnode(dev) + def circuit(x, y): + qml.RX(x, wires=0) + qml.RY(y, wires=0) + return qml.expval(qml.PauliZ(wires=0)) + ``` + ```pycon + >>> batched_partial_circuit = qml.batch_partial(circuit, x=np.array(np.pi / 2)) + >>> y = np.array([0.2, 0.3, 0.4]) + >>> batched_partial_circuit(y=y) + tensor([0.69301172, 0.67552491, 0.65128847], requires_grad=True) + ``` +

Improvements

* The developer-facing `pow` method has been added to `Operator` with concrete implementations @@ -111,7 +132,7 @@

Bug fixes

-* `QNode`'s now can interpret variations on the interface name, like `"tensorflow"` or `"jax-jit"`, when requesting backpropagation. +* `QNode`'s now can interpret variations on the interface name, like `"tensorflow"` or `"jax-jit"`, when requesting backpropagation. [(#2591)](https://github.com/PennyLaneAI/pennylane/pull/2591) * Fixed a bug for `diff_method="adjoint"` where incorrect gradients were From 24fab51a239caf7495925e2b7b28b76a3451169c Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 09:33:04 -0400 Subject: [PATCH 09/26] Update pennylane/transforms/batch_partial.py Co-authored-by: David Wierichs --- pennylane/transforms/batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index b5b4c7855d0..d1bdab7c61c 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -112,7 +112,7 @@ def circuit(x, y): >>> batched_lambda_circuit(y0) tensor([0.97891628, 0.9316158 , 0.85593241, 0.75638669], requires_grad=True) - Gradients can be computed in this scenario as well: + Jacobians can be computed in this scenario as well: >>> qml.jacobian(batched_lambda_circuit)(y0) array([[-0.06402847, 0. , 0. , 0. ], From 26f207961f4388499aebfc5295da33bbbd402bb1 Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 09:33:24 -0400 Subject: [PATCH 10/26] Update pennylane/transforms/batch_partial.py Co-authored-by: David Wierichs --- pennylane/transforms/batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index d1bdab7c61c..e069fd581f7 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -87,7 +87,7 @@ def circuit(x, y): >>> batched_partial_circuit(x) tensor([0.9316158 , 0.91092081, 0.87405565, 0.82167473], requires_grad=True) - Gradients can be computed for the arguments of the wrapper function, but + Jacobians can be computed for the arguments of the wrapper function, but not for any partially evaluated arguments passed to ``qml.batch_partial``: >>> qml.jacobian(batched_partial_circuit)(x) From fd76986911fd30b83c3ff826ff471897cb6ee804 Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 09:34:28 -0400 Subject: [PATCH 11/26] Update pennylane/transforms/batch_partial.py Co-authored-by: David Wierichs --- pennylane/transforms/batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index e069fd581f7..7ea226cd1cc 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -161,7 +161,7 @@ def wrapper(*args, **kwargs): else: batch_dim = qml.math.shape(list(kwargs.values())[0])[0] except IndexError: - raise ValueError("Batch dimension must be provided") from None + raise ValueError("Parameter with batch dimension must be provided") from None for key, val in partial_kwargs.items(): if callable(val): From b826dc014f1dad02cdf71ae5f6328e998c0769e1 Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 09:37:49 -0400 Subject: [PATCH 12/26] Update doc/releases/changelog-dev.md Co-authored-by: David Wierichs --- doc/releases/changelog-dev.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index b5be04930ee..f08f2ab5d28 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -40,7 +40,7 @@ * Added new transform `qml.batch_partial` which behaves similarly to `functools.partial` but supports batching in the unevaluated parameters. [(#2585)](https://github.com/PennyLaneAI/pennylane/pull/2585) - This is useful for batching circuit executions with some identical parameters but not others: + This is useful for executing a circuit with a batch dimension in some of its parameters: ```python dev = qml.device("default.qubit", wires=1) From 7aaf29ec13c2f3c1da29e59650421c96acd31d6d Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Fri, 20 May 2022 09:44:45 -0400 Subject: [PATCH 13/26] Change example in docstring --- pennylane/transforms/batch_partial.py | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index 7ea226cd1cc..1d1f456674c 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -69,14 +69,13 @@ def batch_partial(qnode, all_operations=False, **partial_kwargs): @qml.qnode(dev) def circuit(x, y): qml.RX(x, wires=0) - qml.RY(y[..., 0], wires=0) - qml.RY(y[..., 1], wires=1) + qml.RY(y, wires=1) return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) The ``qml.batch_partial`` decorator allows us to create a partially evaluated function that wraps the QNode. For example, - >>> y = np.array([0.2, 0.3]) + >>> y = np.array(0.2) >>> batched_partial_circuit = qml.batch_partial(circuit, y=y) The unevaluated arguments of the resulting function must now have a batch @@ -85,23 +84,23 @@ def circuit(x, y): >>> batch_size = 4 >>> x = np.linspace(0.1, 0.5, batch_size) >>> batched_partial_circuit(x) - tensor([0.9316158 , 0.91092081, 0.87405565, 0.82167473], requires_grad=True) + tensor([0.97517033, 0.95350781, 0.91491915, 0.86008934], requires_grad=True) Jacobians can be computed for the arguments of the wrapper function, but not for any partially evaluated arguments passed to ``qml.batch_partial``: >>> qml.jacobian(batched_partial_circuit)(x) - array([[-0.09347337, 0. , 0. , 0. ], - [ 0. , -0.21649144, 0. , 0. ], - [ 0. , 0. , -0.33566648, 0. ], - [ 0. , 0. , 0. , -0.44888295]]) + array([[-0.0978434 , 0. , 0. , 0. ], + [ 0. , -0.22661276, 0. , 0. ], + [ 0. , 0. , -0.35135943, 0. ], + [ 0. , 0. , 0. , -0.46986895]]) The same ``qml.batch_partial`` decorator can also be used to replace arguments of a QNode with functions, and calling the wrapper would evaluate those functions and pass the results into the QNode. For example, >>> x = np.array(0.1) - >>> y_fn = lambda y0: y0 * np.array([0.2, 0.3]) + >>> y_fn = lambda y0: y0 * 0.2 + 0.3 >>> batched_lambda_circuit = qml.batch_partial(circuit, x=x, y=y_fn) The wrapped function ``batched_lambda_circuit`` also expects arguments to @@ -110,15 +109,15 @@ def circuit(x, y): >>> batch_size = 4 >>> y0 = np.linspace(0.5, 2, batch_size) >>> batched_lambda_circuit(y0) - tensor([0.97891628, 0.9316158 , 0.85593241, 0.75638669], requires_grad=True) + tensor([0.91645953, 0.8731983 , 0.82121237, 0.76102116], requires_grad=True) Jacobians can be computed in this scenario as well: >>> qml.jacobian(batched_lambda_circuit)(y0) - array([[-0.06402847, 0. , 0. , 0. ], - [ 0. , -0.12422434, 0. , 0. ], - [ 0. , 0. , -0.17699293, 0. ], - [ 0. , 0. , 0. , -0.21920062]]) + array([[-0.07749457, 0. , 0. , 0. ], + [ 0. , -0.09540608, 0. , 0. ], + [ 0. , 0. , -0.11236432, 0. ], + [ 0. , 0. , 0. , -0.12819986]]) """ qnode = qml.batch_params(qnode, all_operations=all_operations) From b5cfe4bc43f64d65db2b9c739eb1cff353a9c5de Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Fri, 20 May 2022 10:00:52 -0400 Subject: [PATCH 14/26] Change test to match new error message --- tests/transforms/test_batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index d32413bc6c4..9b0e1f7b0be 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -590,7 +590,7 @@ def circuit(x, y): batched_partial_circuit = qml.batch_partial(circuit, y=y) - with pytest.raises(ValueError, match="Batch dimension must be provided"): + with pytest.raises(ValueError, match="Parameter with batch dimension must be provided"): out = batched_partial_circuit(x=x) From 797160ecac3ec2a12e1480987276fa33e5a4cd60 Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 16:16:22 -0400 Subject: [PATCH 15/26] Update tests/transforms/test_batch_partial.py Co-authored-by: antalszava --- tests/transforms/test_batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index 9b0e1f7b0be..2a0dc94b49e 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -205,7 +205,7 @@ def circuit(x, y): @pytest.mark.torch -@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_partial_evaluation_torch(diff_method): """Test gradient of partial evaluation matches gradients of individual full evaluations using PyTorch""" From 365a63c74208e4c501a42a5e6a1d17b769d50a98 Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 16:16:29 -0400 Subject: [PATCH 16/26] Update tests/transforms/test_batch_partial.py Co-authored-by: antalszava --- tests/transforms/test_batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index 2a0dc94b49e..d95974c6f9e 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -158,7 +158,7 @@ def circuit(x, y): @pytest.mark.tf -@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_partial_evaluation_tf(diff_method): """Test gradient of partial evaluation matches gradients of individual full evaluations using TF""" From 5434172330949da74a09d63eab0b1267fe55d5f4 Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 16:16:36 -0400 Subject: [PATCH 17/26] Update tests/transforms/test_batch_partial.py Co-authored-by: antalszava --- tests/transforms/test_batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index d95974c6f9e..a3229892d2b 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -119,7 +119,7 @@ def circuit(x, y): @pytest.mark.jax -@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_partial_evaluation_jax(diff_method): """Test gradient of partial evaluation matches gradients of individual full evaluations using jax""" From 037b673aa0cc7559bf538ef6aece785a119750cc Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 16:16:41 -0400 Subject: [PATCH 18/26] Update tests/transforms/test_batch_partial.py Co-authored-by: antalszava --- tests/transforms/test_batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index a3229892d2b..1ae143ea237 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -83,7 +83,7 @@ def circuit(x, y): @pytest.mark.autograd -@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_partial_evaluation_autograd(diff_method): """Test gradient of partial evaluation matches gradients of individual full evaluations using autograd""" From ed1bd584acaddffe0fd3ad769660afbd5f2cc120 Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 16:17:22 -0400 Subject: [PATCH 19/26] Update tests/transforms/test_batch_partial.py Co-authored-by: antalszava --- tests/transforms/test_batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index 1ae143ea237..9a1515d10bf 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -53,7 +53,7 @@ def circuit(x, y): def test_partial_evaluation_kwargs(): """Test partial evaluation matches individual full evaluations - when kwargs are used""" + when the keyword syntax is used to call the partial object""" dev = qml.device("default.qubit", wires=2) @qml.qnode(dev) From 111b22e5d32f0a1cd0ba130f35f2c0dbc59fd22e Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Fri, 20 May 2022 16:18:26 -0400 Subject: [PATCH 20/26] Update pennylane/transforms/batch_partial.py Co-authored-by: antalszava --- pennylane/transforms/batch_partial.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index 1d1f456674c..1e791048558 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -95,7 +95,7 @@ def circuit(x, y): [ 0. , 0. , -0.35135943, 0. ], [ 0. , 0. , 0. , -0.46986895]]) - The same ``qml.batch_partial`` decorator can also be used to replace arguments + The same ``qml.batch_partial`` function can also be used to replace arguments of a QNode with functions, and calling the wrapper would evaluate those functions and pass the results into the QNode. For example, From 5d2b11dfe30d9422df09ea82674a55030631f3d8 Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Sun, 22 May 2022 17:34:36 -0400 Subject: [PATCH 21/26] Update pennylane/transforms/batch_partial.py Co-authored-by: antalszava --- pennylane/transforms/batch_partial.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index 1e791048558..77b4d9d75dd 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -41,8 +41,9 @@ def _convert_to_args(func, args, kwargs): def batch_partial(qnode, all_operations=False, **partial_kwargs): """ - Create a wrapper function around the QNode with partially - evaluated parameters, which supports an initial batch dimension. + Create a batched partial callable object from the QNode specified. + + This transform provides functionality akin to `functools.partial` and allows batching the arguments used for calling the batched partial object. Args: qnode (pennylane.QNode): QNode to partially evaluate From 04a99edb934726324e7e3168c6f48d63a9636c61 Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Tue, 24 May 2022 12:46:30 -0400 Subject: [PATCH 22/26] Changes for review --- pennylane/transforms/batch_partial.py | 71 +++++++++------ tests/transforms/test_batch_partial.py | 121 ++++++++++++++++++++++--- 2 files changed, 152 insertions(+), 40 deletions(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index 77b4d9d75dd..fa62ca8446e 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -12,50 +12,50 @@ # See the License for the specific language governing permissions and # limitations under the License. """ -Contains the batch dimension transform. +Contains the batch dimension transform for partial use of QNodes. """ import functools import inspect import pennylane as qml +from pennylane import numpy as np -def _convert_to_args(func, args, kwargs): +def _convert_to_args(sig, args, kwargs): """ - Given a function, convert the positional and + Given the signature of a function, convert the positional and keyword arguments to purely positional arguments. """ - sig = inspect.signature(func).parameters - new_args = [] for i, param in enumerate(sig): if param in kwargs: - # first check if the name is provided in kwargs + # first check if the name is provided in the keyword arguments new_args.append(kwargs[param]) - elif i < len(sig): - # next check if the argnum is provided + else: + # if not, then the argument must be positional new_args.append(args[i]) return tuple(new_args) -def batch_partial(qnode, all_operations=False, **partial_kwargs): +def batch_partial(qnode, all_operations=False, preprocess=None, **partial_kwargs): """ Create a batched partial callable object from the QNode specified. - - This transform provides functionality akin to `functools.partial` and allows batching the arguments used for calling the batched partial object. + + This transform provides functionality akin to `functools.partial` and + allows batching the arguments used for calling the batched partial object. Args: - qnode (pennylane.QNode): QNode to partially evaluate + qnode (pennylane.QNode): QNode to pre-supply arguments to all_operations (bool): If ``True``, a batch dimension will be added to *all* operations in the QNode, rather than just trainable QNode parameters. - partial_kwargs (dict): partially-evaluated parameters to pass to the QNode + partial_kwargs (dict): pre-supplied arguments to pass to the QNode. Returns: - func: Function which accepts the same arguments as the QNode minus the - partially evaluated arguments provided, and behaves the same as the QNode - called with both the partially evaluated arguments and the extra arguments. - However, the first dimension of each argument of the returned function + func: Function which wraps the QNode and accepts the same arguments minus the + pre-supplied arguments provided, and behaves the same as the QNode called with + both the pre-supplied arguments and the other arguments passed to this wrapper + function. However, the first dimension of each argument of the wrapper function will be treated as a batch dimension. The function output will also contain an initial batch dimension. @@ -122,13 +122,23 @@ def circuit(x, y): """ qnode = qml.batch_params(qnode, all_operations=all_operations) + preprocess = {} if preprocess is None else preprocess + # store whether this decorator is being used as a pure # analog of functools.partial, or whether it is used to # wrap a QNode in a more complex lambda statement - is_partial = False - if not any(callable(val) for val in partial_kwargs.values()): - # none of the kwargs passed in are callable - is_partial = True + is_partial = preprocess == {} + + # determine which arguments need to be stacked along the batch dimension + to_stack = [] + for key, val in partial_kwargs.items(): + try: + # check if the value is a tensor + if qml.math.asarray(val).dtype != object: + to_stack.append(key) + except ImportError: + # autoray can't find a backend for val, so it cannot be stacked + pass sig = inspect.signature(qnode).parameters if is_partial: @@ -139,7 +149,7 @@ def circuit(x, y): else: # if used to wrap a QNode in a lambda statement, then check that # all arguments are provided - if len(sig) > len(partial_kwargs): + if len(sig) > len(partial_kwargs) + len(preprocess): raise ValueError("Callable argument requires all other arguments to QNode be provided") @functools.wraps(qnode) @@ -163,18 +173,21 @@ def wrapper(*args, **kwargs): except IndexError: raise ValueError("Parameter with batch dimension must be provided") from None + for key, val in preprocess.items(): + unstacked_args = (qml.math.unstack(arg) for arg in args) + val = qml.math.stack([val(*a) for a in zip(*unstacked_args)]) + kwargs[key] = val + for key, val in partial_kwargs.items(): - if callable(val): - unstacked_args = (qml.math.unstack(arg) for arg in args) - val = qml.math.stack([val(*a) for a in zip(*unstacked_args)]) - kwargs[key] = val - else: + if key in to_stack: kwargs[key] = qml.math.stack([val] * batch_dim) + else: + kwargs[key] = val if is_partial: - return qnode(*_convert_to_args(qnode, args, kwargs)) + return qnode(*_convert_to_args(sig, args, kwargs)) # don't pass the arguments to the lambda itself into the QNode - return qnode(*_convert_to_args(qnode, (), kwargs)) + return qnode(*_convert_to_args(sig, (), kwargs)) return wrapper diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index 9a1515d10bf..1879c8a55e8 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -82,6 +82,103 @@ def circuit(x, y): assert np.allclose(res, indiv_res) +def test_partial_evaluation_multi_args(): + """Test partial evaluation matches individual full evaluations + for multiple pre-supplied arguments""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y, z): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + qml.RX(z, wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial arguments to construct a new circuit with + y = np.random.uniform(size=2) + z = np.random.uniform(size=()) + + # the batched argument to the new partial circuit + x = np.random.uniform(size=batch_size) + + batched_partial_circuit = qml.batch_partial(circuit, y=y, z=z) + res = batched_partial_circuit(x) + + # check the results against individually executed circuits + indiv_res = [] + for x_indiv in x: + indiv_res.append(circuit(x_indiv, y, z)) + + assert np.allclose(res, indiv_res) + + +def test_partial_evaluation_nonnumeric1(): + """Test partial evaluation matches individual full evaluations + for non-numeric pre-supplied arguments""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y, measurement): + qml.RX(x, wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.apply(measurement) + + batch_size = 4 + + # the partial arguments to construct a new circuit with + y = np.random.uniform(size=2) + measurement = qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + # the batched argument to the new partial circuit + x = np.random.uniform(size=batch_size) + + batched_partial_circuit = qml.batch_partial(circuit, y=y, measurement=measurement) + res = batched_partial_circuit(x) + + # check the results against individually executed circuits + indiv_res = [] + for x_indiv in x: + indiv_res.append(circuit(x_indiv, y, measurement)) + + assert np.allclose(res, indiv_res) + + +def test_partial_evaluation_nonnumeric2(): + """Test partial evaluation matches individual full evaluations + for non-numeric pre-supplied arguments""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y, func): + qml.RX(func(x), wires=0) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial arguments to construct a new circuit with + y = np.random.uniform(size=2) + func = np.cos + + # the batched argument to the new partial circuit + x = np.random.uniform(size=batch_size) + + batched_partial_circuit = qml.batch_partial(circuit, y=y, func=func) + res = batched_partial_circuit(x) + + # check the results against individually executed circuits + indiv_res = [] + for x_indiv in x: + indiv_res.append(circuit(x_indiv, y, func)) + + assert np.allclose(res, indiv_res) + + @pytest.mark.autograd @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_partial_evaluation_autograd(diff_method): @@ -275,7 +372,7 @@ def circuit(x, y): # values for the second argument y0 = np.random.uniform(size=batch_size) - batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + batched_partial_circuit = qml.batch_partial(circuit, x=x, preprocess={"y": fn}) res = batched_partial_circuit(y0) # check the results against individually executed circuits @@ -287,7 +384,7 @@ def circuit(x, y): @pytest.mark.autograd -@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_lambda_evaluation_autograd(diff_method): """Test gradient of lambda argument replacement matches gradients of individual full evaluations using autograd""" @@ -314,7 +411,7 @@ def circuit(x, y): # values for the second argument y0 = np.random.uniform(size=batch_size, requires_grad=True) - batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + batched_partial_circuit = qml.batch_partial(circuit, x=x, preprocess={"y": fn}) # we could also sum over the batch dimension and use the regular # gradient instead of the jacobian, but either works @@ -331,7 +428,7 @@ def circuit(x, y): @pytest.mark.jax -@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_lambda_evaluation_jax(diff_method): """Test gradient of lambda argument replacement matches gradients of individual full evaluations using JAX""" @@ -361,7 +458,9 @@ def circuit(x, y): # values for the second argument y0 = jnp.asarray(np.random.uniform(size=batch_size)) - batched_partial_circuit = qml.batch_partial(circuit, all_operations=True, x=x, y=fn) + batched_partial_circuit = qml.batch_partial( + circuit, all_operations=True, x=x, preprocess={"y": fn} + ) # we could also sum over the batch dimension and use the regular # gradient instead of the jacobian, but either works @@ -378,7 +477,7 @@ def circuit(x, y): @pytest.mark.tf -@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_lambda_evaluation_tf(diff_method): """Test gradient of lambda argument replacement matches gradients of individual full evaluations using TF""" @@ -407,7 +506,7 @@ def circuit(x, y): # values for the second argument y0 = tf.Variable(np.random.uniform(size=batch_size), trainable=True) - batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + batched_partial_circuit = qml.batch_partial(circuit, x=x, preprocess={"y": fn}) with tf.GradientTape() as tape: out = batched_partial_circuit(y0) @@ -432,7 +531,7 @@ def circuit(x, y): @pytest.mark.torch -@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift"]) +@pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_lambda_evaluation_torch(diff_method): """Test gradient of lambda argument replacement matches gradients of individual full evaluations using PyTorch""" @@ -462,7 +561,7 @@ def circuit(x, y): # values for the second argument y0 = torch.tensor(np.random.uniform(size=batch_size), requires_grad=True) - batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + batched_partial_circuit = qml.batch_partial(circuit, x=x, preprocess={"y": fn}) # we could also sum over the batch dimension and use the regular # gradient instead of the jacobian, but either works @@ -531,7 +630,7 @@ def circuit(x, y): with pytest.raises( ValueError, match="Callable argument requires all other arguments to QNode be provided" ): - batched_partial_circuit = qml.batch_partial(circuit, y=fn) + batched_partial_circuit = qml.batch_partial(circuit, preprocess={"y": fn}) def test_kwargs_callable_error(): @@ -555,7 +654,7 @@ def circuit(x, y): fn = lambda y0: y + y0 * np.ones(2) y0 = np.random.uniform(size=batch_size) - batched_partial_circuit = qml.batch_partial(circuit, x=x, y=fn) + batched_partial_circuit = qml.batch_partial(circuit, x=x, preprocess={"y": fn}) with pytest.raises( ValueError, From cce077b278bd25699853d65d8f8ba7663cc7c48a Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Tue, 24 May 2022 13:40:57 -0400 Subject: [PATCH 23/26] Add test for coverage --- tests/transforms/test_batch_partial.py | 32 ++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/tests/transforms/test_batch_partial.py b/tests/transforms/test_batch_partial.py index 1879c8a55e8..f4a3dba440c 100644 --- a/tests/transforms/test_batch_partial.py +++ b/tests/transforms/test_batch_partial.py @@ -179,6 +179,38 @@ def circuit(x, y, func): assert np.allclose(res, indiv_res) +def test_partial_evaluation_nonnumeric3(): + """Test partial evaluation matches individual full evaluations + for non-numeric pre-supplied arguments""" + dev = qml.device("default.qubit", wires=2) + + @qml.qnode(dev) + def circuit(x, y, op): + qml.apply(op(x, wires=0)) + qml.RY(y[..., 0], wires=0) + qml.RY(y[..., 1], wires=1) + return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) + + batch_size = 4 + + # the partial arguments to construct a new circuit with + y = np.random.uniform(size=2) + op = qml.RX + + # the batched argument to the new partial circuit + x = np.random.uniform(size=batch_size) + + batched_partial_circuit = qml.batch_partial(circuit, y=y, op=op) + res = batched_partial_circuit(x) + + # check the results against individually executed circuits + indiv_res = [] + for x_indiv in x: + indiv_res.append(circuit(x_indiv, y, op)) + + assert np.allclose(res, indiv_res) + + @pytest.mark.autograd @pytest.mark.parametrize("diff_method", ["backprop", "adjoint", "parameter-shift", "finite-diff"]) def test_partial_evaluation_autograd(diff_method): From f127987823ae4ba50cd48906913e2bd79d6d0e2a Mon Sep 17 00:00:00 2001 From: Edward Jiang <34989448+eddddddy@users.noreply.github.com> Date: Thu, 26 May 2022 09:40:46 -0400 Subject: [PATCH 24/26] Apply suggestions from code review Co-authored-by: antalszava --- pennylane/transforms/batch_partial.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index fa62ca8446e..700e2f694d4 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -73,8 +73,8 @@ def circuit(x, y): qml.RY(y, wires=1) return qml.expval(qml.PauliZ(wires=0) @ qml.PauliZ(wires=1)) - The ``qml.batch_partial`` decorator allows us to create a partially evaluated - function that wraps the QNode. For example, + The ``qml.batch_partial`` decorator allows us to create a partial callable + object that wraps the QNode. For example, >>> y = np.array(0.2) >>> batched_partial_circuit = qml.batch_partial(circuit, y=y) @@ -88,7 +88,7 @@ def circuit(x, y): tensor([0.97517033, 0.95350781, 0.91491915, 0.86008934], requires_grad=True) Jacobians can be computed for the arguments of the wrapper function, but - not for any partially evaluated arguments passed to ``qml.batch_partial``: + not for any pre-supplied argument passed to ``qml.batch_partial``: >>> qml.jacobian(batched_partial_circuit)(x) array([[-0.0978434 , 0. , 0. , 0. ], @@ -142,7 +142,7 @@ def circuit(x, y): sig = inspect.signature(qnode).parameters if is_partial: - # the partially evaluated function must have at least one more + # the batched partial function must have at least one more # parameter, otherwise batching doesn't make sense if len(sig) <= len(partial_kwargs): raise ValueError("Partial evaluation must leave at least one unevaluated parameter") From 92df7be37a459eadedff7a2fe248cb0c06bf5cbc Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Thu, 26 May 2022 12:07:10 -0400 Subject: [PATCH 25/26] rerun ci From 37c2b5b4a7d3fa778198e36df5ac328a90db481b Mon Sep 17 00:00:00 2001 From: Edward Jiang Date: Tue, 31 May 2022 10:33:22 -0400 Subject: [PATCH 26/26] Remove unused import --- pennylane/transforms/batch_partial.py | 1 - 1 file changed, 1 deletion(-) diff --git a/pennylane/transforms/batch_partial.py b/pennylane/transforms/batch_partial.py index 700e2f694d4..5c726a88060 100644 --- a/pennylane/transforms/batch_partial.py +++ b/pennylane/transforms/batch_partial.py @@ -18,7 +18,6 @@ import inspect import pennylane as qml -from pennylane import numpy as np def _convert_to_args(sig, args, kwargs):