[BUG] `qml.state` does not output right results with `shots_vector`

[rmoyard]

Expected behavior

dev = qml.device("lightning.qubit", wires=nwires, shots=[(10, 1), (10, 2)])

@qml.qnode(dev)
def circuit():
    return qml.state()

It should output

[[1.+0.j 0.+0.j 0.+0.j 0.+0.j]
 [1.+0.j 0.+0.j 0.+0.j 0.+0.j]
 [1.+0.j 0.+0.j 0.+0.j 0.+0.j]]

Also

dev = qml.device("lightning.qubit", wires=nwires, shots=[(10, 2), (10, 2)])

@qml.qnode(dev)
def circuit():
    return qml.state()

it should output

[[1.+0.j 0.+0.j 0.+0.j 0.+0.j]
 [1.+0.j 0.+0.j 0.+0.j 0.+0.j]
 [1.+0.j 0.+0.j 0.+0.j 0.+0.j]
 [1.+0.j 0.+0.j 0.+0.j 0.+0.j]]

Actual behavior

The first QNode outputs an error, see thee traceback below.

The second QNode outputs wrong results

[1.+0.j 0.+0.j 0.+0.j 0.+0.j 1.+0.j 0.+0.j 0.+0.j 0.+0.j]

Additional information

No response

Source code

dev = qml.device("lightning.qubit", wires=nwires, shots=[(10, 1), (10, 2)])

@qml.qnode(dev)
def circuit():
    return qml.state()

Tracebacks

File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/qnode.py", line 581, in __call__
    res = qml.execute(
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/interfaces/execution.py", line 378, in execute
    res = _execute(
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/interfaces/autograd.py", line 64, in execute
    return _execute(
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/penv/lib/python3.9/site-packages/autograd/tracer.py", line 48, in f_wrapped
    return f_raw(*args, **kwargs)
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/interfaces/autograd.py", line 108, in _execute
    res, jacs = execute_fn(tapes, **gradient_kwargs)
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/interfaces/execution.py", line 135, in wrapper
    res = fn(execution_tapes.values(), **kwargs)
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/interfaces/execution.py", line 87, in fn
    return original_fn(tapes, **kwargs)
  File "/usr/local/Cellar/python@3.9/3.9.7/Frameworks/Python.framework/Versions/3.9/lib/python3.9/contextlib.py", line 79, in inner
    return func(*args, **kwds)
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/_qubit_device.py", line 345, in batch_execute
    res = self.execute(circuit)
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/_qubit_device.py", line 291, in execute
    results = qml.math.stack(results)
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/math/multi_dispatch.py", line 178, in wrapper
    return fn(*args, **kwargs)
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane/pennylane/math/multi_dispatch.py", line 495, in stack
    return np.stack(values, axis=axis, like=like)
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/penv/lib/python3.9/site-packages/autoray/autoray.py", line 84, in do
    return get_lib_fn(backend, fn)(*args, **kwargs)
  File "<__array_function__ internals>", line 5, in stack
  File "/Users/romainmoyard/PennyLaneMaster/PennyLane/penv/lib/python3.9/site-packages/numpy/core/shape_base.py", line 426, in stack
    raise ValueError('all input arrays must have the same shape')
ValueError: all input arrays must have the same shape

System information

Name: PennyLane
Version: 0.23.0.dev0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author: 
Author-email: 
License: Apache License 2.0
Location: /Users/romainmoyard/PennyLaneMaster/PennyLane/pennylane
Requires: numpy, scipy, networkx, retworkx, autograd, toml, appdirs, semantic_version, autoray, cachetools, pennylane-lightning
Required-by: PennyLane-Lightning, PennyLane-Qchem
Platform info:           macOS-10.15.7-x86_64-i386-64bit
Python version:          3.9.7
Numpy version:           1.21.0
Scipy version:           1.7.0
Installed devices:
- default.gaussian (PennyLane-0.23.0.dev0)
- default.mixed (PennyLane-0.23.0.dev0)
- default.qubit (PennyLane-0.23.0.dev0)
- default.qubit.autograd (PennyLane-0.23.0.dev0)
- default.qubit.jax (PennyLane-0.23.0.dev0)
- default.qubit.tf (PennyLane-0.23.0.dev0)
- default.qubit.torch (PennyLane-0.23.0.dev0)
- lightning.qubit (PennyLane-Lightning-0.22.0)
None

Existing GitHub issues

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

[timmysilv]

Running this example today actually fails with an error: ValueError: all input arrays must have the same shape (called from here). That said, I'm not sure what our current stance is on users providing ShotTuple-like lists as arguments, so I'm not 100% sure what the expected behaviour here is.