From 74bc82b8c07a756f6a8acd4b99ea84119c4d4b8d Mon Sep 17 00:00:00 2001 From: dwierichs Date: Thu, 5 Sep 2024 19:22:20 +0200 Subject: [PATCH 1/5] fix bug --- pennylane/gradients/jvp.py | 11 +++++++--- pennylane/workflow/jacobian_products.py | 20 +++++++++--------- tests/gradients/core/test_jvp.py | 28 +++++++++++++++++++++++++ 3 files changed, 46 insertions(+), 13 deletions(-) diff --git a/pennylane/gradients/jvp.py b/pennylane/gradients/jvp.py index e3634cce492..1247bbe37a9 100644 --- a/pennylane/gradients/jvp.py +++ b/pennylane/gradients/jvp.py @@ -295,11 +295,16 @@ def jvp(tape, tangent, gradient_fn, gradient_kwargs=None): if len(tape.trainable_params) == 0: # The tape has no trainable parameters; the JVP # is simply none. - def zero_vjp(_): - res = tuple(np.zeros(mp.shape(None, tape.shots)) for mp in tape.measurements) + def zero_jvp_for_single_shots(s): + res = tuple(np.zeros(mp.shape(shots=s)) for mp in tape.measurements) return res[0] if len(tape.measurements) == 1 else res - return tuple(), zero_vjp + def zero_jvp(_): + if tape.shots.has_partitioned_shots: + return tuple(zero_jvp_for_single_shots(s) for s in tape.shots) + return zero_jvp_for_single_shots(tape.shots) + + return tuple(), zero_jvp multi_m = len(tape.measurements) > 1 diff --git a/pennylane/workflow/jacobian_products.py b/pennylane/workflow/jacobian_products.py index f8163813366..688e786c1c5 100644 --- a/pennylane/workflow/jacobian_products.py +++ b/pennylane/workflow/jacobian_products.py @@ -46,6 +46,15 @@ def _compute_vjps(jacs, dys, tapes): return tuple(vjps) +def _zero_jvp_single_shots(shots, tape): + jvp = tuple(np.zeros(mp.shape(shots=shots), dtype=mp.numeric_type) for mp in tape.measurements) + return jvp[0] if len(tape.measurements) == 1 else jvp + +def _zero_jvp(tape): + if tape.shots.has_partitioned_shots: + return tuple(_zero_jvp_single_shots(s, tape) for s in tape.shots) + return _zero_jvp_single_shots(tape.shots, tape) + def _compute_jvps(jacs, tangents, tapes): """Compute the jvps of multiple tapes, directly for a Jacobian and tangents.""" f = {True: qml.gradients.compute_jvp_multi, False: qml.gradients.compute_jvp_single} @@ -54,16 +63,7 @@ def _compute_jvps(jacs, tangents, tapes): for jac, dx, t in zip(jacs, tangents, tapes): multi = len(t.measurements) > 1 if len(t.trainable_params) == 0: - empty_shots = qml.measurements.Shots(None) - zeros_jvp = tuple( - np.zeros(mp.shape(None, empty_shots), dtype=mp.numeric_type) - for mp in t.measurements - ) - zeros_jvp = zeros_jvp[0] if len(t.measurements) == 1 else zeros_jvp - if t.shots.has_partitioned_shots: - jvps.append(tuple(zeros_jvp for _ in range(t.shots.num_copies))) - else: - jvps.append(zeros_jvp) + jvps.append(_zero_jvp(t)) elif t.shots.has_partitioned_shots: jvps.append(tuple(f[multi](dx, j) for j in jac)) else: diff --git a/tests/gradients/core/test_jvp.py b/tests/gradients/core/test_jvp.py index 54bdb051572..bf694740b6a 100644 --- a/tests/gradients/core/test_jvp.py +++ b/tests/gradients/core/test_jvp.py @@ -860,6 +860,34 @@ def test_all_tapes_no_trainable_parameters(self): assert qml.math.allclose(fn([])[0], np.array(0.0)) assert qml.math.allclose(fn([])[1], np.array(0.0)) + def test_some_tapes_no_trainable_parameters(self): + """If some tapes have no trainable parameters all outputs will be None""" + + with qml.queuing.AnnotatedQueue() as q1: + qml.RX(0.4, wires=0) + qml.expval(qml.PauliZ(0)) + + tape1 = qml.tape.QuantumScript.from_queue(q1) + with qml.queuing.AnnotatedQueue() as q2: + qml.RX(0.4, wires=0) + qml.RX(0.6, wires=0) + qml.CNOT(wires=[0, 1]) + qml.expval(qml.PauliZ(0)) + + tape2 = qml.tape.QuantumScript.from_queue(q2) + tape1.trainable_params = {0} + tape2.trainable_params = set() + + tapes = [tape1, tape2] + tangents = [np.array([0.7]), np.array([1.0, 0.0])] + + v_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, param_shift) + + assert len(v_tapes) == 2 + ps_res = [np.cos(0.4+np.pi/2), np.cos(0.4-np.pi/2)] + assert qml.math.allclose(fn(ps_res)[0], -np.sin(0.4) * 0.7) + assert qml.math.allclose(fn(ps_res)[1], np.array(0.0)) + def test_zero_tangent(self): """A zero dy vector will return no tapes and a zero matrix""" dev = qml.device("default.qubit", wires=2) From a7a53de4ad0cfb00c6809d6c86448a5bb5825b2b Mon Sep 17 00:00:00 2001 From: dwierichs Date: Thu, 5 Sep 2024 19:25:57 +0200 Subject: [PATCH 2/5] -a --- pennylane/workflow/jacobian_products.py | 2 ++ tests/gradients/core/test_jvp.py | 45 ++++++------------------- 2 files changed, 12 insertions(+), 35 deletions(-) diff --git a/pennylane/workflow/jacobian_products.py b/pennylane/workflow/jacobian_products.py index 688e786c1c5..44188773221 100644 --- a/pennylane/workflow/jacobian_products.py +++ b/pennylane/workflow/jacobian_products.py @@ -50,11 +50,13 @@ def _zero_jvp_single_shots(shots, tape): jvp = tuple(np.zeros(mp.shape(shots=shots), dtype=mp.numeric_type) for mp in tape.measurements) return jvp[0] if len(tape.measurements) == 1 else jvp + def _zero_jvp(tape): if tape.shots.has_partitioned_shots: return tuple(_zero_jvp_single_shots(s, tape) for s in tape.shots) return _zero_jvp_single_shots(tape.shots, tape) + def _compute_jvps(jacs, tangents, tapes): """Compute the jvps of multiple tapes, directly for a Jacobian and tangents.""" f = {True: qml.gradients.compute_jvp_multi, False: qml.gradients.compute_jvp_single} diff --git a/tests/gradients/core/test_jvp.py b/tests/gradients/core/test_jvp.py index bf694740b6a..8ffd1a7f9bd 100644 --- a/tests/gradients/core/test_jvp.py +++ b/tests/gradients/core/test_jvp.py @@ -17,6 +17,7 @@ import pennylane as qml from pennylane import numpy as np from pennylane.gradients import param_shift +from pennylane.measurements.shots import Shots _x = np.arange(12).reshape((2, 3, 2)) @@ -799,7 +800,8 @@ def cost_fn(params, tangent): class TestBatchJVP: """Tests for the batch JVP function""" - def test_one_tape_no_trainable_parameters(self): + @pytest.mark.parametrize("shots", [Shots(None), Shots(10), Shots([20, 10])]) + def test_one_tape_no_trainable_parameters(self, shots): """A tape with no trainable parameters will simply return None""" dev = qml.device("default.qubit", wires=2) @@ -808,14 +810,14 @@ def test_one_tape_no_trainable_parameters(self): qml.CNOT(wires=[0, 1]) qml.expval(qml.PauliZ(0)) - tape1 = qml.tape.QuantumScript.from_queue(q1) + tape1 = qml.tape.QuantumScript.from_queue(q1, shots=shots) with qml.queuing.AnnotatedQueue() as q2: qml.RX(0.4, wires=0) qml.RX(0.6, wires=0) qml.CNOT(wires=[0, 1]) qml.expval(qml.PauliZ(0)) - tape2 = qml.tape.QuantumScript.from_queue(q2) + tape2 = qml.tape.QuantumScript.from_queue(q2, shots=shots) tape1.trainable_params = {} tape2.trainable_params = {0, 1} @@ -823,16 +825,17 @@ def test_one_tape_no_trainable_parameters(self): tangents = [np.array([1.0, 1.0]), np.array([1.0, 1.0])] v_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, param_shift) - assert len(v_tapes) == 4 # Even though there are 3 parameters, only two contribute # to the JVP, so only 2*2=4 quantum evals + assert len(v_tapes) == 4 res = fn(dev.execute(v_tapes)) assert qml.math.allclose(res[0], np.array(0.0)) assert res[1] is not None - def test_all_tapes_no_trainable_parameters(self): + @pytest.mark.parametrize("shots", [Shots(None), Shots(10), Shots([20, 10])]) + def test_all_tapes_no_trainable_parameters(self, shots): """If all tapes have no trainable parameters all outputs will be None""" with qml.queuing.AnnotatedQueue() as q1: @@ -840,14 +843,14 @@ def test_all_tapes_no_trainable_parameters(self): qml.CNOT(wires=[0, 1]) qml.expval(qml.PauliZ(0)) - tape1 = qml.tape.QuantumScript.from_queue(q1) + tape1 = qml.tape.QuantumScript.from_queue(q1, shots=shots) with qml.queuing.AnnotatedQueue() as q2: qml.RX(0.4, wires=0) qml.RX(0.6, wires=0) qml.CNOT(wires=[0, 1]) qml.expval(qml.PauliZ(0)) - tape2 = qml.tape.QuantumScript.from_queue(q2) + tape2 = qml.tape.QuantumScript.from_queue(q2, shots=shots) tape1.trainable_params = set() tape2.trainable_params = set() @@ -860,34 +863,6 @@ def test_all_tapes_no_trainable_parameters(self): assert qml.math.allclose(fn([])[0], np.array(0.0)) assert qml.math.allclose(fn([])[1], np.array(0.0)) - def test_some_tapes_no_trainable_parameters(self): - """If some tapes have no trainable parameters all outputs will be None""" - - with qml.queuing.AnnotatedQueue() as q1: - qml.RX(0.4, wires=0) - qml.expval(qml.PauliZ(0)) - - tape1 = qml.tape.QuantumScript.from_queue(q1) - with qml.queuing.AnnotatedQueue() as q2: - qml.RX(0.4, wires=0) - qml.RX(0.6, wires=0) - qml.CNOT(wires=[0, 1]) - qml.expval(qml.PauliZ(0)) - - tape2 = qml.tape.QuantumScript.from_queue(q2) - tape1.trainable_params = {0} - tape2.trainable_params = set() - - tapes = [tape1, tape2] - tangents = [np.array([0.7]), np.array([1.0, 0.0])] - - v_tapes, fn = qml.gradients.batch_jvp(tapes, tangents, param_shift) - - assert len(v_tapes) == 2 - ps_res = [np.cos(0.4+np.pi/2), np.cos(0.4-np.pi/2)] - assert qml.math.allclose(fn(ps_res)[0], -np.sin(0.4) * 0.7) - assert qml.math.allclose(fn(ps_res)[1], np.array(0.0)) - def test_zero_tangent(self): """A zero dy vector will return no tapes and a zero matrix""" dev = qml.device("default.qubit", wires=2) From d78e4e152fd507c79e443737eae4b6c08f5ab90f Mon Sep 17 00:00:00 2001 From: dwierichs Date: Thu, 5 Sep 2024 19:28:45 +0200 Subject: [PATCH 3/5] changelog --- doc/releases/changelog-dev.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 813ee6374e4..379d0c3a868 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -28,6 +28,9 @@

