[BUG] Output shape with `probs` and shot vector depends on the values of the shot vector

[dwierichs]

Expected behavior

When using a shot vector, there should be an output shape depending on the length of the shot vector, but not on the particular entries. This should hold for probs.

In particular, in the code example the outputs should have the same shape.

Actual behavior

The shape differs between a shot vector with 2 equal entries and a shot vector with differing entries.

Additional information

The issue seems to be caused by the different treatment of the two cases by range and bin_size in the shot vector logic.

Source code

import pennylane as qml
import jax
jnp = jax.numpy
jax.config.update("jax_enable_x64", True)

qml.enable_return()
for shots in ([100, 100], [100, 99]):
    dev = qml.device("default.qubit.jax", wires=1, shots=shots)

    @qml.qnode(dev, interface="jax")
    def circuit(params):
        qml.RX(params[0], 0)
        return qml.probs(wires=0)
    params = [jnp.array(0.4)]
    print(*circuit(params), sep="\n")

Tracebacks

(Array([[0.95, 0.97],
       [0.05, 0.03]], dtype=float64),) # For two equal shot numbers, the probs are stacked
(Array([[0.91],
       [0.09]], dtype=float64), Array([[0.93939394],
       [0.06060606]], dtype=float64)) # For different shot numbers they are not stacked

System information

Name: PennyLane
Version: 0.30.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: /home/david/venvs/dev/lib/python3.10/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, retworkx, scipy, semantic-version, toml
Required-by: PennyLane-Lightning, PennyLane-qiskit

Platform info:           Linux-5.19.0-32-generic-x86_64-with-glibc2.35
Python version:          3.10.6
Numpy version:           1.23.5
Scipy version:           1.9.3
Installed devices:
- default.gaussian (PennyLane-0.30.0.dev0)
- default.mixed (PennyLane-0.30.0.dev0)
- default.qubit (PennyLane-0.30.0.dev0)
- default.qubit.autograd (PennyLane-0.30.0.dev0)
- default.qubit.jax (PennyLane-0.30.0.dev0)
- default.qubit.tf (PennyLane-0.30.0.dev0)
- default.qubit.torch (PennyLane-0.30.0.dev0)
- default.qutrit (PennyLane-0.30.0.dev0)
- null.qubit (PennyLane-0.30.0.dev0)

Existing GitHub issues

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

[josh146]

@dwierichs I re-ran the example and it seems like this bug is no longer present. But I may have missed something -- do you mind double checking, and closing the bug if so?

[dwierichs]

Hi @josh146. Did you run it with "default.qubit" instead of "default.qubit.jax"? Then I also don't see it anymore. With "default.qubit.jax", which is of course deprecated, it is still there.

[josh146]

Okay! let's close it then