[BUG] Counts does not work with broadcasting

[rmoyard]

Expected behavior

Broadcasted parameters should work with counts.

Actual behavior

The following code returns:

Counts: {'0000011010010010100000000001000000010110000000000000000101001110100011010000000001010111001001100001': 1, '0000111100001000111010000101010001101000010001010001000110010010011000101000010001111000010001010000': 1}
Counts: {'1010011011000010100111010011101000000100001000000000001100000000110000000000100100010011000000010111': 1}

Additional information

Source code

import pennylane as qml

dev = qml.device('default.qubit', wires=2, shots=100)

@qml.qnode(dev, interface="autograd", diff_method=None)
def circuit(param):
    # These two gates represent our QML model.
    qml.RX(param, wires=0)
    qml.CNOT(wires=[0, 1])
    return qml.counts(wires=0)


counts = circuit(qml.numpy.array([1.123, 1.1]))
print(f"Counts: {counts}")
counts = circuit(qml.numpy.array([1.123]))
print(f"Counts: {counts}")

Tracebacks

No response

System information

-

Existing GitHub issues

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