Consistency of QNode return value shape after qml.gradients.param_shift_hessian() for QNodes that return a one element list

Feature details

For the special case of QNodes returning a list with a single expval the shape returned after transforming such QNode with qml.gradients.param_shift_hessian() is a bit “inconsistent” with how the untransformed return value is treated.

What I mean is best explained with the example below.

I think it would be nice if this minor “inconsistency” in the user interface could be fixed.

Implementation

No response

How important would you say this feature is?

1: Not important. Would be nice to have.

Additional information

For me the expected behavior would have been that all asserts pass in the following code:

import pennylane as qml
from pennylane import numpy as np

wires = 3
dev = qml.device('default.qubit', wires=wires)

def ansatz(params):
    for wire in range(wires):
        qml.Rot(*params[0][wire], wires=wire)
    qml.CNOT(wires=[0,1])
    for wire in range(wires):
        qml.Rot(*params[1][wire], wires=wire)
    qml.CNOT(wires=[1,2])

@qml.qnode(dev)
def cost1(params):
    ansatz(params)
    return qml.expval(qml.PauliZ(0))

@qml.qnode(dev)
def cost2(params):
    ansatz(params)
    return [qml.expval(qml.PauliZ(0))]

@qml.qnode(dev)
def cost3(params):
    ansatz(params)
    return [qml.expval(qml.PauliZ(0)), qml.expval(qml.PauliZ(1))]


params = np.random.rand(2, wires, 3)

expval1 = cost1(params)
expval2 = cost2(params)
expval3 = cost3(params)

assert expval1 == expval2[0]
assert expval1 == expval3[0]

hessian1 = qml.gradients.param_shift_hessian(cost1)(params)
hessian2 = qml.gradients.param_shift_hessian(cost2)(params)
hessian3 = qml.gradients.param_shift_hessian(cost3)(params)

print("hessian1.shape", hessian1.shape, "hessian2.shape", hessian2.shape, "hessian3.shape", hessian3.shape)
assert hessian1.shape == hessian3[0].shape
assert hessian1.shape == hessian2[0].shape, "For consistency with how a QNode with 'return [qml.expval(qml.PauliZ(0))] behaves when it is evaluated (it reutrns a one element tensor) it would be nice that if after transforming it with gradients.param_shift_hessian() it should also return a tensor, whose only element is the hessian. Instead it returns the plain hessian directly.'"
assert hessian1 == hessian2[0]

However hessian1.shape == hessian3.shape != hessian2[0].shape.

[cvjjm]

I would be in favor of not unpacking outputs of QNode that return a list with a single item as in return [qml.expval(...)].

The rational and use case is that of QNodes that are dynamically created by some higher level code so that the length of the list or returned expvals is not known in advance and whose output is to be processed by further code. That further code needs to get the QNode output in a consistent shape irrespective of whether the list container 4, twelve, or a single expval. It would be nice if the first dimension would always be the length of the list after return.

[dime10]

Hi @cvjjm, thank you very much for pointing out this inconsistency!

After investigating I’ve realised the problem is actually not specific to param_shift_hessian, param_shift for example runs into same issue (but the two are also inconsistent w.r.t. each other in how they behave). Essentially, param_shift_hessian will squeeze any dimensions of size 1 from the result, while param_shift squeezes in some circumstances but not in others, as shown in the table below:

QNode return statement	QNode output shape	param_shift output shape	param_shift_hessian output shape	batch execution result shape
scalar (e.g. expval)	(0,)	(1, *) !	(*, *)	(#, 1) !
single element (scalar) list/tuple (e.g. [expval])	(1,)	(1, *)	(*, *) !	(#, 1)
multi element (scalar) list/tuple (e.g. [expval, expval])	(2,)	(2, *)	(2, *, *)	(#, 2)
non-scalar (e.g. probs(0))	(2,)	(2, *)	(2, *, *)	(#, 1, 2) !
single element (non-scalar) list/tuple (e.g. [probs(0)])	(1, 2)	(2, *) !	(2, *, *) !	(#, 1, 2)
multi element (non-scalar) list/tuple (e.g. [probs(0), probs(1)])	(2, 2)	(2, 2, *)	(2, 2, *, *)	(#, 2, 2)
bold ! values deviate from expected		* = QNode args shape	* = QNode args shape	# = number of tapes

Unfortunately, I don’t think either transform is able to detect whether, for example, a single expectation value was wrapped in a list or not. It appears that the QNode output shape that the transforms receive is determined by the batched tape execution (batch_transform or maybe qml.execute?). @josh146 do you know if it would be possible to make changes such that the tape execution results produce the same QNode output dimensions as when running the QNode directly?