[BUG] output shape and type of gradient_transform applied to QNodes depend on classical processing

[dwierichs]

Expected behavior

When applying a gradient_transform to a QNode, the output shape and type should not differ whether or not there is classical processing.

Actual behavior

If there is no classical processing (and other conditions apply, such as "there is only one QNode argument"), the output shape and type differs from cases where there is classical processing:

import pennylane as qml
from jax import numpy as np
from pennylane import numpy as np

dev = qml.device("default.qubit")

@qml.qnode(dev)
def circuit(w, factor=1.):
    qml.Rot(*(factor * w), 0)
    return qml.expval(qml.PauliZ(0))

w = np.array([0.612, 0.42, 1.2])

print(qml.gradients.hadamard_grad(circuit)(w, factor=1.))
# (tensor(-0., requires_grad=True), tensor(-0.40776045, requires_grad=True), tensor(-0., requires_grad=True))
print(qml.gradients.hadamard_grad(circuit)(w, factor=2.))
# [ 0.         -1.48928624  0.        ]

Another example, this time causing a transpose, is

@qml.qnode(dev)
def circuit(w, factor=1.):
    qml.Rot(*(factor * w), 0)
    return qml.probs([0, 1])

w = np.array([0.612, 0.42, 1.2])

print(qml.gradients.hadamard_grad(circuit)(w, factor=1.))
# (tensor([0., 0., 0., 0.], requires_grad=True), tensor([-0.20388023,  0.        ,  0.20388023,  0.        ], requires_grad=True), tensor([0., 0., 0., 0.], requires_grad=True))
# -> Could be interpreted as shape (3, 4)
print(qml.gradients.hadamard_grad(circuit)(w, factor=2.))
# [[ 0.         -0.74464312  0.        ]
#  [ 0.          0.          0.        ]
#  [ 0.          0.74464312  0.        ]
#  [ 0.          0.          0.        ]]
# Has shape (4, 3)

Additional information

This is due to a special branch in _contract_qjac_with_cjac in gradients/gradient_transform.py, which is activated whenever the classical Jacobian is the identity.
While it is correct that no numerics have to be performed in this case, the contraction that would be performed without this branch (and yields the correct behaviour) also leads to transposing the output. As the contraction is skipped, the transpose is as well, leading to the observed behaviour.

Source code

No response

Tracebacks

No response

System information

pl dev; master

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.