Differentiable parameter-shift gradient transform #1479

josh146 · 2021-07-27T15:57:38Z

Context: As part of the roadmap for supporting differentiable batch execution, gradient logic will be moved out of the subclasses and into a module of pure functions. This is the second PR after #1476; here, we move the finite-difference logic out of QubitParamShiftTape and into a new gradients package.

Description of the Change:

Adds a function expval_param_shift for computing the parameter-shift gradient of a tape terminating in expectation values. Directly equivalent to QubitParamShift.parameter_shift.
Adds a function var_param_shift for computing the parameter-shift gradient of a tape terminating in one or more variances. Directly equivalent to QubitParamShift.parameter_shift_var.
Adds a wrapper function param_shift, which does the following:
- Basic input validation
- Performs static analysis of the tape parameters, to determine which support parameter-shift rules.
- A fallback mode; finite_diff is called for any unsupported parameters
- Dispatches to one of expval_param_shift or var_param_shift depending on the structure of the tape
- Provides a processing function to combine the gradients computed via fallback and the parameter-shift methods.

Benefits:

The parameter-shift logic is now much more user and dev accessible
The parameter-shift logic is now differentiable, allowing higher-order derivatives to be accessed no matter the gradient recipe.
Redundant and zero terms are automatically removed from gradient recipes
If the output value of the unshifted, input tape is known, and the gradient recipe contains an unshifted component, this can be provided to the parameter-shift rule to reduce the number of evaluations required.

Possible Drawbacks:

See drawbacks in Differentiable finite-difference gradient transform #1476
For variance support, the only non-involutory observable supported is qml.Hermitian. We should consider extending support to qml.Hamiltonian, now that expval(H) is supported.
We require JacobianTape as input, since we continue to rely on that subclasses static gradient analysis methods
To cache the static gradient analysis, it is appended to the input tape - a better approach should be considered.

Related GitHub Issues:

…inite-diff

Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

github-actions · 2021-07-27T15:57:54Z

Hello. You may have forgotten to update the changelog!
Please edit .github/CHANGELOG.md with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

codecov · 2021-07-27T16:08:06Z

Codecov Report

Merging #1479 (3009847) into master (f2bb1a0) will increase coverage by 0.01%.
The diff coverage is 99.36%.

@@            Coverage Diff             @@
##           master    #1479      +/-   ##
==========================================
+ Coverage   98.32%   98.34%   +0.01%     
==========================================
  Files         180      181       +1     
  Lines       12741    12899     +158     
==========================================
+ Hits        12528    12685     +157     
- Misses        213      214       +1

Impacted Files	Coverage Δ
pennylane/gradients/parameter_shift.py	`99.35% <99.35%> (ø)`
pennylane/gradients/__init__.py	`100.00% <100.00%> (ø)`
pennylane/gradients/finite_difference.py	`100.00% <100.00%> (ø)`

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update f2bb1a0...3009847. Read the comment docs.

…inite-diff

josh146 · 2021-07-29T09:35:34Z

@antalszava @glassnotes: I fixed the variance rule to take into account Projectors, and also added a test, in commits https://github.com/PennyLaneAI/pennylane/pull/1479/files/be46c0ab7019b80cfae4f932c41a2c8028391191..0c5666c7bdb3bb390ca967660d272953973eed1b

antalszava

Just have some questions/suggestions, otherwise it's looking awesome! 🎉

After our chat yesterday, curious: how would we handle template/tape expansions?

As you've suggested, the following now works like a charm:

n_layers = 4
n_wires = 4

dev = qml.device('default.qubit', wires=n_wires)

rng = np.random.default_rng(seed=42)
params = np.array(rng.standard_normal((n_layers, n_wires)))

with qml.tape.JacobianTape() as tape:
    qml.templates.BasicEntanglerLayers(params, wires=range(n_wires)).expand()
    qml.expval(qml.PauliZ(0))
    
tape = tape.expand(stop_at=lambda obj: not isinstance(obj, qml.tape.QuantumTape) and dev.supports_operation(obj.name))

tape.trainable_params = set(range(n_layers * n_wires))

Not sure, however, if it will be intuitive to users to call expand on the template and then call tape.expand too. Maybe it's for another PR?

antalszava · 2021-07-29T11:25:21Z

