Hi there. I updated from 0.18 to 0.20 and it seems that circuits with qml.sample() outputs now throw the error InterfaceUnsupportedError: Only Variance and Expectation returns are supported for the jittable JAX interface, given ObservableReturnTypes.Sample. The example code given at the end of the documentation (https://pennylane.readthedocs.io/en/stable/introduction/interfaces/jax.html) also returns this error. I am using JAX 0.3.0 if that is of any help.
Hi @PatrickHuembeli, thanks for catching this!
I think this is due to a configuration change we made in PennyLane; if you are now jitting a QNode that returns samples, you need to explicitly specify diff_method=None.
This was done I believe because previous versions of PennyLane would assign diff_method="parameter-shift" by default, which could lead to odd behaviour and errors if attempting to differentiate the QNode.
But good catch, the documentation there needs to be updated