[bugfix] Probabilities do not sum to one with Torch

doc/releases/changelog-dev.md

@@ -303,6 +303,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* The probabilities now sum to one using the `torch` interface with `default_dtype` set to `torch.float32`. 
+  [(#5462)](https://github.com/PennyLaneAI/pennylane/pull/5462)
+
 * Tensorflow can now handle devices with float32 results but float64 input parameters.
   [(#5446)](https://github.com/PennyLaneAI/pennylane/pull/5446)
 

pennylane/devices/qubit/sampling.py

@@ -448,10 +448,34 @@ def sample_state(
     with qml.queuing.QueuingManager.stop_recording():
         probs = qml.probs(wires=wires_to_sample).process_state(flat_state, state_wires)
 
+    # when using the torch interface with float32 as default dtype,
+    # probabilities must be renormalized as they may not sum to one
+    # see https://github.com/PennyLaneAI/pennylane/issues/5444
+    norm = qml.math.sum(probs, axis=-1)
+    abs_diff = np.abs(norm - 1.0)
+    cutoff = 1e-07
+
     if is_state_batched:
+
+        normalize_condition = False
+
+        for s in abs_diff:
+            if s != 0:
+                normalize_condition = True
+            if s > cutoff:
+                normalize_condition = False
+                break
+
+        if normalize_condition:
+            probs = probs / norm[:, np.newaxis] if norm.shape else probs / norm
+
         # rng.choice doesn't support broadcasting
         samples = np.stack([rng.choice(basis_states, shots, p=p) for p in probs])
     else:
+
+        if 0 < abs_diff < cutoff:
+            probs /= norm
+
         samples = rng.choice(basis_states, shots, p=probs)
 
     powers_of_two = 1 << np.arange(num_wires, dtype=np.int64)[::-1]

tests/devices/qubit/test_sampling.py

@@ -649,6 +649,63 @@ def test_nan_shadow_expval(self, H, interface, shots):
                 assert qml.math.all(qml.math.isnan(r))
 
 
+two_qubit_state_to_be_normalized = np.array([[0, 1.0000000005j], [-1, 0]]) / np.sqrt(2)
+two_qubit_state_not_normalized = np.array([[0, 1.0000005j], [-1.00000001, 0]]) / np.sqrt(2)
+
+batched_state_to_be_normalized = np.stack(
+    [
+        np.array([[0, 0], [0, 1.000000000009]]),
+        np.array([[1.00000004, 0], [1, 0]]) / np.sqrt(2),
+        np.array([[1, 1], [1, 0.99999995]]) / 2,
+    ]
+)
+batched_state_not_normalized = np.stack(
+    [
+        np.array([[0, 0], [0, 1]]),
+        np.array([[1.0000004, 0], [1, 0]]) / np.sqrt(2),
+        np.array([[1, 1], [1, 0.9999995]]) / 2,
+    ]
+)
+
+
+class TestRenormalization:
+    """Test suite for renormalization functionality."""
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    def test_sample_state_renorm(self, interface):
+        """Test renormalization for a non-batched state."""
+
+        state = qml.math.array(two_qubit_state_to_be_normalized, like=interface)
+        _ = sample_state(state, 10)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    def test_sample_state_renorm_error(self, interface):
+        """Test that renormalization does not occur if the error is too large."""
+
+        state = qml.math.array(two_qubit_state_not_normalized, like=interface)
+        with pytest.raises(ValueError, match="probabilities do not sum to 1"):
+            _ = sample_state(state, 10)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    def test_sample_batched_state_renorm(self, interface):
+        """Test renormalization for a batched state."""
+
+        state = qml.math.array(batched_state_to_be_normalized, like=interface)
+        _ = sample_state(state, 10, is_state_batched=True)
+
+    @pytest.mark.all_interfaces
+    @pytest.mark.parametrize("interface", ["numpy", "jax", "torch", "tensorflow"])
+    def test_sample_batched_state_renorm_error(self, interface):
+        """Test that renormalization does not occur if the error is too large."""
+
+        state = qml.math.array(batched_state_not_normalized, like=interface)
+        with pytest.raises(ValueError, match="probabilities do not sum to 1"):
+            _ = sample_state(state, 10, is_state_batched=True)
+
+
 class TestBroadcasting:
     """Test that measurements work when the state has a batch dim"""