Add support for arbitrary linear combination gradient recipes

[antalszava]

Context:
The "gradient recipe" of a gate parameter stores the coefficient (c) and the shift (s) used to compute the partial derivative of an expectation value using the parameter shift rule (∂f(x) = c*(f(x+s) -f(x-s))). The Operation.grad_recipe attribute is used for this. It turns out, that certain gates require more advanced gradient recipes for computing analytic gradients.

While calculating the gradients that involve quantum circuits containing quantum channels, we'd further require that the gate parameter is not only shifted by a scalar constant but also scaled (at times with a scaling factor of 0, such that we only evaluate using the constant shift).

Bearing these points in mind, it would be good for a gradient recipe to have the ability to:

• Define more than two (unique) coefficients and shifts for each gate parameter
• Define a scaling factor for each of such terms

Description of the Change:

The gradient recipe has been changed so that it can support computations like ∂f(x) = Σi c_i f(a_i x + b_i). This is even more flexible, as it allows us to represent terms that do not depend on x (1.2*f(0.2)):

# corresponds to 0.1*f(x+5) - 0.65*f(2x-pi) + 1.2*f(0.2)
grad_recipe = ([[0.1, 1, 5], [-0.65, 2, -pi], [1.2, 0, 0.2]],)

Notes on the changes:

• Made the grad_recipe immutable (tuple), each element of the grad_recipe represents the recipe for a gate parameter; immutability was added as the original recipe is not mutated, but queried for further computation.
• Made the elements of grad_recipe nested lists. Each list that is nested describes a term in the sum making up ∂f(x) and is of the form of [coefficient, scaling factor, shift].
  The reason why lists are used is that we may potentially be post-processing the multiplier and the shift values in non-tape mode:
  

  pennylane/pennylane/operation.py

  Line 644 in 6d489ea

  multiplier *= var_mult

  
  Therefore, it made sense to choose a mutable collection rather than a tuple.
• In the CV case, we have several gates, which define gradient recipes for multiple parameters. These are now made more explicit as for example grad_recipe = ([[multiplier, a, shift], [-multiplier, a, -shift]], None), with [[multiplier, a, shift], [-multiplier, a, -shift]] being the recipe for the first, None being the recipe for the second parameter.
• Restricted the order-2 version of the CV gradient rule so that we can only have two shifts.

Benefits:

• More complex gradient recipes can be expressed.
• Gradient recipes for controlled rotations and quantum channels can be expressed.

Possible Drawbacks:

• More complex gradient recipes.

Related GitHub Issues:
N/A

TODO:

• Update grad_recipe attributes to have contain pos_multiplier, pos_shift, neg_multiplier, neg_shift
• Update get_parameter_shift
• Update docs & docstrings explaining shifts
• Update developers guide
• Add tests for new edge cases where a NotImplementedError is raised

[codecov]

Codecov Report

Merging #909 (a9a8122) into master (5e01b7f) will decrease coverage by 0.01%.
The diff coverage is 97.82%.

@@            Coverage Diff             @@
##           master     #909      +/-   ##
==========================================
- Coverage   97.89%   97.88%   -0.02%     
==========================================
  Files         147      147              
  Lines       10069    10100      +31     
==========================================
+ Hits         9857     9886      +29     
- Misses        212      214       +2     

Impacted Files	Coverage Δ
pennylane/tape/qnode.py	98.74% <ø> (ø)
pennylane/tape/tapes/jacobian_tape.py	99.23% <ø> (ø)
pennylane/qnodes/cv.py	97.56% <91.66%> (-0.75%)	⬇️
pennylane/tape/tapes/cv_param_shift.py	98.71% <95.00%> (-0.63%)	⬇️
pennylane/operation.py	95.60% <100.00%> (+0.11%)	⬆️
pennylane/ops/cv.py	100.00% <100.00%> (ø)
pennylane/qnodes/qubit.py	98.79% <100.00%> (-0.01%)	⬇️
pennylane/tape/tapes/qubit_param_shift.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 5e01b7f...a9a8122. Read the comment docs.

[josh146]

Thanks @antalszava! I think this is going in the right direction, but I think it needs to be even more general than this.

For example, rather than thinking about positive shifts and negative shifts, we want to simply provide a framework for defining gradient recipes that are arbitrary linear combinations.

For example, say a gate has the gradient recipe ∂f(x) = 0.1*f(x+5) - 0.65*f(x-pi) + 1.2*f(x+0.2). We want to represent each term in the linear combination, so the gradient recipe would look like this:

grad_recipe = [(0.1, 5), (-0.65, -pi), (1.2, 0.2)]

That is, the grad recipe should be of the form [(c0, s0), (c1, s1), (c2, s2), ...], where c_i is the coefficient of the ith term, and s_i is the shift of the ith term. Notably, the length of the grad_recipe list is flexible; a gate might have 2 terms, 3 terms, or more.

