[BUG] `qml.sample` with no observable fails for shot vectors

[antalszava]

Expected behavior

qml.sample with no observable works well with shot vectors.

Actual behavior

Error raised.

Additional information

A more bigger scope change is upcoming, that will likely contain a bug fix to this.

Source code

import pennylane as qml

dev = qml.device("default.qubit", wires=3, shots=[10, 1000])

@qml.qnode(device=dev)
def circuit(x):
    qml.Hadamard(wires=[0])
    qml.CRX(x, wires=[0, 1])
    return qml.sample(wires=[0])

circuit(0.5)

Tracebacks

~/pennylane/_qubit_device.py in sample(self, observable, shot_range, bin_size, counts)
   1062             ]
   1063         return (
-> 1064             samples.reshape((3, bin_size, -1))
   1065             if no_observable_provided
   1066             else samples.reshape((bin_size, -1))

ValueError: cannot reshape array of size 10 into shape (3,10,newaxis)

System information

PennyLane master branch

Existing GitHub issues

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

[mudit2812]

I think this can be addressed by doing the following:

In pennylane/_qubit_device.py:1058, make the following changes:

        num_wires = len(device_wires) if len(device_wires) > 0 else self.num_wires
        if counts:
            shape = (-1, bin_size, num_wires) if no_observable_provided else (-1, bin_size)
            return [
                _samples_to_counts(bin_sample, no_observable_provided)
                for bin_sample in samples.reshape(shape)
            ]
        return (
            samples.reshape((num_wires, bin_size, -1))
            if no_observable_provided
            else samples.reshape((bin_size, -1))
        )

instead of

        if counts:
            shape = (-1, bin_size, 3) if no_observable_provided else (-1, bin_size)
            return [
                _samples_to_counts(bin_sample, no_observable_provided)
                for bin_sample in samples.reshape(shape)
            ]
        return (
            samples.reshape((3, bin_size, -1))
            if no_observable_provided
            else samples.reshape((bin_size, -1))
        )

Hardcoding the shape of the samples when raw samples are returned to be (3, bin_size, -1) only works when the number of wires being sampled is 3. To this effect, the testing suite in tests/test_measurements.py:TestCounts should also be expanded to include tests where varying numbers of wires are being sampled when counts=True.

Once I made the changes I proposed locally, I found another related bug. I noticed that only the counts for the first item in the shot vector show. An example is shown below:

>>> dev = qml.device("default.qubit", wires=4, shots=[10, 5])
>>> @qml.qnode(dev)
... def test():
...     qml.PauliX(wires=0)
...     qml.PauliY(wires=1)
...     qml.CNOT(wires=[1, 2])
...     return qml.sample(counts=True)
>>> test()
{'1110': 10}

Upon following the trace, I found that the output was being changed at pennylane/qnode.py:653:

        if self._qfunc_output.return_type is qml.measurements.Counts:
            # return a dictionary with counts not as a single-element array
            return res[0]

Since only the first item in res is returned, the dictionaries for all shot vectors except for the first one are removed.

Hope this is helpful!

[antalszava]

Note that this was resolved as per Mudit's points. Thanks for that!

The code runs without errors, it does warn VisibleDeprecationWarning: Creating an ndarray from ragged nested sequences that should be cleared by the new return types system.