Fix bug where a QNode with a single probability output was not correctly differentiated

[josh146]

Context:

QNodes with a single probability output, of the form

@qml.qnode(dev)
def circuit(x, y):
    qml.RX(x, wires=[0])
    qml.RY(y, wires=[1])
    qml.CNOT(wires=[0, 1])
    return qml.probs(wires=[1])

would result in the following error when differentiated in tape-mode using the parameter-shift rule:

  File "/home/pennylane/tape/tapes/qubit_param_shift.py", line 172, in processing_fn
    return np.dot(coeffs, results)
  File "<__array_function__ internals>", line 6, in dot
ValueError: shapes (2,) and (2,1,2) not aligned: 2 (dim 0) != 1 (dim 1)

This is due to the fact that we are using np.dot(coeffs, results) as a shorthand for [c*r for c, r in zip(coeffs, results)]. However, in the case of returning a single probability array, NumPy attempts to broadcast the dot product instead.

Note that this bug does not affect the case return qml.probs(wires=0), qml.probs(wires=1) (e.g., when multiple probability arrays are returned), as np.dot behaviour coincides again with [c*r for c, r in zip(coeffs, results)].

Description of the Change:

Changes the line to np.dot(coeffs, np.squeeze(results)), to avoid broadcasting issues.

Benefits:

The use case specified above now works.

Possible Drawbacks: n/a

Related GitHub Issues: n/a

[codecov]

Codecov Report

Merging #1007 (4d31820) into master (1fa2bc4) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #1007   +/-   ##
=======================================
  Coverage   97.89%   97.89%           
=======================================
  Files         151      151           
  Lines       11059    11059           
=======================================
  Hits        10826    10826           
  Misses        233      233           

Impacted Files	Coverage Δ
pennylane/tape/tapes/qubit_param_shift.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 1fa2bc4...4d31820. Read the comment docs.

[antalszava]

How does this change affect the hotfix in the tape QNode? Would that still be required?

[josh146]

Yep, that will still be required. Note that the squeeze here is functionally independent of the one in the QNode.

• The one in QNode.__call__ is designed so that the QNode output matches the users return function. E.g., if the user specifies return qml.expval(qml.PauliZ(0)), squeezing occurs so that a float is returned. If instead the user specifies return [qml.expval(qml.PauliZ(0))] or return qml.expval(qml.PauliZ(0)), qml.var(qml.PauliX(0)), then no squeezing occurs and a list and/or tuple is returned.

• The squeezing here is low level, and not connected to the QNode return statement. Here, we are simply ensuring that np.dot() behaves as we need it to.

[antalszava]

It seemed that squeezing is slightly different for each interface (autograd, jax and Torch/TF) as per the previously linked hotfix, could it be worth parametrizing based on the interface?

[josh146]

The reason why the QNode must use the ML framework for squeezing, and here we don't, is simply due to where we are with respect to the interface.

• In QNode.__call__, we are outside the interface. All array/tensor manipulation must be done via the correct ML framework to ensure differentiability is not broken.

• In tape.parameter_shift, we are inside the interface. Inside the interface, no ML frameworks are used; everything has been unwrapped into a NumPy array. Hence why we can use np.dot, and don't have to worry about torch.matmul, tf.tensordot, etc 🙂

[josh146]

Tangent: now that we have a maths module, we should add qml.math.squeeze, so that we can remove the if statements in QNode.__call__ to support squeezing in autograd/torch/tf.

[antalszava]

Looks good to me! 💯 Thanks for the details.

[thisac]

Thanks for fixing this quickly @josh146!

[josh146]

Note: @antalszava has pointed out that we could fix this 'closer to the source', by instead squeezing the linked line in QubitDevice.

However, this would require framework-independent squeezing, which necessitates creating qml.math.squeeze. Since this is a lot of extra work, and this bug needs to be fixed to allow us to port the demos to tape mode, we should merge this in as a temp fix, and a new issue #1009 has been created for the low-level hotfix.