[BUG] Restore ability to have observables returning types other than float about pennylane

[josh146]

Thanks @cvjjm for catching this! Something else that I think is important to identify: why the existence of the test for this feature did not catch this regression. My suspicion is that the test is not being run 🤔

[astralcai]

Thanks @cvjjm for catching this! Something else that I think is important to identify: why the existence of the test for this feature did not catch this regression. My suspicion is that the test is not being run 🤔

The test is being run. The reason why the current test case is passing is in how the device is defined. In the current version of the test case, the test device overrides the _asarray method (line 1432 of test_autograd.py).

@staticmethod
def _asarray(arr, dtype=None):
    # pylint: disable=unused-argument
    return arr

which bypasses the issue.

[albi3ro]

Thanks for the issue.

I see two main differences that have changed. I think PennyLane will actually currently be better at supporting arbitrary results now. We used to always smash all the results into a ragged array, but now we can handle any type of result object.

You can add a custom measurement, with any custom return, by inheriting from qml.measurements.StateMeasurement and/or qml.measurements.SampleMeasurement. See documentation on this here

Then with DefaultQubit, you can simply do something like:

class SpecialObject:
    def __init__(self, val):
        self.val = val

    def __repr__(self):
        return f"SpecialObject({self.val})"

class CustomMeas(qml.measurements.StateMeasurement):

    def process_state(self, state, wire_order):
        return SpecialObject(state) # or some other postprocessing of the state

@qml.qnode(qml.device('default.qubit'), interface=None)
def circuit():
    qml.PauliX(0)
    return CustomMeas()

circuit()

SpecialObject([0. 1.])

If you want to maintain your own device and not just extend DefaultQubit, you should simply need to make sure execute (or old interface batch_execute) returns an object of the right type.

Update:

I've also managed to get the Legacy Interface code to mostly work by swapping qml.expval with a CustomExpval as well.

class CustomExpval(qml.measurements.StateMeasurement):
    
    def process_state(self, state, wire_order):
        expval = qml.expval(qml.PauliZ(0)).process_state(state, wire_order)
        return SpecialObject(expval)

@qml.qnode(dev, diff_method="parameter-shift")
def qnode(x):
    qml.RY(x, wires=0)
    return CustomExpval(obs=SpecialObservable(wires=0))

We currently assume that the result of an ExpectationMP is a float of either size zero or with a batch dimension. By having a different measurement process, we can circumvent that assumption.

[cvjjm]

Nice! That was super useful. I will first try to get it to work with v0.30 with overwriting _asarray() and then work on compatibility of our devices with the new device interface. I will close this for now.