From 22da9a05167fff94bdad0f9154b5d822f1f517e4 Mon Sep 17 00:00:00 2001 From: Astral Cai Date: Thu, 6 Jun 2024 14:52:00 -0400 Subject: [PATCH] Set `grad_method=None` for `ControlledSequence`, `Reflection`, `AmplitudeAmplification`, and `Qubitization`. (#5806) **Context:** Templates that are not actually supported by `parameter_shift` should have `grad_method=None` so that they are decomposed by `_expand_transform_param_shift` **Description of the Change:** 1. Adds the `data` of components of the templates to the `data` of the templates such that trainable parameters are tracked 2. Adds `grad_method=None` for `ControlledSequence`, `Reflection`, `AmplitudeAmplification`, and `Qubitization`. **Related GitHub Issues:** Fixes https://github.com/PennyLaneAI/pennylane/issues/5802 [sc-64967] --- doc/releases/changelog-dev.md | 3 ++ .../subroutines/amplitude_amplification.py | 6 ++- .../subroutines/controlled_sequence.py | 2 + .../templates/subroutines/qubitization.py | 4 +- pennylane/templates/subroutines/reflection.py | 4 +- .../test_amplitude_amplification.py | 31 +++++++++---- .../test_controlled_sequence.py | 45 ++++++++++++++----- .../test_subroutines/test_qubitization.py | 11 +++-- .../test_subroutines/test_reflection.py | 40 ++++++++++++----- 9 files changed, 108 insertions(+), 38 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index e57878e54e9..86f6889eadc 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -332,6 +332,9 @@ * `CNOT` and `Toffoli` now have an `arithmetic_depth` of `1`, as they are controlled operations. [(#5797)](https://github.com/PennyLaneAI/pennylane/pull/5797) +* Fixes a bug where the gradient of `ControlledSequence`, `Reflection`, `AmplitudeAmplification`, and `Qubitization` is incorrect on `default.qubit.legacy` with `parameter_shift`. + [(#5806)](https://github.com/PennyLaneAI/pennylane/pull/5806) +

Contributors ✍️

This release contains contributions from (in alphabetical order): diff --git a/pennylane/templates/subroutines/amplitude_amplification.py b/pennylane/templates/subroutines/amplitude_amplification.py index 08e050ff73f..fd89816b984 100644 --- a/pennylane/templates/subroutines/amplitude_amplification.py +++ b/pennylane/templates/subroutines/amplitude_amplification.py @@ -102,6 +102,8 @@ def circuit(): [0.013, 0.013, 0.91, 0.013, 0.013, 0.013, 0.013, 0.013] """ + grad_method = None + def _flatten(self): data = (self.hyperparameters["U"], self.hyperparameters["O"]) metadata = tuple(item for item in self.hyperparameters.items() if item[0] not in ["O", "U"]) @@ -141,11 +143,11 @@ def __init__( self.hyperparameters["p_min"] = p_min self.hyperparameters["reflection_wires"] = qml.wires.Wires(reflection_wires) - super().__init__(wires=wires) + super().__init__(*U.data, *O.data, wires=wires) # pylint:disable=arguments-differ @staticmethod - def compute_decomposition(**kwargs): + def compute_decomposition(*_, **kwargs): U = kwargs["U"] O = kwargs["O"] iters = kwargs["iters"] diff --git a/pennylane/templates/subroutines/controlled_sequence.py b/pennylane/templates/subroutines/controlled_sequence.py index e9efad65281..517129556cc 100644 --- a/pennylane/templates/subroutines/controlled_sequence.py +++ b/pennylane/templates/subroutines/controlled_sequence.py @@ -69,6 +69,8 @@ def circuit(): """ + grad_method = None + def _flatten(self): return (self.base,), (self.control,) diff --git a/pennylane/templates/subroutines/qubitization.py b/pennylane/templates/subroutines/qubitization.py index ae45487404a..5847d64c652 100644 --- a/pennylane/templates/subroutines/qubitization.py +++ b/pennylane/templates/subroutines/qubitization.py @@ -92,6 +92,8 @@ def circuit(): eigenvalue: 0.7 """ + grad_method = None + @classmethod def _primitive_bind_call(cls, *args, **kwargs): return cls._primitive.bind(*args, **kwargs) @@ -104,7 +106,7 @@ def __init__(self, hamiltonian, control, id=None): "control": qml.wires.Wires(control), } - super().__init__(wires=wires, id=id) + super().__init__(*hamiltonian.data, wires=wires, id=id) def _flatten(self): data = (self.hyperparameters["hamiltonian"],) diff --git a/pennylane/templates/subroutines/reflection.py b/pennylane/templates/subroutines/reflection.py index a39cf7cb3af..bad3316b3fe 100644 --- a/pennylane/templates/subroutines/reflection.py +++ b/pennylane/templates/subroutines/reflection.py @@ -104,6 +104,8 @@ def circuit(): """ + grad_method = None + @classmethod def _primitive_bind_call(cls, *args, **kwargs): return cls._primitive.bind(*args, **kwargs) @@ -136,7 +138,7 @@ def __init__(self, U, alpha=np.pi, reflection_wires=None, id=None): "reflection_wires": tuple(reflection_wires), } - super().__init__(alpha, wires=wires, id=id) + super().__init__(alpha, *U.data, wires=wires, id=id) def map_wires(self, wire_map: dict): # pylint: disable=protected-access diff --git a/tests/templates/test_subroutines/test_amplitude_amplification.py b/tests/templates/test_subroutines/test_amplitude_amplification.py index 487d0f3dbe4..9dff939ee6d 100644 --- a/tests/templates/test_subroutines/test_amplitude_amplification.py +++ b/tests/templates/test_subroutines/test_amplitude_amplification.py @@ -145,7 +145,7 @@ def circuit(params): qml.RZ(params[1], wires=0), iters=3, fixed_point=True, - work_wire=3, + work_wire=2, ) return qml.expval(qml.PauliZ(0)) @@ -156,28 +156,36 @@ def circuit(params): params = np.array([0.9, 0.1]) @pytest.mark.autograd - def test_qnode_autograd(self): + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + @pytest.mark.parametrize("shots", [None, 50000]) + def test_qnode_autograd(self, device, shots): """Test that the QNode executes with Autograd.""" - dev = qml.device("default.qubit") - qnode = qml.QNode(self.circuit, dev, interface="autograd") + dev = qml.device(device, wires=3, shots=shots) + diff_method = "backprop" if shots is None else "parameter-shift" + qnode = qml.QNode(self.circuit, dev, interface="autograd", diff_method=diff_method) params = qml.numpy.array(self.params, requires_grad=True) res = qml.grad(qnode)(params) assert qml.math.shape(res) == (2,) - assert np.allclose(res, self.exp_grad, atol=1e-5) + assert np.allclose(res, self.exp_grad, atol=0.01) @pytest.mark.jax @pytest.mark.parametrize("use_jit", [False, True]) @pytest.mark.parametrize("shots", [None, 50000]) - def test_qnode_jax(self, shots, use_jit): + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + def test_qnode_jax(self, shots, use_jit, device): """Test that the QNode executes and is differentiable with JAX. The shots argument controls whether autodiff or parameter-shift gradients are used.""" import jax jax.config.update("jax_enable_x64", True) - dev = qml.device("default.qubit", shots=shots, seed=10) + if device == "default.qubit": + dev = qml.device("default.qubit", shots=shots, seed=10) + else: + dev = qml.device("default.qubit.legacy", shots=shots, wires=3) + diff_method = "backprop" if shots is None else "parameter-shift" qnode = qml.QNode(self.circuit, dev, interface="jax", diff_method=diff_method) if use_jit: @@ -195,12 +203,17 @@ def test_qnode_jax(self, shots, use_jit): @pytest.mark.torch @pytest.mark.parametrize("shots", [None, 50000]) - def test_qnode_torch(self, shots): + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + def test_qnode_torch(self, shots, device): """Test that the QNode executes and is differentiable with Torch. The shots argument controls whether autodiff or parameter-shift gradients are used.""" import torch - dev = qml.device("default.qubit", shots=shots, seed=10) + if device == "default.qubit": + dev = qml.device("default.qubit", shots=shots, seed=10) + else: + dev = qml.device("default.qubit.legacy", shots=shots, wires=3) + diff_method = "backprop" if shots is None else "parameter-shift" qnode = qml.QNode(self.circuit, dev, interface="torch", diff_method=diff_method) diff --git a/tests/templates/test_subroutines/test_controlled_sequence.py b/tests/templates/test_subroutines/test_controlled_sequence.py index 6286af4e577..f7dc9c8a233 100644 --- a/tests/templates/test_subroutines/test_controlled_sequence.py +++ b/tests/templates/test_subroutines/test_controlled_sequence.py @@ -31,6 +31,7 @@ def test_standard_validity(): class TestInitialization: + def test_id(self): """Tests that the id attribute can be set.""" op = qml.ControlledSequence(qml.RX(0.25, wires=3), control=[0, 1, 2], id="a") @@ -57,6 +58,7 @@ def test_name(self): class TestProperties: + def test_hash(self): """Test that op.hash uniquely describes a ControlledSequence""" @@ -97,6 +99,7 @@ def test_has_matrix(self): class TestMethods: + def test_repr(self): """Test that the operator repr is as expected""" op = qml.ControlledSequence(qml.RX(0.25, wires=3), control=[0, 1, 2]) @@ -189,28 +192,41 @@ def test_qnode_numpy(self): assert np.allclose(res, self.exp_result, atol=0.002) @pytest.mark.autograd - def test_qnode_autograd(self): + @pytest.mark.parametrize("shots", [None, 50000]) + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + def test_qnode_autograd(self, shots, device): """Test that the QNode executes with Autograd.""" - dev = qml.device("default.qubit") - qnode = qml.QNode(self.circuit, dev, interface="autograd") - + dev = qml.device(device, wires=4, shots=shots) + diff_method = "backprop" if shots is None else "parameter-shift" + qnode = qml.QNode(self.circuit, dev, interface="autograd", diff_method=diff_method) x = qml.numpy.array(self.x, requires_grad=True) + res = qnode(x) assert qml.math.shape(res) == (16,) assert np.allclose(res, self.exp_result, atol=0.002) + res = qml.jacobian(qnode)(x) + assert np.shape(res) == (16,) + assert np.allclose(res, self.exp_jac, atol=0.005) + @pytest.mark.jax @pytest.mark.parametrize("use_jit", [False, True]) - @pytest.mark.parametrize("shots", [None, 10000]) - def test_qnode_jax(self, shots, use_jit): + @pytest.mark.parametrize("shots", [None, 50000]) + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + def test_qnode_jax(self, shots, use_jit, device): """Test that the QNode executes and is differentiable with JAX. The shots argument controls whether autodiff or parameter-shift gradients are used.""" + import jax jax.config.update("jax_enable_x64", True) - dev = qml.device("default.qubit", shots=shots, seed=10) + if device == "default.qubit": + dev = qml.device("default.qubit", shots=shots, seed=10) + else: + dev = qml.device("default.qubit.legacy", shots=shots, wires=4) + diff_method = "backprop" if shots is None else "parameter-shift" qnode = qml.QNode(self.circuit, dev, interface="jax", diff_method=diff_method) if use_jit: @@ -230,13 +246,19 @@ def test_qnode_jax(self, shots, use_jit): assert np.allclose(jac, self.exp_jac, atol=0.006) @pytest.mark.torch - @pytest.mark.parametrize("shots", [None, 10000]) - def test_qnode_torch(self, shots): + @pytest.mark.parametrize("shots", [None, 50000]) + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + def test_qnode_torch(self, shots, device): """Test that the QNode executes and is differentiable with Torch. The shots argument controls whether autodiff or parameter-shift gradients are used.""" + import torch - dev = qml.device("default.qubit", shots=shots, seed=10) + if device == "default.qubit": + dev = qml.device("default.qubit", shots=shots, seed=10) + else: + dev = qml.device("default.qubit.legacy", shots=shots, wires=4) + diff_method = "backprop" if shots is None else "parameter-shift" qnode = qml.QNode(self.circuit, dev, interface="torch", diff_method=diff_method) @@ -247,7 +269,7 @@ def test_qnode_torch(self, shots): jac = torch.autograd.functional.jacobian(qnode, x) assert qml.math.shape(jac) == (16,) - assert qml.math.allclose(jac, self.exp_jac, atol=0.006) + assert qml.math.allclose(jac, self.exp_jac, atol=0.005) @pytest.mark.tf @pytest.mark.parametrize("shots", [None, 10000]) @@ -255,6 +277,7 @@ def test_qnode_torch(self, shots): def test_qnode_tf(self, shots): """Test that the QNode executes and is differentiable with TensorFlow. The shots argument controls whether autodiff or parameter-shift gradients are used.""" + import tensorflow as tf dev = qml.device("default.qubit", shots=shots, seed=10) diff --git a/tests/templates/test_subroutines/test_qubitization.py b/tests/templates/test_subroutines/test_qubitization.py index 430382f011e..124445f0ee8 100644 --- a/tests/templates/test_subroutines/test_qubitization.py +++ b/tests/templates/test_subroutines/test_qubitization.py @@ -227,14 +227,19 @@ def test_qnode_autograd(self): "use_jit , shots", ((False, None), (True, None), (False, 50000)), ) # TODO: (True, 50000) fails because jax.jit on jax.grad does not work with AmplitudeEmbedding - def test_qnode_jax(self, shots, use_jit): + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + def test_qnode_jax(self, shots, use_jit, device): """ "Test that the QNode executes and is differentiable with JAX. The shots argument controls whether autodiff or parameter-shift gradients are used.""" import jax jax.config.update("jax_enable_x64", True) - dev = qml.device("default.qubit", shots=shots, seed=10) + if device == "default.qubit": + dev = qml.device("default.qubit", shots=shots, seed=10) + else: + dev = qml.device("default.qubit.legacy", shots=shots, wires=5) + diff_method = "backprop" if shots is None else "parameter-shift" qnode = qml.QNode(self.circuit, dev, interface="jax", diff_method=diff_method) if use_jit: @@ -248,7 +253,7 @@ def test_qnode_jax(self, shots, use_jit): jac = jac_fn(params) assert jac.shape == (4,) - assert np.allclose(jac, self.exp_grad, atol=0.01) + assert np.allclose(jac, self.exp_grad, atol=0.05) @pytest.mark.torch @pytest.mark.parametrize( diff --git a/tests/templates/test_subroutines/test_reflection.py b/tests/templates/test_subroutines/test_reflection.py index 417c3d715d0..5d7cefb729c 100644 --- a/tests/templates/test_subroutines/test_reflection.py +++ b/tests/templates/test_subroutines/test_reflection.py @@ -163,28 +163,40 @@ def test_lightning_qubit(self): assert np.allclose(res, self.exp_result, atol=0.002) @pytest.mark.autograd - def test_qnode_autograd(self): + @pytest.mark.parametrize("shots", [None, 50000]) + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + def test_qnode_autograd(self, shots, device): """Test that the QNode executes with Autograd.""" - dev = qml.device("default.qubit") - qnode = qml.QNode(self.circuit, dev, interface="autograd") + dev = qml.device(device, shots=shots, wires=3) + diff_method = "backprop" if shots is None else "parameter-shift" + qnode = qml.QNode(self.circuit, dev, interface="autograd", diff_method=diff_method) x = qml.numpy.array(self.x, requires_grad=True) res = qnode(x) assert qml.math.shape(res) == (8,) - assert np.allclose(res, self.exp_result, atol=0.002) + assert np.allclose(res, self.exp_result, atol=0.005) + + res = qml.jacobian(qnode)(x) + assert np.shape(res) == (8,) + assert np.allclose(res, self.exp_jac, atol=0.005) @pytest.mark.jax @pytest.mark.parametrize("use_jit", [False, True]) @pytest.mark.parametrize("shots", [None, 50000]) - def test_qnode_jax(self, shots, use_jit): + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + def test_qnode_jax(self, shots, use_jit, device): """Test that the QNode executes and is differentiable with JAX. The shots argument controls whether autodiff or parameter-shift gradients are used.""" import jax jax.config.update("jax_enable_x64", True) - dev = qml.device("default.qubit", shots=shots, seed=10) + if device == "default.qubit": + dev = qml.device("default.qubit", shots=shots, seed=10) + else: + dev = qml.device("default.qubit.legacy", shots=shots, wires=3) + diff_method = "backprop" if shots is None else "parameter-shift" qnode = qml.QNode(self.circuit, dev, interface="jax", diff_method=diff_method) if use_jit: @@ -201,27 +213,33 @@ def test_qnode_jax(self, shots, use_jit): jac = jac_fn(x) assert jac.shape == (8,) - assert np.allclose(jac, self.exp_jac, atol=0.006) + assert np.allclose(jac, self.exp_jac, atol=0.005) @pytest.mark.torch @pytest.mark.parametrize("shots", [None, 50000]) - def test_qnode_torch(self, shots): + @pytest.mark.parametrize("device", ["default.qubit", "default.qubit.legacy"]) + def test_qnode_torch(self, shots, device): """Test that the QNode executes and is differentiable with Torch. The shots argument controls whether autodiff or parameter-shift gradients are used.""" + import torch - dev = qml.device("default.qubit", shots=shots, seed=10) + if device == "default.qubit": + dev = qml.device("default.qubit", shots=shots, seed=10) + else: + dev = qml.device("default.qubit.legacy", shots=shots, wires=3) + diff_method = "backprop" if shots is None else "parameter-shift" qnode = qml.QNode(self.circuit, dev, interface="torch", diff_method=diff_method) x = torch.tensor(self.x, requires_grad=True) res = qnode(x) assert qml.math.shape(res) == (8,) - assert qml.math.allclose(res, self.exp_result, atol=0.002) + assert qml.math.allclose(res, self.exp_result, atol=0.005) jac = torch.autograd.functional.jacobian(qnode, x) assert qml.math.shape(jac) == (8,) - assert qml.math.allclose(jac, self.exp_jac, atol=0.006) + assert qml.math.allclose(jac, self.exp_jac, atol=0.005) @pytest.mark.tf @pytest.mark.parametrize("shots", [None, 50000])