Question: we could go even further, and support grad recipes of the form ∂f(x) = Σi c_i f(a_i x + b_i). This is even more flexible, as it allows us to represent terms that do not depend on x:

# corresponds to 0.1*f(x+5) - 0.65*f(2x-pi) + 1.2*f(0.2)
grad_recipe = [(0.1, 1, 5), (-0.65, 2, -pi), (1.2, 0, 0.2)]

Another question: Do we want the grad recipe to be a function?

For example, if a gradient recipe has a degree of freedom, we could easily represent this:

# corresponds to [f(x+s) - f(x-s)] / 2sin(s)
grad_recipe = lambda s: [(1/(2*sin(s)), 1, s), (1/(2*sin(s)), 1, -s)]

[antalszava]

The workaround proposed by numpy/numpy#7225 (comment).

[josh146]

🤔 I'm guessing this is because c has a dtype that is causing casting issues?

Would it be better to simply cast/ensure that c and shifted are the correct dtype?

[antalszava]

Hi @josh146, thanks for the suggestions! Eventually went with

# corresponds to 0.1*f(x+5) - 0.65*f(2x-pi) + 1.2*f(0.2)
grad_recipe = [(0.1, 1, 5), (-0.65, 2, -pi), (1.2, 0, 0.2)]

After checking with Tom and Juan Miguel, this would be motivated by works with quantum channels. There, we'd like to have the ability to specify something like (1, 0, 0.2) to compute 1*f(0.2).

Notes on the current version:

• Made the elements of grad_recipe lists. The reason is that we may potentially be post-processing the multiplier and the shift values in non-tape mode:
  

  pennylane/pennylane/operation.py

  Line 644 in 6d489ea

  multiplier *= var_mult

  
  Therefore, it made sense to choose a mutable collection rather than a tuple.
• In the CV case, we have several gates, which define gradient recipes for multiple parameters. These are now made more explicit as for example grad_recipe = ([[multiplier, a, shift], [-multiplier, a, -shift]], None), with [[multiplier, a, shift], [-multiplier, a, -shift]] being the recipe for the first, None being the recipe for the second parameter.
• Restricted the order-2 version of the CV gradient rule so that we can only have two shifts. Was not quite sure, if this would be applicable for further terms.

Updating the docstrings, the documentation as well as revisiting parts of the code for refactoring purposes remain as TODO (thought first it's worth for you to see if this would something that fits our case).

[antalszava]

The coverage report of the files where there was addition/deletion:

pennylane/operation.py                                                    411      5    99%   352, 551-553, 870
pennylane/ops/cv.py                                                       252      0   100%
pennylane/qnodes/cv.py                                                    120      1    99%   85
pennylane/qnodes/qubit.py                                                 166      2    99%   378, 429
pennylane/tape/tapes/cv_param_shift.py                                    167      0   100%

The missing lines do not seem to be affected by this PR and Codecov doesn't seem to list where the drop would have occured.

[josh146]

Very nice way of putting this! I find it very clear to read and understand.

[antalszava]

Thanks so much for the comments @trbromley! Updated the formulae and further docs based on your suggestions. Thanks, @josh146 for merging the qubit case in here!

[josh146]

💯

[josh146]

Suggested change

	* Updated how gradient recipes are stored for operations, allowing for
	gradient recipe definitions involving custom multipliers and with more than
	two terms.
	[(#909)](https://github.com/PennyLaneAI/pennylane/pull/909)
	* Updated how parameter-shift gradient recipes are defined for operations, allowing for
	gradient recipes that are specified as an arbitrary number of terms.
	[(#909)](https://github.com/PennyLaneAI/pennylane/pull/909)
	Previously, `Operation.grad_recipe` was restricted to two-term parameter-shift formulas.
	With this change, the gradient recipe now contains elements of the form
	$[c_i, a_i, s_i]$, resulting in a gradient recipe of
	$$\frac{\partial}{\partial\phi_k}f(\phi_k) = \sum_{i} c_i * f(a_i * \phi_k+s_i).$$
	As this is a breaking change, all custom operations with defined gradient recipes must be
	updated to continue working with PennyLane 0.13. Note though that if `grad_recipe = None`, the
	default gradient recipe remains unchanged, and corresponds to the two terms $[c_0, a_0, s_0]=[1/2, 1, \pi/2]$
	and $[c_1, a_1, s_1]=[-1/2, 1, -\pi/2]$ for every parameter.

[josh146]

Thought it was best to add some more details! Might have to double check that the math renders correctly.

[antalszava]

Thanks! Checked and updated.

Update: enclosing with $$ didn't seem to work, added it with <img src="https://render.githubusercontent.com/render/math?math=">.

[trbromley]

Awesome work, thanks @antalszava 💯