tests/gradients/test_parameter_shift.py

+        assert res.shape == (1, 2)
+
+        # only called for parameter 0
+        assert spy.call_args[0][0:2] == (tape, [0])


Oh right! So the first two elements are checked for spy.call_args[0], though would we want to make sure that there's no spy.call_args[1] here? Or that the function was called only once as per # only called for parameter 0.

pennylane/gradients/parameter_shift.py

antalszava · 2021-07-29T11:36:55Z

pennylane/gradients/parameter_shift.py

+    _gradient_analysis(tape)
+    gradient_tapes = []
+
+    # TODO: replace the JacobianTape._grad_method_validation


Hmm, I think I still don't fully grasp the comment. Would we not deprecate JacobianTape as is and JacobianTape._grad_method_validation with it? How come we'd need to revisit this particular spot?

pennylane/gradients/parameter_shift.py

tests/gradients/test_parameter_shift.py

pennylane/gradients/parameter_shift.py

josh146 · 2021-07-29T13:08:13Z

Not sure, however, if it will be intuitive to users to call expand on the template and then call tape.expand too. Maybe it's for another PR?

@antalszava yes exactly 🙂 This PR simply adds the low-level tape transform. A future PR will add support via QNodes, which is always the user-facing level in PL.

glassnotes

Have made it through all the test cases, mostly just caught typos this time around.

One question I have is what happens to the gates that use the four-term shift rule, like the controlled rotations. They're included in some test cases, but the docs for param_shift are pretty clear about the shift value only being valid for the two-term rule, but the controlled rotation tests don't pass any custom parameters or gradient recipes.

pennylane/gradients/parameter_shift.py

glassnotes · 2021-07-29T13:18:11Z

tests/gradients/test_parameter_shift.py

+            qml.RY(-0.654, wires=[0])
+            qml.expval(qml.PauliZ(0))
+
+        gradient_recipes = [[[-1e7, 1, 0], [1e7, 1, 1e7]], [[-1e7, 1, 0], [1e7, 1, 1e7]]]


Why are these numbers so large? 😨

oh this is forward finite differences haha. So the coefficients are ±1/h=±1/1e-7 =1e7 😆

glassnotes · 2021-07-29T13:21:15Z

tests/gradients/test_parameter_shift.py

+
+    def test_independent_parameters_analytic(self):
+        """Test the case where expectation values are independent of some parameters. For those
+        parameters, the gradient should be evaluated to zero without executing the device."""


This is actually such an awesome feature 😁

It's a contentious feature! Because it uses networkx which can be slow 😬

It depends on your perspective; do you want to save quantum compute at the expense of classical compute?

tests/gradients/test_parameter_shift.py

glassnotes · 2021-07-29T13:34:37Z

tests/gradients/test_parameter_shift.py

+        assert len(tapes) == 4
+
+        res = fn(dev.batch_execute(tapes))
+        assert res.shape == (5, 2)


I don't quite follow how this output shape is obtained, is it because there is 1 expectation value and 4 output probabilities, with a gradient for each of the two parameters?

Yep! The Jacobian is simply reshaped to be a 2D array, so the 1 expval and the 4 probs = 5 outputs. Coupled with the 2 parameters, the output Jacobian has shape (5, 2).

Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

glassnotes

Thanks @josh146 , looks like all the changes are in, happy to approve!

tests/gradients/test_parameter_shift.py

Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

josh146 · 2021-07-30T15:27:10Z

[ch7869]

antalszava · 2021-08-03T15:47:30Z

tests/gradients/test_parameter_shift.py

@@ -413,13 +442,15 @@ def test_variance_gradients_agree_finite_differences(self, tol):
            qml.CNOT(wires=[1, 0])
            qml.RX(params[2], wires=[0])
            qml.CNOT(wires=[0, 1])
-            qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(1))
+            qml.expval(qml.PauliZ(0)), qml.var(qml.PauliZ(0) @ qml.PauliX(1))


How come this was changed?

I don't recall 🤔

antalszava

Looks good to me! 💪 😍 Excited for this!

josh146 and others added 24 commits July 25, 2021 23:19

more

fc67b6c

more

e57f624

more

8786425

more

3f924b2

more

0542bd4

add tests

9f6b53c

more tests

ce46343

more tests

1a28703

