Refactor the tape.jacobian method to batch execute tapes

[mariaschuld]

Note: the below is written by @josh146

Context:

In the current core, the gradient methods freely intersperse quantum computation and classical processing. That is, the quantum device will be executed internally and directly whenever a quantum result is needed.

It is advantageous to instead split up the quantum and classical processing:

1. Generate all the quantum tapes corresponding to quantum processing required

2. Execute all tapes

3. Classically post-process the evaluated tape results.

This allows us finer control over the quantum execution --- for example, if a device supports execution of multiple circuits at once.

Description of changes:

• The Operation class now has the __copy__ special method defined. This allows us to ensure that, when a shallow copy is performed of an operation, the mutable list storing the operation parameters is also shallow copied. Both the old operation and the copied operation will continue to share the same parameter data,

  >>> import copy
  >>> op = qml.RX(0.2, wires=0)
  >>> op2 = copy.copy(op)
  >>> op.data[0] is op2.data[0]
  True

  however the list container is not a reference:

  >>> op.data is op2.data
  False

  This allows us to modify the parameters of the copied operation, without mutating
  the parameters of the original operation:

  >>> op2.data[0] = 1
  >>> op
  RX(0.2, wires=[0])
  >>> op2
  RX(1, wires=[0])

• Likewise, the Tensor.__copy__ method is defined to make sure that the obs list is shallow copied

• Likewise, the MeasurementProcess.__copy__ method is defined to make sure that the obs list is shallow copied

• The QuantumTape.copy method has been tweaked so that:

  • Optionally, the tape's operations are shallow copied in addition to the tape by passing the copy_operations=True boolean flag. This allows the copied tape's parameters to be mutated without affecting the original tape's parameters. (Note: the two tapes will share parameter data until one of the tapes has their parameter list modified.)

  • Copied tapes continue to share the same caching dictionary as the original tape.

  • Copied tapes can be cast to another QuantumTape subclass by passing the tape_cls keyword argument.

• All gradient methods, including finite-differences, the qubit/CV parameter-shift rules, and reversible differentiation, have been modified. Rather than returning the evaluated gradients directly, they now return a tuple containing the required quantum and classical processing steps:

  def gradient_method(idx, param, **options):
      # generate the quantum tapes that must be computed
      # to determine the quantum gradient
      tapes = quantum_gradient_tapes(self)
  
      def processing_fn(results):
          # perform classical processing on the evaluated tapes
          # returning the evaluated quantum gradient
          return classical_processing(results)
  
      return tapes, processing_fn

  This is similar in structure to TensorFlow's custom_gradient API.

  In addition, we make use of the ability to cast copied tapes to different QuantumTape subclasses to ensure that any copied tape created by the gradient methods no longer has an interface. This is because the we are currently behind the interface boundary; we can safely assume that all tape parameters have already been unwrapped where needed.

• Finally, the tape.jacobian() method has been modified to follow the following structure:

  • Loop over all parameters, and accumulate the (a) quantum tapes, and (b) classical post-processing functions required.

  • Evaluate the accumulated quantum tapes

  • Apply the post-processing functions to the evaluated tape results

Benefits:

• Quantum and classical processing is cleanly separated; in the future, we can take advantage of devices that support batched execution.

• Tapes can now be safely shallow copied and parameters mutated without affecting all tapes.

• The new form of the gradient methods suggests a rough API that we might want to standardize at some point, to make it easier for developers to contribute new gradient methods.

• When the tapes are copied, we are essentially doing two shallow copies; a shallow copy of the tape structure, and a shallow copy of the operation list. This is much more efficient than attempting to do a recursive deep copy.

• The reversible method is now very cleanly written as a tape transform with state output, rather than a list of operations and calls to private device methods.

• The gradient methods are now written in the form tape -> tape(s) + classical processing. In the future, we might want to elevate them as user-facing QNode transforms, for use by advanced users that want to see the generated quantum gradient tapes.

Drawbacks:

• While removing the device/execution specific logic from the gradient methods was generally pretty smooth, both the CV parameter-shift and reversible backprop methods continue to require access to the device.

  • The CV parameter shift rule requires the device to ensure that the PolyXP observable is supported if second order shifts are required. If not, it falls back to finite-differences.

  • The reversible method requires access to the device's pre-rotated statevector.

  Both methods also require the device wire labeling, as they manipulate observables.

  Currently, we get around this by having the tape.jacobian method simply pass the device as a keyword argument to the gradient methods. We proposing leaving this for now; while inelegant, there is no clear and quick solution. (A long term solution would perhaps be to have the gradient methods register the operations/conditions they support.)

