fully use status in circuit.sample

Author: refraction-ray

CHANGELOG.md

@@ -34,6 +34,8 @@
 
 - `MPSCircuit` now will first try to transform `QuVector` input to tensors directly instead of evaluating it to dense wavefunction first.
 
+- Fix to use `status` for `circuit.sample` when `allow_state=True`.
+
 ### Changed
 
 - The order of arguments of `tc.timeevol.ed_evol` are changed for consistent interface with other evolution methods.

tensorcircuit/backends/tensorflow_backend.py

@@ -533,6 +533,9 @@ def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
     def argsort(self, a: Tensor, axis: int = -1) -> Tensor:
         return tf.argsort(a, axis=axis)
 
+    def shape_tuple(self, a: Tensor) -> Tuple[int, ...]:
+        return tuple(a.shape)
+
     def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         r = tf.unique_with_counts(a)
         order = tf.argsort(r.y)

tensorcircuit/basecircuit.py

@@ -576,7 +576,8 @@ def sample(
         :param random_generator: random generator,  defaults to None
         :type random_generator: Optional[Any], optional
         :param status: external randomness given by tensor uniformly from [0, 1],
-            if set, can overwrite random_generator
+            if set, can overwrite random_generator, shape [batch] for `allow_state=True`
+            and shape [batch, nqubits] for `allow_state=False` using perfect sampling implementation
         :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
@@ -597,20 +598,24 @@ def sample(
                     return r
                 r = [r]  # type: ignore
             else:
+                r = []  # type: ignore
+                if status is not None:
+                    assert backend.shape_tuple(status)[0] == batch
+                    for seed in status:
+                        r.append(self.perfect_sampling(seed))  # type: ignore
 
-                @backend.jit
-                def perfect_sampling(key: Any) -> Any:
-                    backend.set_random_state(key)
-                    return self.perfect_sampling()
+                else:
 
-                # TODO(@refraction-ray): status is not used here
+                    @backend.jit
+                    def perfect_sampling(key: Any) -> Any:
+                        backend.set_random_state(key)
+                        return self.perfect_sampling()
 
-                r = []  # type: ignore
+                    subkey = random_generator
+                    for _ in range(batch):
+                        key, subkey = backend.random_split(subkey)
+                        r.append(perfect_sampling(key))  # type: ignore
 
-                subkey = random_generator
-                for _ in range(batch):
-                    key, subkey = backend.random_split(subkey)
-                    r.append(perfect_sampling(key))  # type: ignore
             if format is None:
                 return r
             r = backend.stack([ri[0] for ri in r])  # type: ignore

tests/test_circuit.py

@@ -1285,6 +1285,7 @@ def test_batch_sample(backend):
     c.H(0)
     c.cnot(0, 1)
     print(c.sample())
+    print(c.sample(batch=8, status=np.random.uniform(size=[8, 3])))
     print(c.sample(batch=8))
     print(c.sample(random_generator=tc.backend.get_random_state(42)))
     print(c.sample(allow_state=True))
@@ -1302,6 +1303,14 @@ def test_batch_sample(backend):
             format="sample_bin",
         )
     )
+    print(
+        c.sample(
+            batch=8,
+            allow_state=False,
+            status=np.random.uniform(size=[8, 3]),
+            format="sample_bin",
+        )
+    )
 
 
 def test_expectation_y_bug():