Fix jacobian shape in VJP for measurements with shape and batched inputs

[majafranz]

Context:
This PR should fix a bug related to computing the gradient of a circuit which supports:

1. batching in non-trainable data
2. Measurement with a shape (i.e. probs)
3. More than one trainable parameter in the circuit

Localised example that fails, without the fix:

@qml.qnode(qml.device('default.qubit'), diff_method="parameter-shift")
def circuit(x, data):
    qml.RX(x[0], 0)
    qml.RX(x[1], 0)
    qml.RY(data, 0)
    return qml.probs(wires=0)

x = qml.numpy.array([0.5, 0.8], requires_grad=True)
data = qml.numpy.array([1.2, 2.3, 3.4], requires_grad=False)
circuit(x, data)
qml.jacobian(circuit)(x, data)

Reshaping the jacobian, when computing the VJP does resolve the issue.

Description of the Change:
Jacobian is reshaped according to the shape of dy when computing the VJP for measurements with dimension.

Benefits:
One less bug :)

Possible Drawbacks:
None

Related GitHub Issues:
This PR fixes #5979

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.65%. Comparing base (f9adf90) to head (f34fea3).
Report is 2 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5986      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         430      430              
  Lines       41544    41281     -263     
==========================================
- Hits        41404    41140     -264     
- Misses        140      141       +1     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[josh146]

Hi @majafranz, thank you for this PR/bugfix! Someone from the team will be able to review it shortly :)

[mudit2812]

Thanks for opening this PR @majafranz ! I just have one comment, and other than that, this looks great 🎉 .

Also, please ignore the one failing tensorflow test, that is unrelated to this PR and is being addressed.

[trbromley]

@Alex-Preciado who do you recommend as a second reviewer for this PR?

[dwierichs]

Hi @majafranz! Thank you for fixing this bug, it's a very good catch 🎉
I basically have two comments:

1. The test could be reduced a bit regarding its dependency on other parts of the library (QNode and VJP->Jacobian logic). I suggested an alternative, let me know what you think!
2. There seems to be a similar bug in the JVP logic. Would you be interested in fixing this as well? It can be taken care of separately, but it also would fit neatly into this PR :)

[dwierichs]

Thanks for your contribution, @majafranz 🎉 💯

[mudit2812]

@majafranz Could you fix the formatting errors preventing this PR from being merged? Thanks!