• The gradient methods continue to act on single parameters in isolation, with tape.jacobian() performing a loop over parameters and attempting to compute fallback methods. This is not ideal, for two reasons:

  • If a particular gradient method has a weird constraint that determines fallback (e.g., the CV param shift rule), the method needs to take it into account, or override tape.jacobian(). This spreads logic for fallback between both tape.jacobian and the gradient method.

  • More importantly, a lot of gradient methods require knowledge from previous parameter gradient calculations; for example, where the gradient of parameter n can be made more efficient by re-using info from the calculation for parameter n-1, n-2, etc.. The current approach, where gradient methods are computed in isolation per gradient method call, does not allow for this.

    First-order finite differences (requiring a pre-execution to determine y0) and the variance parameter-shift methods, are two current examples that are hampered by the existing approach.

  Thus, there is flexibility and efficiency gains in allowing gradient method to generate tapes for all parameters, rather than single parameters in isolation. Of course, a downside is that we lose the ability for users to fine-grain the method manually per parameter. Although, I'm not so convinced this is a common/wanted use-case; it is more likely that users want to fine-grain the method per gate, not per parameter.

[codecov]

Codecov Report

Merging #840 into master will increase coverage by 0.02%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master     #840      +/-   ##
==========================================
+ Coverage   97.79%   97.82%   +0.02%     
==========================================
  Files         137      137              
  Lines        9207     9358     +151     
==========================================
+ Hits         9004     9154     +150     
- Misses        203      204       +1     

Impacted Files	Coverage Δ
pennylane/tape/qnode.py	98.17% <ø> (-0.61%)	⬇️
pennylane/_device.py	96.00% <100.00%> (+0.14%)	⬆️
pennylane/_qubit_device.py	98.87% <100.00%> (+0.04%)	⬆️
pennylane/operation.py	95.69% <100.00%> (+0.24%)	⬆️
pennylane/tape/measure.py	97.14% <100.00%> (+0.26%)	⬆️
pennylane/tape/tapes/cv_param_shift.py	99.34% <100.00%> (+0.13%)	⬆️
pennylane/tape/tapes/jacobian_tape.py	99.21% <100.00%> (+0.23%)	⬆️
pennylane/tape/tapes/qubit_param_shift.py	100.00% <100.00%> (ø)
pennylane/tape/tapes/reversible.py	100.00% <100.00%> (ø)
pennylane/tape/tapes/tape.py	99.25% <100.00%> (+0.04%)	⬆️
... and 1 more

---

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 939b80f...8250b54. Read the comment docs.

[mariaschuld]

But that now makes sense, I was wondering why the device is used.

One thing I was wondering too is hardware compatibility. Since we can always write a qubit circuit with PaulZ observables (by rotating the basis), and PauliZ is unitary, can't we implement reversible diff by applying the original unitary, observable, and diffed unitary inverse and measure the overlap with the zero state?

[co9olguy]

I wonder if we'll ever hit a point where this nonstandard copying behaviour impacts developers or power users?

[josh146]

Good question. By default, tape.copy() will perform a single shallow copy --- as users or developers would expect. It is only when copy_operations=True (so something the user/developer will have to explicitly request) does the second shallow copy of operations take place. Since this is not default behaviour, I'm not too worried.

I'm more worried about the non-standard behaviour of deep copying. Users/developers expect deep copy to copy everything; however the parameters are not copied (since PyTorch doesn't allow this). I think fixing this is a good thing; it's easy to get everything working with copy.deepcopy(tape), and then realize only later that while autograd and TF work, Torch doesn't.

Perhaps the deep copy could try and copy the parameters, and only skip them if an exception occurs (e.g., Torch)?

[co9olguy]

Awesome tests 💪

[trbromley]

Is there an issue/forum post to see more context?

[trbromley]

would this be 2h rather than h?

[josh146]

No, should be h here, because the shift is +- h/2 😆

I have to admit, I did a double take, since I'm so used to seeing 2h in the denominator of centered finite differences.