Add DeviceDerivatives jacobian product calculator
#4527
Conversation
Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>
Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>
Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>
Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>
Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>
Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>
albi3ro
requested review from
rmoyard
and removed request for
mudit2812 and
timmysilv
There was a problem hiding this comment.
Thanks for this PR! Two high level questions, how do we integrate adjoint_vjp and adjoint_jvp that do not use the multiplication? Second question: how do we choose between the tmethods of this class in the interfaces when they are multiple choise like (adjoint_vjp and vjp with multiplication)?
Thanks for the review :)
Each interface will prefer a different method of the class. Autograd and torch will just use Jax will use An instance of the relevant class will be created in the |
Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
This PR follows on from #4435 , and adds a child class that can handle device-provided derivatives. It has two private caches, `_results_cache` and `_jacs_cache`. This handles storing the jacobian between the forward pass and the backward pass. The caches will also store results and jacobians between successive calls to the same instance. The class can handle both device interfaces, `qml.Device` and `qml.devices.Device`. Name for the class is still up in the air. This class takes the full jacobian from the device and then performs the dot product. The class in #4637 takes the jacobian product itself from the device. --------- Co-authored-by: David Wierichs <david.wierichs@xanadu.ai> Co-authored-by: Mudit Pandey <mudit.pandey@xanadu.ai>
This PR follows on from #4527 and #4637 . Since we still need to make performance improvements to `DefaultQubit.compute_vjp`, I added an opt-in keyword argument `use_device_jacobian_product` to `qml.execute`. This allows us to easily compare between using the device jacobian product and not. Fixes #3217 --------- Co-authored-by: David Wierichs <david.wierichs@xanadu.ai> Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>
This PR follows on from #4435 , and adds a child class that can handle device-provided derivatives.
It has two private caches,
_results_cacheand_jacs_cache. This handles storing the jacobian between the forward pass and the backward pass. The caches will also store results and jacobians between successive calls to the same instance.The class can handle both device interfaces,
qml.Deviceandqml.devices.Device.Name for the class is still up in the air. This class takes the full jacobian from the device and then performs the dot product. The class in #4637 takes the jacobian product itself from the device.