batch and jit support for sample_expectation_ps

Author: refraction-ray

CHANGELOG.md

@@ -14,10 +14,14 @@
 
 - Add `parameter_shift.py` script in examples
 
+- Add jit support and external random management for `tc.quantum.measurement_counts`
+
 ### Changed
 
 - `rxx`, `ryy`, `rzz` gates now has 1/2 factor before theta consitent with `rx`, `ry`, `rz` gates. (breaking change)
 
+- replace `status` arguments in `sample` method as `random_generator` (new convention: status for 0, 1 uniform randomness and random_generator for random key) (breaking change)
+
 ## 0.3.1
 
 ### Added

docs/source/advance.rst

@@ -171,4 +171,6 @@ And a more neat approach to achieve this is as follows:
 
 It is worth noting that since ``Circuit.unitary_kraus`` and ``Circuit.general_kraus`` call ``implicit_rand*`` API, the correct usage of these APIs is the same as above.
 
-One may wonder why random numbers are dealt in such a complicated way, please refer to the `Jax design note <https://github.com/google/jax/blob/main/docs/design_notes/prng.md>`_ for some hints.
+One may wonder why random numbers are dealt in such a complicated way, please refer to the `Jax design note <https://github.com/google/jax/blob/main/docs/design_notes/prng.md>`_ for some hints.
+
+If vmap is also involved apart from jit, I currently find no way to maintain the backend agnosticity as TensorFlow seems to have no support of vmap over random keys (ping me on GitHub if you think you have a way to do this). I strongly recommend the users using Jax backend in the vmap+random setup.

examples/parameter_shift.py

@@ -19,7 +19,7 @@ def f1(param):
             c.cnot(i, i + 1)
         for i in range(n):
             c.rx(i, theta=param[i, j])
