add probability_sample method for backend

Author: refraction-ray

CHANGELOG.md

@@ -20,6 +20,8 @@
 
 - Add `searchsorted` method for backend
 
+- Add `probability_sample` method for backend as an alternative for `random_choice` since it supports `status` as external randomness format
+
 ### Changed
 
 - The inner mechanism for `sample_expectation_ps` is changed to sample representation from count representation for a fast speed

tensorcircuit/backends/abstract_backend.py

@@ -1067,6 +1067,41 @@ def stateful_randc(
             "Backend '{}' has not implemented `stateful_randc`.".format(self.name)
         )
 
+    def probability_sample(
+        self: Any, shots: int, p: Tensor, status: Optional[Tensor] = None, g: Any = None
+    ) -> Tensor:
+        """
+        Drawn ``shots`` samples from probability distribution p, given the external randomness
+        determined by uniform distributed ``status`` tensor or backend random generator ``g``.
+        This method is similar with ``stateful_randc``, but it supports ``status`` beyond ``g``,
+        which is convenient when jit or vmap
+
+        :param shots: Number of samples to draw with replacement
+        :type shots: int
+        :param p: prbability vector
+        :type p: Tensor
+        :param status: external randomness as a tensor with each element drawn uniformly from [0, 1],
+            defaults to None
+        :type status: Optional[Tensor], optional
+        :param g: backend random genrator, defaults to None
+        :type g: Any, optional
+        :return: The drawn sample as an int tensor
+        :rtype: Tensor
+        """
+        if status is not None:
+            status = self.convert_to_tensor(status)
+        elif g is not None:
+            status = self.stateful_randu(g, shape=[shots])
+        else:
+            status = self.implicit_randu(shape=[shots])
+        p = p / self.sum(p)
+        p_cuml = self.cumsum(p)
+        r = p_cuml[-1] * (1 - self.cast(status, p.dtype))
+        ind = self.searchsorted(p_cuml, r)
+        a = self.arange(shots)
+        res = self.gather1d(a, ind)
+        return res
+
     def gather1d(self: Any, operand: Tensor, indices: Tensor) -> Tensor:
         """
         Return ``operand[indices]``, both ``operand`` and ``indices`` are rank-1 tensor.

tests/test_backends.py

@@ -283,6 +283,16 @@ def test_backend_methods_2(backend):
     values = tc.backend.convert_to_tensor(np.array([0.0, 4.1, 12.0], dtype=np.float32))
     r = tc.backend.numpy(tc.backend.searchsorted(edges, values))
     np.testing.assert_allclose(r, np.array([1, 2, 4]))
+    p = tc.backend.convert_to_tensor(
+        np.array(
+            [0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.05, 0.2, 0.4], dtype=np.float32
+        )
+    )
+    r = tc.backend.probability_sample(10000, p, status=np.random.uniform(size=[10000]))
+    _, r = np.unique(r, return_counts=True)
+    np.testing.assert_allclose(
+        r - tc.backend.numpy(p) * 10000.0, np.zeros([10]), atol=100, rtol=1
+    )
 
 
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])