Bug fixes 🐛

+* Fix a bug where zero-valued JVPs were calculated wrongly in the presence of shot vectors. + [(#6219)](https://github.com/PennyLaneAI/pennylane/pull/6219) + * Fix Pytree serialization of operators with empty shot vectors: [(#6155)](https://github.com/PennyLaneAI/pennylane/pull/6155) @@ -46,3 +49,4 @@ Utkarsh Azad Jack Brown Christina Lee William Maxwell +David Wierichs From e1a1fc53286b060c55d8fbc03555e70d1e81c170 Mon Sep 17 00:00:00 2001 From: David Wierichs Date: Mon, 9 Sep 2024 12:27:48 +0200 Subject: [PATCH 4/5] Apply suggestions from code review Co-authored-by: Christina Lee --- pennylane/gradients/jvp.py | 4 ++-- pennylane/workflow/jacobian_products.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pennylane/gradients/jvp.py b/pennylane/gradients/jvp.py index 1247bbe37a9..4f4f97f3d01 100644 --- a/pennylane/gradients/jvp.py +++ b/pennylane/gradients/jvp.py @@ -296,13 +296,13 @@ def jvp(tape, tangent, gradient_fn, gradient_kwargs=None): # The tape has no trainable parameters; the JVP # is simply none. def zero_jvp_for_single_shots(s): - res = tuple(np.zeros(mp.shape(shots=s)) for mp in tape.measurements) + res = tuple(np.zeros(mp.shape(shots=s), dtype=mp.numeric_type) for mp in tape.measurements) return res[0] if len(tape.measurements) == 1 else res def zero_jvp(_): if tape.shots.has_partitioned_shots: return tuple(zero_jvp_for_single_shots(s) for s in tape.shots) - return zero_jvp_for_single_shots(tape.shots) + return zero_jvp_for_single_shots(tape.shots.total_shots) return tuple(), zero_jvp diff --git a/pennylane/workflow/jacobian_products.py b/pennylane/workflow/jacobian_products.py index 44188773221..d5ede1227a5 100644 --- a/pennylane/workflow/jacobian_products.py +++ b/pennylane/workflow/jacobian_products.py @@ -54,7 +54,7 @@ def _zero_jvp_single_shots(shots, tape): def _zero_jvp(tape): if tape.shots.has_partitioned_shots: return tuple(_zero_jvp_single_shots(s, tape) for s in tape.shots) - return _zero_jvp_single_shots(tape.shots, tape) + return _zero_jvp_single_shots(tape.shots.total_shots, tape) def _compute_jvps(jacs, tangents, tapes): From 1bf10f0224e1879a6679abbade2e4cf81dceab1c Mon Sep 17 00:00:00 2001 From: dwierichs Date: Tue, 10 Sep 2024 09:30:11 +0200 Subject: [PATCH 5/5] black --- pennylane/gradients/jvp.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pennylane/gradients/jvp.py b/pennylane/gradients/jvp.py index 4f4f97f3d01..67428ab5c68 100644 --- a/pennylane/gradients/jvp.py +++ b/pennylane/gradients/jvp.py @@ -296,7 +296,9 @@ def jvp(tape, tangent, gradient_fn, gradient_kwargs=None): # The tape has no trainable parameters; the JVP # is simply none. def zero_jvp_for_single_shots(s): - res = tuple(np.zeros(mp.shape(shots=s), dtype=mp.numeric_type) for mp in tape.measurements) + res = tuple( + np.zeros(mp.shape(shots=s), dtype=mp.numeric_type) for mp in tape.measurements + ) return res[0] if len(tape.measurements) == 1 else res def zero_jvp(_):