Gradients of non-linear post-processing functions

Context We have been stumbling over this problem during the classical shadows PRs https://github.com/PennyLaneAI/pennylane/pull/2871 as well as the counts https://github.com/PennyLaneAI/pennylane/issues/2932 . I believe it is also related to this older issue https://github.com/PennyLaneAI/pennylane/issues/1722.

Here is an abstract example where a qnode is executed and returns a value $g(x)$, which is then post-processed in $f(g)$.

@qml.qnode(dev)
def qnode(x):
    ..
    qml.RX(x, 0)
    ..
    return g

def f(g):
    return non-linear(g)

When derivative of $f$ is given by $\partial_x f(g(x)) = f’(g(x)) \partial_x g(x) = f’(g(x)) (g(x+\pi/2) - g(x-\pi/2))/2$.

This is also what PennyLane is doing. However, there are situations where I want to be able to do something different. For example, in classical shadows, $g$ are some abstract bits that inform the classical representation of the snapshot states. From this I can post-process the expectation value, $f$, which is a non-linear function of $g$, the bits and recipes (measurement choices).

When I call qml.grad(f)(x), PennyLane does the above logic (I am actually not sure what it computes for the gradients of the bits and recipes). But what I would want instead is

$\partial_x f = (f(g(x+\pi/2)) - f(g(x-\pi/2)))/2$

Solutions One solution is to perform the post-processing inside the qnode, i.e. registering a new measurement. As far as I understand, by default, calling grad on this measurement outcome does exactly what I want.

Instead of registering new measurements, it would be great to have the possibility to register custom functions as outputs of qnodes. That custom function would internally call a valid measurement, and post-process it to the desired output, which is then returned by the qnode.

A different solution would be to register custom vjps, but as far as I understand, this is not something very desirable.

Disclaimer While writing this I realized a lot of false conceptions on my end, so I am not 100% sure I am making sense. Please understand this issue more as a discussion forum rather than a contained and explicit feature request.

[josh146]

Instead of registering new measurements, it would be great to have the possibility to register custom functions as outputs of qnodes. That custom function would internally call a valid measurement, and post-process it to the desired output, which is then returned by the qnode.

I think a solution where we make it really easy for users and developers to define their own measurement process, which can be automatically used by devices without needing to modify devices, is the key thing we are missing

[josh146]

The idea of registering callables instead of MeasurementProcesses in QNodes is pretty cool, but I don’t think it would be very maintainable…

@dwierichs I’m curious why you think this? To me, this feels like the most natural solution, especially if we build it into the device API.

For a long time, we have actually wanted to:

• Move the default implementations of expval, var, etc. out of QubitDevice and into the corresponding MeasurementProcess

• Allow devices to choose whether they use their own expval, var, etc. logic, or utilize the logic provided by a MeasurementProcess.

The advantage here is that any developers/user can define their own measurement process and have it be used by any device, without needing to modify the device. They also get gradient support via calling some_autodiff_framework.grad(...) for free, even if they use NumPy* 🙂 Currently, any time we want to add new measurement support to PL, we have to manually add it to every device we want to support as a new device method, which is a big overhead.

* Assuming their output of the measurement process supports the parameter-shift rule.