Take wire ordering in the wires argument to qml.density_matrix into account

[dwierichs]

Context:
#878 added the functionality to return reduced density matrices by introducing qml.density_matrix(wires=...) as a return object in a QNode. While the reduction by tracing out the correct wires is respected, the ordering of the given wires is not taken into account.
The following code therefore produces the same result for wires=[0, 1] and wires=[1, 0], although it should differ:

import pennylane as qml

wires = [0, 1]
#wires = [1, 0]
dev = qml.device("default.qubit", wires=2)
@qml.qnode(dev)
def circuit():
    qml.Hadamard(0)
    return qml.density_matrix(wires=wires)

circuit()

Description of the Change:
This PR takes the wire ordering into account as well by applying pennylane.math.matrix_manipulation.py::_permute_dense_matrix to the computed density matrix in pennylane.math.quantum.py::reduced_dm.

In addition, this PR raises an error when broadcasting is detected in reduced_dm, because this is not supported and currently crashes in some cases somewhat incomprehensibly.

Benefits:
Bug fix of the above. Transparency of broadcasting scenario that is not supported anyways.

Possible Drawbacks:

Related GitHub Issues:

[codecov]

Codecov Report

Merging #4072 (521fa52) into master (8af8514) will decrease coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #4072      +/-   ##
==========================================
- Coverage   99.77%   99.77%   -0.01%     
==========================================
  Files         340      340              
  Lines       30565    30564       -1     
==========================================
- Hits        30496    30495       -1     
  Misses         69       69              

Impacted Files	Coverage Δ
pennylane/math/quantum.py	100.00% <100.00%> (ø)

[mudit2812]

Thanks for the changes @dwierichs and for the rigorous testing. I have concerns about the condition that checks for broadcasting, but other than that, this looks good!

[mudit2812]

I'm happy to approve this PR in it's current state. One way of avoiding overriding QubitDevice.density_matrix is to check what type of device is being used in QubitDevice.density_matrix itself:

input_is_dm = self.short_name[-5:] == "mixed"
return qml.math.reduced_dm(state, indices=wires, c_dtype=self.C_DTYPE, input_is_dm=input_is_dm)

I'm hopeful this change wouldn't interfere with plugins as I'm not sure if any of the plugin devices support state measurements, but it's probably something that should be checked before my suggested change is potentially implemented.

[eddddddy]

Looks good @dwierichs, thanks for this 💯

[eddddddy]

Thanks for separating out the broadcasting logic! Happy to approve now, but this is conditional on @albi3ro's comment being addressed.

[dwierichs]

@albi3ro , could you check whether I understood you correctly and implemented the checks you were thinking of? :)

Thanks @eddddddy for the quick reviews and discussions!

[albi3ro]

Looks good to me, and can be merged as is.