[BUG] Probabilities do not sum to one with Torch

[trbromley]

Expected behavior

The following code executes regardless of whether the Torch dtype is float64 or float32.

Actual behavior

Error.

Additional information

No response

Source code

import pennylane as qml
import torch

torch.set_default_dtype(torch.float32)

dev = qml.device("default.qubit", shots=1000)

@qml.qnode(dev, interface="torch")
def f(x):
    qml.RX(x, 0)
    return qml.expval(qml.PauliZ(0))

x = torch.tensor([0.4, 0.6])
f(x)

Tracebacks

ValueError: probabilities do not sum to 1

System information

Dev branch of PL and torch v2.2.1.

Existing GitHub issues

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

[albi3ro]

Adding the lines

    norm = qml.math.sum(probs, axis=-1)
    probs = probs/norm

before performing rng.choice seems to solve the problem, but I guess the question is whether we're happy with this solution. Maybe we could add a check to see in the norm is "almost 1" and only re-normalize if it's sufficiently close.