pennylane icon indicating copy to clipboard operation
pennylane copied to clipboard

`QNode` defers diff method validation to the device

Open albi3ro opened this issue 1 year ago • 4 comments

[sc-56106]

Context:

The new device interface allows devices to support a variety of device gradient methods, each with different requirements on the circuit. For example, backprop is considered "best" for default.qubit if we have finite shots and no expectation values of SparseHamiltonian. Adjoint will be considered "best" for lightning.qubit if the circuit is analytic and contains all expectation values. We can also imagine other devices with different types of differentiation, and different restrictions on each one.

Historically, all these restrictions were hardcoded in various private methods in the QNode, such as QNode._validate_backprop_method, QNode._validate_device_method, and QNode._validate_adjoint_method. But this architecture is very rigid and hard to extend. It was also fully determined on QNode initialization time, and difficult to adjust to changing circuits and shots.

Description of the Change:

This PR changes the following QNode methods to dynamically depend on the tape and defer to the qml.devices.Device interface when possible

  • QNode._update_gradient_fn
  • QNode.get_gradient_fn
  • QNode.get_best_method

The private validation methods are now only applicable to the qml.Device interface, and ValueError's are raised with the new device interface:

  • QNode._validate_backprop_method
  • QNode._validate_adjoint_method
  • QNode._validate_device_method

QNode._update_gradient_fn is now run on each call, instead of on initialization, on overridden shots, and interface update. Making sure this got ran in the correct place forced several other pieces of code to move around. This also causes backprop device switching to occur more often, forcing a few test changes. I'm not too concerned about these changes, as default.qubit.legacy is now a legacy device.

Error changes:

Some former DeviceError's are now QuantumFunctionError's, and some string outputs are a little different.

Other Minor Changes:

Now default.qubit includes validation for backprop and qml.SparseHamiltonian. The QNode now only needs to validate this for the legacy devices.

I moved the hybrid transform program classical component computation onto the TransformProgram, to minimize the number of moving pieces in QNode.__call__.

Benefits:

Possible Drawbacks:

albi3ro avatar Feb 07 '24 23:02 albi3ro

Codecov Report

All modified and coverable lines are covered by tests :white_check_mark:

Project coverage is 99.63%. Comparing base (a4f0fff) to head (a2f1998). Report is 1 commits behind head on master.

Additional details and impacted files
@@            Coverage Diff             @@
##           master    #5176      +/-   ##
==========================================
- Coverage   99.63%   99.63%   -0.01%     
==========================================
  Files         400      400              
  Lines       37156    36880     -276     
==========================================
- Hits        37022    36745     -277     
- Misses        134      135       +1     

:umbrella: View full report in Codecov by Sentry.
:loudspeaker: Have feedback on the report? Share it here.

codecov-commenter avatar Feb 09 '24 22:02 codecov-commenter

Hello. You may have forgotten to update the changelog! Please edit doc/releases/changelog-dev.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.

github-actions[bot] avatar Feb 09 '24 22:02 github-actions[bot]

Do we expect any performance overhead from doing validation on each call?

trbromley avatar Feb 12 '24 18:02 trbromley

Do we expect any performance overhead from doing validation on each call?

So for an extremely small circuit where classical overheads will play more of a role:

@qml.qnode(qml.device('default.qubit'))
def circuit(x):
    qml.RX(x, wires=0)
    return qml.expval(qml.PauliZ(0))

x = qml.numpy.array(0.5)

%snakeviz circuit(x)

We are spending 2.61% of the time in QNode._update_gradient_fn. Though I would note that 1.71% is spent constructing the TransformProgram that we immediately throw away. If we were concerned about the time to do this, I would recommend splitting the transform program construction from the diff method calculation. But I'm not particularly concerned about those milliseconds.

Screenshot 2024-02-12 at 4 28 25 PM

albi3ro avatar Feb 12 '24 21:02 albi3ro