Fix cases for qml.sample(..., counts=True)

[antalszava]

Context:
In a recent PR it was found that the qml.shots(..., counts=True) behaves in unexpected ways in some cases. The original addition was made in this PR.

Description of the Change:

Fixed the use of returning counts when:

• For varying number of wires (other than 3);
• Using the interfaces;
• Counts are returned as part of multiple measurements.

Note that counts now no longer return tensor objects, but rather raw dictionaries or sequences that contain dictionaries.

Benefits:

Returning counts works in more cases as expected.

Possible Drawbacks:
N/A

Related GitHub Issues:
N/A

[codecov]

Codecov Report

Merging #2839 (d034c72) into master (14485da) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff            @@
##           master    #2839    +/-   ##
========================================
  Coverage   99.63%   99.63%            
========================================
  Files         251      252     +1     
  Lines       20572    20757   +185     
========================================
+ Hits        20496    20681   +185     
  Misses         76       76            

Impacted Files	Coverage Δ
pennylane/measurements.py	99.63% <ø> (ø)
pennylane/ops/qubit/observables.py	100.00% <ø> (ø)
pennylane/_qubit_device.py	99.45% <100.00%> (+<0.01%)	⬆️
pennylane/drawer/tape_mpl.py	100.00% <100.00%> (ø)
pennylane/drawer/tape_text.py	100.00% <100.00%> (ø)
pennylane/interfaces/autograd.py	100.00% <100.00%> (ø)
pennylane/interfaces/jax.py	100.00% <100.00%> (ø)
pennylane/interfaces/tensorflow.py	100.00% <100.00%> (ø)
pennylane/interfaces/torch.py	100.00% <100.00%> (ø)
pennylane/ops/op_math/__init__.py	100.00% <100.00%> (ø)
... and 2 more

---

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 52d656d...d034c72. Read the comment docs.

[rmoyard]

Great changes @antalszava, I am glad that it is working as expected 💯 I've left some comments. One thing that is out of scope of this PR is that I would find great that we always return the dictionary for every state:

{'001': 7, '111': 3} could become {'000': 0, 001': 7, '010': 0, '011': 0, '100': 0, '101': 0, '110': 0 '111': 3}

[antalszava]

@rmoyard thanks for the comments!

One thing that is out of scope of this PR is that I would find great that we always return the dictionary for every state:{'001': 7, '111': 3} could become {'000': 0, 001': 7, '010': 0, '011': 0, '100': 0, '101': 0, '110': 0 '111': 3}

Good point! Could it be user-dependent, i.e., some prefer this some prefer that? Would it be worth having kwarg for it?

[albi3ro]

@rmoyard thanks for the comments!

One thing that is out of scope of this PR is that I would find great that we always return the dictionary for every state:{'001': 7, '111': 3} could become {'000': 0, 001': 7, '010': 0, '011': 0, '100': 0, '101': 0, '110': 0 '111': 3}

Good point! Could it be user-dependent, i.e., some prefer this some prefer that? Would it be worth having kwarg for it?

That could be done with defaultdict, like used in specs.

[albi3ro]

Left some minor, optional potential improvements.

Excellent job getting this to work!

[antalszava]

[sc-23475]

[rmoyard]

Thanks @antalszava, it looks good on my side 💯

[albi3ro]

Looks good to go from my side.

[antalszava]

Created an issue for returning or states with counts:
#2864

From my side, I'm still wondering about the refactor to the QuantumTape.is_sampled attribute, but happy to create an issue for that too if we have a clear vision.