-    return c.expectation_ps(y=[n // 2])
+    return K.real(c.expectation_ps(y=[n // 2]))
 
 
 g1f1 = K.jit(K.grad(f1))
@@ -43,21 +43,59 @@ def f2(paramzz, paramx):
             c.rzz(i, i + 1, theta=paramzz[i, j])
         for i in range(n):
             c.rx(i, theta=paramx[i, j])
-    return c.expectation_ps(y=[n // 2])
+    return K.real(c.expectation_ps(y=[n // 2]))
 
 
 g1f2 = K.jit(K.grad(f2, argnums=(0, 1)))
 
-r1, ts, tr = tc.utils.benchmark(
+r12, ts, tr = tc.utils.benchmark(
     g1f2, K.ones([n, m], dtype="float32"), K.ones([n, m], dtype="float32")
 )
 
 g2f2 = K.jit(E.parameter_shift_grad(f2, argnums=(0, 1)))
 
-r2, ts, tr = tc.utils.benchmark(
+r22, ts, tr = tc.utils.benchmark(
     g2f2, K.ones([n, m], dtype="float32"), K.ones([n, m], dtype="float32")
 )
 
-np.testing.assert_allclose(r1[0], r2[0], atol=1e-5)
-np.testing.assert_allclose(r1[1], r2[1], atol=1e-5)
-print("equality test passed!")
+np.testing.assert_allclose(r12[0], r22[0], atol=1e-5)
+np.testing.assert_allclose(r12[1], r22[1], atol=1e-5)
+print("mutilple weight inputs: equality test passed!")
+
+# sampled expectation version
+
+
+def f3(param):
+    c = tc.Circuit(n)
+    for j in range(m):
+        for i in range(n - 1):
+            c.cnot(i, i + 1)
+        for i in range(n):
+            c.rx(i, theta=param[i, j])
+    return K.real(c.sample_expectation_ps(y=[n // 2]))
+
+
+g2f3 = K.jit(E.parameter_shift_grad(f3))
+
+r2, ts, tr = tc.utils.benchmark(g2f3, K.ones([n, m], dtype="float32"))
+
+np.testing.assert_allclose(r1, r2, atol=1e-5)
+print("analytical sampled expectation: equality test passed!")
+
+
+# def f3(param):
+#     c = tc.Circuit(n)
+#     for j in range(m):
+#         for i in range(n - 1):
+#             c.cnot(i, i + 1)
+#         for i in range(n):
+#             c.rx(i, theta=param[i, j])
+#     return K.real(c.sample_expectation_ps(y=[n // 2], shots=81920))
+
+
+# g2f3 = K.jit(E.parameter_shift_grad(f3))
+
+# r2, ts, tr = tc.utils.benchmark(g2f3, K.ones([n, m], dtype="float32"))
+# print(r1 - r2)
+# np.testing.assert_allclose(r1 - r2, np.zeros_like(r1), atol=1e-3)
+# print("finite sampled expectation: equality test passed!")

tensorcircuit/backends/abstract_backend.py

@@ -534,7 +534,7 @@ def argmin(self: Any, a: Tensor, axis: int = 0) -> Tensor:
             "Backend '{}' has not implemented `argmin`.".format(self.name)
         )
 
-    def unique_with_counts(self: Any, a: Tensor) -> Tuple[Tensor, Tensor]:
+    def unique_with_counts(self: Any, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         """
         Find the unique elements and their corresponding counts of the given tensor ``a``.
 

tensorcircuit/backends/jax_backend.py

@@ -367,8 +367,10 @@ def argmax(self, a: Tensor, axis: int = 0) -> Tensor:
     def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
         return jnp.argmin(a, axis=axis)
 
-    def unique_with_counts(self, a: Tensor) -> Tuple[Tensor, Tensor]:
-        return jnp.unique(a, return_counts=True)  # type: ignore
+    def unique_with_counts(
+        self, a: Tensor, *, size: Optional[int] = None, fill_value: Optional[int] = None
+    ) -> Tuple[Tensor, Tensor]:
+        return jnp.unique(a, return_counts=True, size=size, fill_value=fill_value)  # type: ignore
 
     def sigmoid(self, a: Tensor) -> Tensor:
         return libjax.nn.sigmoid(a)

tensorcircuit/backends/numpy_backend.py

@@ -164,7 +164,7 @@ def mean(
     ) -> Tensor:
         return np.mean(a, axis=axis, keepdims=keepdims)
 
-    def unique_with_counts(self, a: Tensor) -> Tuple[Tensor, Tensor]:
+    def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         return np.unique(a, return_counts=True)  # type: ignore
 
     def min(self, a: Tensor, axis: Optional[int] = None) -> Tensor:

tensorcircuit/backends/pytorch_backend.py

@@ -349,7 +349,7 @@ def argmax(self, a: Tensor, axis: int = 0) -> Tensor:
     def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
         return torchlib.argmin(a, dim=axis)
 
-    def unique_with_counts(self, a: Tensor) -> Tuple[Tensor, Tensor]:
+    def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         return torchlib.unique(a, return_counts=True)  # type: ignore
 
     def sigmoid(self, a: Tensor) -> Tensor:

tensorcircuit/backends/tensorflow_backend.py

@@ -355,7 +355,7 @@ def argmax(self, a: Tensor, axis: int = 0) -> Tensor:
     def argmin(self, a: Tensor, axis: int = 0) -> Tensor:
         return tf.math.argmin(a, axis=axis)
 
-    def unique_with_counts(self, a: Tensor) -> Tuple[Tensor, Tensor]:
+    def unique_with_counts(self, a: Tensor, **kws: Any) -> Tuple[Tensor, Tensor]:
         r = tf.unique_with_counts(a)
         order = tf.argsort(r.y)
         return tf.gather(r.y, order), tf.gather(r.count, order)

tensorcircuit/basecircuit.py

@@ -4,7 +4,6 @@
 # pylint: disable=invalid-name
 
 from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
-from copy import deepcopy
 
 import numpy as np
 import graphviz
@@ -543,7 +542,7 @@ def sample(
         self,
         batch: Optional[int] = None,
         allow_state: bool = False,
-        status: Optional[Tensor] = None,
+        random_generator: Optional[Any] = None,
     ) -> Any:
         """
         batched sampling from state or circuit tensor network directly
@@ -553,18 +552,18 @@ def sample(
         :param allow_state: if true, we sample from the final state
             if memory allsows, True is prefered, defaults to False
         :type allow_state: bool, optional
-        :param status: random generator,  defaults to None
-        :type status: Optional[Tensor], optional
+        :param random_generator: random generator,  defaults to None
+        :type random_generator: Optional[Any], optional
         :return: List (if batch) of tuple (binary configuration tensor and correponding probability)
         :rtype: Any
         """
         # allow_state = False is compatibility issue
         if not allow_state:
-            if status is None:
-                status = backend.get_random_state()
+            if random_generator is None:
+                random_generator = backend.get_random_state()
 
             if batch is None:
-                seed = backend.stateful_randu(status, shape=[self._nqubits])
+                seed = backend.stateful_randu(random_generator, shape=[self._nqubits])
                 return self.perfect_sampling(seed)
 
             @backend.jit  # type: ignore
@@ -574,7 +573,7 @@ def perfect_sampling(key: Any) -> Any:
 
             r = []
 
-            subkey = status
+            subkey = random_generator
             for _ in range(batch):
                 key, subkey = backend.random_split(subkey)
                 r.append(perfect_sampling(key))
@@ -589,12 +588,12 @@ def perfect_sampling(key: Any) -> Any:
         if self.is_dm is False:
             p = backend.abs(s) ** 2
         else:
-            p = backend.real(backend.diagonal(s))
-        if status is None:
+            p = backend.abs(backend.diagonal(s))
+        if random_generator is None:
             ch = backend.implicit_randc(a=2**self._nqubits, shape=[nbatch], p=p)
         else:
             ch = backend.stateful_randc(
-                status, a=2**self._nqubits, shape=[nbatch], p=p
+                random_generator, a=2**self._nqubits, shape=[nbatch], p=p
             )
         prob = backend.gather1d(p, ch)
         confg = backend.mod(
@@ -614,6 +613,7 @@ def sample_expectation_ps(
         y: Optional[Sequence[int]] = None,
         z: Optional[Sequence[int]] = None,
         shots: Optional[int] = None,
+        random_generator: Optional[Any] = None,
         **kws: Any,
     ) -> Tensor:
         """
@@ -635,10 +635,15 @@ def sample_expectation_ps(
         :type z: Optional[Sequence[int]], optional
         :param shots: number of measurement shots, defaults to None, indicating analytical result
         :type shots: Optional[int], optional
+        :param random_generator: random_generator, defaults to None
+        :type random_general: Optional[Any]
         :return: [description]
         :rtype: Tensor
         """
-        c = deepcopy(self)
+        if self.is_dm is False:
+            c = type(self)(self._nqubits, mps_inputs=self.quvector())  # type: ignore
+        else:
+            c = type(self)(self._nqubits, mpo_dminputs=self.get_dm_as_quoperator())  # type: ignore
         if x is None:
             x = []
         if y is None:
@@ -653,9 +658,17 @@ def sample_expectation_ps(
         if c.is_dm is False:
             p = backend.abs(s) ** 2
         else:
-            p = backend.real(backend.diagonal(s))
+            p = backend.abs(backend.diagonal(s))
         # readout error can be processed here later
-        mc = measurement_counts(p, counts=shots, sparse=False, is_prob=True)
+        # TODO(@refraction-ray): explicit management on randomness
+        mc = measurement_counts(
+            p,
+            counts=shots,
+            sparse=False,
+            is_prob=True,
+            random_generator=random_generator,
+            jittable=True,
+        )
         x = list(x)
         y = list(y)
         z = list(z)

tensorcircuit/experimental.py

@@ -226,6 +226,7 @@ def parameter_shift_grad(
     :return: the grad function
     :rtype: Callable[..., Tensor]
     """
+    # TODO(@refraction-ray): finite shot sample_expectation_ps not supported well for now
     if jit is True:
         f = backend.jit(f)
 

tensorcircuit/quantum.py

@@ -1878,6 +1878,8 @@ def measurement_counts(
     counts: Optional[int] = 8192,
     sparse: bool = True,
     is_prob: bool = False,
+    random_generator: Optional[Any] = None,
+    jittable: bool = False,
 ) -> Union[Tuple[Tensor, Tensor], Tensor]:
     """
     Simulate the measuring of each qubit of ``p`` in the computational basis,
@@ -1900,6 +1902,11 @@ def measurement_counts(
     :param sparse: Defaults True. The bool indicating whether
         the return form is in the form of two array or one of the same length as the ``state`` (if ``sparse=False``).
     :type sparse: bool
+    :param is_prob: if True, the `state` is directly regarded as a probability list,
+        defaults to be False
+    :type is_prob: bool
+    :param random_generator: random_generator, defaults to None
+    :type random_general: Optional[Any]
     :return: The counts for each bit string measured.
     :rtype: Tuple[]
     """
@@ -1908,21 +1915,30 @@ def measurement_counts(
     else:
         if len(state.shape) == 2:
             state /= backend.trace(state)
-            pi = backend.real(backend.diagonal(state))
+            pi = backend.abs(backend.diagonal(state))
         else:
             state /= backend.norm(state)
             pi = backend.real(backend.conj(state) * state)
         pi = backend.reshape(pi, [-1])
-    d = int(pi.shape[0])
+    d = int(backend.shape_tuple(pi)[0])
+    drange = backend.arange(d)
     # raw counts in terms of integers
     if (counts is None) or counts <= 0:
         if not sparse:
             return pi
         else:
             return counts_d2s(pi)
     else:
-        raw_counts = backend.implicit_randc(d, shape=counts, p=pi)
-        results = backend.unique_with_counts(raw_counts)
+        if random_generator is None:
+            raw_counts = backend.implicit_randc(drange, shape=counts, p=pi)
+        else:
+            raw_counts = backend.stateful_randc(
+                random_generator, a=drange, shape=counts, p=pi
+            )
+        if not jittable:
+            results = backend.unique_with_counts(raw_counts)  # non-jittable
+        else:  # jax specified
+            results = backend.unique_with_counts(raw_counts, size=d, fill_value=-1)
         if sparse:
             return results  # type: ignore
         dense_results = counts_s2d(results, d)

tests/test_circuit.py

@@ -1002,10 +1002,14 @@ def test_batch_sample(backend):
     c.cnot(0, 1)
     print(c.sample())
     print(c.sample(batch=8))
-    print(c.sample(status=tc.backend.get_random_state(42)))
+    print(c.sample(random_generator=tc.backend.get_random_state(42)))
     print(c.sample(allow_state=True))
     print(c.sample(batch=8, allow_state=True))
-    print(c.sample(batch=8, allow_state=True, status=tc.backend.get_random_state(42)))
+    print(
+        c.sample(
+            batch=8, allow_state=True, random_generator=tc.backend.get_random_state(42)
+        )
+    )
 
 
 def test_expectation_y_bug():

tests/test_dmcircuit.py

@@ -468,3 +468,87 @@ def test_prepend_dmcircuit(backend):
         assert n["name"] == n0
     s = c3.wavefunction()
     np.testing.assert_allclose(s[0], s[1], atol=1e-5)
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_dm_sexpps(backend):
+    c = tc.DMCircuit(1, inputs=1 / np.sqrt(2) * np.array([1.0, 1.0j]))
+    y = c.sample_expectation_ps(y=[0])
+    ye = c.expectation_ps(y=[0])
+    np.testing.assert_allclose(y, 1.0, atol=1e-5)
+    np.testing.assert_allclose(ye, 1.0, atol=1e-5)
+
+    c = tc.DMCircuit(4)
+    c.H(0)
+    c.H(1)
+    c.X(2)
+    c.Y(3)
+    c.cnot(0, 1)
+    c.depolarizing(1, px=0.05, py=0.05, pz=0.1)
+    c.rx(1, theta=0.3)
+    c.ccnot(2, 3, 1)
+    c.depolarizing(0, px=0.05, py=0.0, pz=0.1)
+    c.rzz(0, 3, theta=0.5)
+    c.ry(3, theta=2.2)
+    c.amplitudedamping(2, gamma=0.1, p=0.95)
+    c.s(1)
+    c.td(2)
+    c.cswap(3, 0, 1)
+    y = c.sample_expectation_ps(x=[1], y=[0], z=[2, 3])
+    ye = c.expectation_ps(x=[1], y=[0], z=[2, 3])
+    np.testing.assert_allclose(ye, y, atol=1e-5)
+    y2 = c.sample_expectation_ps(x=[1], y=[0], z=[2, 3], shots=81920)
+    assert np.abs(y2 - y) < 0.01
+
+
+def test_dm_sexpps_jittable_vamppable(jaxb):
+    n = 4
+    m = 2
+
+    def f(param, key):
+        c = tc.DMCircuit(n)
+        for j in range(m):
+            for i in range(n - 1):
+                c.cnot(i, i + 1)
+            for i in range(n):
+                c.rx(i, theta=param[i, j])
+        return tc.backend.real(
+            c.sample_expectation_ps(y=[n // 2], shots=8192, random_generator=key)
+        )
+
+    vf = tc.backend.jit(tc.backend.vmap(f, vectorized_argnums=(0, 1)))
+    r = vf(
+        tc.backend.ones([2, n, m], dtype="float32"),
+        tc.backend.stack(
+            [
+                tc.backend.get_random_state(42),
+                tc.backend.get_random_state(43),
+            ]
+        ),
+    )
+    assert np.abs(r[0] - r[1]) > 1e-4
+
+    print(r)
+
+
+def test_dm_sexpps_jittable_vamppable_tf(tfb):
+    # finally giving up backend agnosticity
+    # and not sure the effciency and the safety of vmap random in tf
+    n = 4
+    m = 2
+
+    def f(param):
+        c = tc.DMCircuit(n)
+        for j in range(m):
+            for i in range(n - 1):
+                c.cnot(i, i + 1)
+            for i in range(n):
+                c.rx(i, theta=param[i, j])
+        return tc.backend.real(c.sample_expectation_ps(y=[n // 2], shots=8192))
+
+    vf = tc.backend.jit(tc.backend.vmap(f, vectorized_argnums=0))
+    r = vf(tc.backend.ones([2, n, m]))
+    r1 = vf(tc.backend.ones([2, n, m]))
+    assert np.abs(r[0] - r[1]) > 1e-5
+    assert np.abs(r[0] - r1[0]) > 1e-5
+    print(r, r1)