Merge branch 'master' into finite-diff

c436a42

changelog

a449fef

Merge branch 'finite-diff' of github.com:PennyLaneAI/pennylane into f…

a61b3d9

…inite-diff

add todo

e63ffbd

fix

83bad03

fix

148cd45

fix

c3a5b41

Apply suggestions from code review

dacc22a

Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

Merge branch 'master' into finite-diff

e460d2e

add parameter-shift

5b2ca47

add parameter-shift

f5908e5

changes

a9269c5

Merge branch 'finite-diff' into parameter-shift

91ccca2

fixes

23454fa

fixes

623aaca

Merge branch 'master' into finite-diff

c855407

josh146 added the WIP 🚧 Work-in-progress label Jul 27, 2021

josh146 added 3 commits July 28, 2021 00:09

fixes

807a921

suggested changes

33d608b

Merge branch 'finite-diff' of github.com:PennyLaneAI/pennylane into f…

32d2daa

…inite-diff

josh146 added 3 commits July 29, 2021 15:54

Merge branch 'finite-diff' of github.com:PennyLaneAI/pennylane into f…

5614fd5

…inite-diff

added multiplier test

49fe608

Merge branch 'finite-diff' into parameter-shift

98e8b19

Base automatically changed from finite-diff to master July 29, 2021 08:19

josh146 added 5 commits July 29, 2021 16:40

Suggested changes

4d0ff45

merge master

80f596e

suggested changes

be46c0a

suggested changes

78ee8f7

typo

0c5666c

antalszava reviewed Jul 29, 2021

View reviewed changes

Merge branch 'master' into parameter-shift

9c141d4

josh146 commented Jul 29, 2021

View reviewed changes

pennylane/gradients/parameter_shift.py Outdated Show resolved Hide resolved

glassnotes reviewed Jul 29, 2021

View reviewed changes

josh146 and others added 2 commits July 29, 2021 21:50

Apply suggestions from code review

bdc10ff

Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

Merge branch 'master' into parameter-shift

8bbe123

josh146 mentioned this pull request Jul 30, 2021

(non)Differentiable CV parameter-shift gradient transform #1486

Merged

suggested changes

6d9636c

josh146 requested review from glassnotes and antalszava July 30, 2021 14:13

Merge branch 'master' into parameter-shift

c1419a6

glassnotes approved these changes Jul 30, 2021

View reviewed changes

tests/gradients/test_parameter_shift.py Outdated Show resolved Hide resolved

Update tests/gradients/test_parameter_shift.py

283d5c8

Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

Merge branch 'master' into parameter-shift

3009847

antalszava reviewed Aug 3, 2021

View reviewed changes

antalszava approved these changes Aug 3, 2021

View reviewed changes

josh146 merged commit ab71001 into master Aug 3, 2021

josh146 deleted the parameter-shift branch August 3, 2021 17:03

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Differentiable parameter-shift gradient transform #1479

Differentiable parameter-shift gradient transform #1479

josh146 commented Jul 27, 2021

github-actions bot commented Jul 27, 2021

codecov bot commented Jul 27, 2021 •

edited

Loading

josh146 commented Jul 29, 2021

antalszava left a comment

antalszava Jul 29, 2021

antalszava Jul 29, 2021

josh146 commented Jul 29, 2021

glassnotes left a comment

glassnotes Jul 29, 2021

josh146 Jul 29, 2021

glassnotes Jul 29, 2021

josh146 Jul 29, 2021

glassnotes Jul 29, 2021

josh146 Jul 29, 2021

glassnotes left a comment

josh146 commented Jul 30, 2021

antalszava Aug 3, 2021

josh146 Aug 3, 2021

antalszava left a comment

Differentiable parameter-shift gradient transform #1479

Differentiable parameter-shift gradient transform #1479

Conversation

josh146 commented Jul 27, 2021

github-actions bot commented Jul 27, 2021

codecov bot commented Jul 27, 2021 • edited Loading

Codecov Report

josh146 commented Jul 29, 2021

antalszava left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

josh146 commented Jul 29, 2021

glassnotes left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

glassnotes left a comment

Choose a reason for hiding this comment

josh146 commented Jul 30, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

antalszava left a comment

Choose a reason for hiding this comment

codecov bot commented Jul 27, 2021 •

edited

Loading