add jittable option in sample

Author: refraction-ray

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Added
+
+- Add `jittable` option in `c.sample()` method, friendly to switch off for large scale sample
+
 ### Changed
 
 - Downgrading some logger warning to info

tensorcircuit/basecircuit.py

@@ -528,6 +528,7 @@ def sample(
         format: Optional[str] = None,
         random_generator: Optional[Any] = None,
         status: Optional[Tensor] = None,
+        jittable: bool = True,
     ) -> Any:
         """
         batched sampling from state or circuit tensor network directly
@@ -547,6 +548,9 @@ def sample(
         :param status: external randomness given by tensor uniformly from [0, 1],
             if set, can overwrite random_generator
         :type status: Optional[Tensor]
+        :param jittable: when converting to count, whether keep the full size. if false, may be conflict
+            external jit, if true, may fail for large scale system with actual limited count results
+        :type jittable: bool, defaults true
         :return: List (if batch) of tuple (binary configuration tensor and corresponding probability)
             if the format is None, and consistent with format when given
         :rtype: Any
@@ -612,7 +616,7 @@ def perfect_sampling(key: Any) -> Any:
                 if batch is None:
                     r = r[0]  # type: ignore
                 return r
-        return sample2all(sample=ch, n=self._nqubits, format=format, jittable=True)
+        return sample2all(sample=ch, n=self._nqubits, format=format, jittable=jittable)
 
     def sample_expectation_ps(
         self,

tests/test_circuit.py

@@ -236,6 +236,21 @@ def f5(key):
             )
 
 
+@pytest.mark.parametrize("backend", [lf("jaxb")])  # too slow for np
+def test_large_scale_sample(backend):
+    L = 30
+    c = tc.Circuit(L)
+    c.h(0)
+    c.cnot([i for i in range(L - 1)], [i + 1 for i in range(L - 1)])
+    results = c.sample(
+        allow_state=False, batch=1024, format="count_dict_bin", jittable=False
+    )
+    assert (
+        results["0" * L] / results["1" * L] < 1.2
+        and results["0" * L] / results["1" * L] > 0.8
+    )
+
+
 @pytest.mark.parametrize("backend", [lf("npb"), lf("cpb")])
 def test_expectation(backend):
     c = tc.Circuit(2)