new sample API and new backend methods

Author: refraction-ray

CHANGELOG.md

@@ -6,6 +6,18 @@
 
 - Add PyTorch nn Module wrapper in `torchnn`
 
+- Add `reverse`, `mod`, `left_shift`, `right_shift`, `arange` methods on backend
+
+- Brand new `sample` API with batch support and sampling from state support
+
+### Fixed
+
+- Fixed bug in merge single gates when all gates are single-qubit ones
+
+### Changed
+
+- The default contractor enable preprocessing feature where single-qubit gates are merged firstly
+
 ## 0.1.3
 
 ### Added

docs/source/quickstart.rst

@@ -239,6 +239,7 @@ and the other part is implemented in `TensorCircuit package <modules.html#module
     'acosh',
     'addition',
     'adjoint',
+    'arange',
     'argmax',
     'argmin',
     'asin',
@@ -255,6 +256,7 @@ and the other part is implemented in `TensorCircuit package <modules.html#module
     'conj',
     'convert_to_tensor',
     'coo_sparse_matrix',
+    'coo_sparse_matrix_from_numpy',
     'copy',
     'cos',
     'cosh',
@@ -267,6 +269,7 @@ and the other part is implemented in `TensorCircuit package <modules.html#module
     'eigs',
     'eigsh',
     'eigsh_lanczos',
+    'eigvalsh',
     'einsum',
     'eps',
     'exp',
@@ -293,12 +296,14 @@ and the other part is implemented in `TensorCircuit package <modules.html#module
     'jit',
     'jvp',
     'kron',
+    'left_shift',
     'log',
     'matmul',
     'max',
     'mean',
     'min',
     'minor',
+    'mod',
     'multiply',
     'name',
     'norm',
@@ -319,6 +324,8 @@ and the other part is implemented in `TensorCircuit package <modules.html#module
     'reshape',
     'reshape2',
     'reshapem',
+    'reverse',
+    'right_shift',
     'rq',
     'scatter',
     'serialize_tensor',
@@ -357,7 +364,9 @@ and the other part is implemented in `TensorCircuit package <modules.html#module
     'to_dense',
     'trace',
     'transpose',
+    'tree_flatten',
     'tree_map',
+    'tree_unflatten',
     'unique_with_counts',
     'value_and_grad',
     'vectorized_value_and_grad',

docs/source/tutorial.rst

@@ -5,6 +5,7 @@ Jupyter Tutorials
 .. toctree::
 
     tutorials/circuit_basics.ipynb
+    tutorials/qaoa.ipynb
     tutorials/tfim_vqe.ipynb
     tutorials/mnist_qml.ipynb
     tutorials/torch_qml.ipynb

examples/sample_benchmark.py

@@ -27,16 +27,19 @@ def construct_circuit(n, nlayers):
         print("n: ", n, " nlayers: ", nlayers)
         c = construct_circuit(n, nlayers)
         time0 = time.time()
-        s = c.state()
+        s = c.sample(allow_state=True)
         time1 = time.time()
-        smp = bin(np.random.choice(range(2**n), p=np.abs(K.numpy(s)) ** 2))
         # print(smp)
         print("state sampling time: ", time1 - time0)
         time0 = time.time()
         smp = c.sample()
         # print(smp)
         time1 = time.time()
         print("nonjit tensor sampling time: ", time1 - time0)
+        time0 = time.time()
+        s = c.sample(allow_state=True, batch=10)
+        time1 = time.time()
+        print("batch state sampling time: ", (time1 - time0) / 10)
 
         @K.jit
         def f(key):
@@ -48,11 +51,10 @@ def f(key):
         time0 = time.time()
         smp = f(key1)
         time1 = time.time()
-        for _ in range(5):
+        for _ in range(10):
             key1, key2 = K.random_split(key2)
             smp = f(key1)
             # print(smp)
         time2 = time.time()
-
-        print("jittable tensor sampling staging time: ", time1 - time0)
-        print("jittable tensor sampling running time: ", (time2 - time1) / 5)
+        print("jittable tensor sampling staginging time: ", time1 - time0)
+        print("jittable tensor sampling running time: ", (time2 - time1) / 10)

tensorcircuit/backends/abstract_backend.py

@@ -644,6 +644,81 @@ def cast(self: Any, a: Tensor, dtype: str) -> Tensor:
             "Backend '{}' has not implemented `cast`.".format(self.name)
         )
 
+    def mod(self: Any, x: Tensor, y: Tensor) -> Tensor:
+        """
+        Compute y-mod of x (negative number behavior is not guaranteed to be consistent)
+
+        :param x: input values
+        :type x: Tensor
+        :param y: mod ``y``
+        :type y: Tensor
+        :return: results
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `mod`.".format(self.name)
+        )
+
+    def reverse(self: Any, a: Tensor) -> Tensor:
+        """
+        return ``a[::-1]``, only 1D tensor is guaranteed for consistent behavior
+
+        :param a: 1D tensor
+        :type a: Tensor
+        :return: 1D tensor in reverse order
+        :rtype: Tensor
+        """
+        return a[::-1]
+
+    def right_shift(self: Any, x: Tensor, y: Tensor) -> Tensor:
+        """
+        Shift the bits of an integer x to the right y bits.
+
+        :param x: input values
+        :type x: Tensor
+        :param y: Number of bits shift to ``x``
+        :type y: Tensor
+        :return: result with the same shape as ``x``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `right_shift`.".format(self.name)
+        )
+
+    def left_shift(self: Any, x: Tensor, y: Tensor) -> Tensor:
+        """
+        Shift the bits of an integer x to the left y bits.
+
+        :param x: input values
+        :type x: Tensor
+        :param y: Number of bits shift to ``x``
+        :type y: Tensor
+        :return: result with the same shape as ``x``
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `left_shift`.".format(self.name)
+        )
+
+    def arange(
+        self: Any, start: int, stop: Optional[int] = None, step: int = 1
+    ) -> Tensor:
+        """
+        Values are generated within the half-open interval [start, stop)
+
+        :param start: start index
+        :type start: int
+        :param stop: end index, defaults to None
+        :type stop: Optional[int], optional
+        :param step: steps, defaults to 1
+        :type step: Optional[int], optional
+        :return: _description_
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `arange`.".format(self.name)
+        )
+
     def solve(self: Any, A: Tensor, b: Tensor, **kws: Any) -> Tensor:
         """
         Solve the linear system Ax=b and return the solution x.

tensorcircuit/backends/jax_backend.py

@@ -315,6 +315,20 @@ def cast(self, a: Tensor, dtype: str) -> Tensor:
             return a.astype(getattr(jnp, dtype))
         return a.astype(dtype)
 
+    def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tensor:
+        if stop is None:
+            return jnp.arange(start=0, stop=start, step=step)
+        return jnp.arange(start=start, stop=stop, step=step)
+
+    def mod(self, x: Tensor, y: Tensor) -> Tensor:
+        return jnp.mod(x, y)
+
+    def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
+        return jnp.right_shift(x, y)
+
+    def left_shift(self, x: Tensor, y: Tensor) -> Tensor:
+        return jnp.left_shift(x, y)
+
     def expm(self, a: Tensor) -> Tensor:
         return jsp.linalg.expm(a)
         # currently expm in jax doesn't support AD, it will raise an AssertError,

tensorcircuit/backends/numpy_backend.py

@@ -211,6 +211,20 @@ def cast(self, a: Tensor, dtype: str) -> Tensor:
             return a.astype(getattr(np, dtype))
         return a.astype(dtype)
 
+    def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tensor:
+        if stop is None:
+            return np.arange(start=0, stop=start, step=step)
+        return np.arange(start=start, stop=stop, step=step)
+
+    def mod(self, x: Tensor, y: Tensor) -> Tensor:
+        return np.mod(x, y)
+
+    def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
+        return np.right_shift(x, y)
+
+    def left_shift(self, x: Tensor, y: Tensor) -> Tensor:
+        return np.left_shift(x, y)
+
     def solve(self, A: Tensor, b: Tensor, assume_a: str = "gen") -> Tensor:
         # gen, sym, her, pos
         # https://stackoverflow.com/questions/44672029/difference-between-numpy-linalg-solve-and-numpy-linalg-lu-solve/44710451

tensorcircuit/backends/pytorch_backend.py

@@ -375,9 +375,26 @@ def cast(self, a: Tensor, dtype: str) -> Tensor:
             return a.type(getattr(torchlib, dtype))
         return a.type(dtype)
 
+    def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tensor:
+        if stop is None:
+            return torchlib.arange(start=0, end=start, step=step)
+        return torchlib.arange(start=start, end=stop, step=step)
+
+    def mod(self, x: Tensor, y: Tensor) -> Tensor:
+        return torchlib.fmod(x, y)
+
+    def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
+        return torchlib.bitwise_right_shift(x, y)
+
+    def left_shift(self, x: Tensor, y: Tensor) -> Tensor:
+        return torchlib.bitwise_left_shift(x, y)
+
     def solve(self, A: Tensor, b: Tensor, **kws: Any) -> Tensor:
         return torchlib.linalg.solve(A, b)
 
+    def reverse(self, a: Tensor) -> Tensor:
+        return torchlib.flip(a, dims=(-1,))
+
     def tree_map(self, f: Callable[..., Any], *pytrees: Any) -> Any:
         # TODO(@refraction-ray): torch not support multiple pytree args
         return torchlib.utils._pytree.tree_map(f, *pytrees)

tensorcircuit/backends/tensorflow_backend.py

@@ -414,6 +414,20 @@ def cast(self, a: Tensor, dtype: str) -> Tensor:
             return tf.cast(a, dtype=getattr(tf, dtype))
         return tf.cast(a, dtype=dtype)
 
+    def arange(self, start: int, stop: Optional[int] = None, step: int = 1) -> Tensor:
+        if stop is None:
+            return tf.range(start=0, limit=start, delta=step)
+        return tf.range(start=start, limit=stop, delta=step)
+
+    def mod(self, x: Tensor, y: Tensor) -> Tensor:
+        return tf.math.mod(x, y)
+
+    def right_shift(self, x: Tensor, y: Tensor) -> Tensor:
+        return tf.bitwise.right_shift(x, y)
+
+    def left_shift(self, x: Tensor, y: Tensor) -> Tensor:
+        return tf.bitwise.left_shift(x, y)
+
     def solve(self, A: Tensor, b: Tensor, **kws: Any) -> Tensor:
         if b.shape[-1] == A.shape[-1]:
             b = b[..., tf.newaxis]

tensorcircuit/circuit.py

@@ -1419,7 +1419,70 @@ def perfect_sampling(self) -> Tuple[str, float]:
         """
         return self.measure_jit(*[i for i in range(self._nqubits)], with_prob=True)
 
-    sample = perfect_sampling
+    # sample = perfect_sampling
+
+    def sample(
+        self,
+        batch: Optional[int] = None,
+        allow_state: bool = False,
+        status: Optional[Tensor] = None,
+    ) -> Any:
+        """
+        batched sampling from state or circuit tensor network directly
+
+        :param batch: number of samples, defaults to None
+        :type batch: Optional[int], optional
+        :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
+        :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 batch is None:
+                return self.perfect_sampling()
+
+            @backend.jit  # type: ignore
+            def perfect_sampling(key: Any) -> Any:
+                backend.set_random_state(key)
+                return self.perfect_sampling()
+
+            r = []
+            if status is None:
+                status = backend.get_random_state()
+            subkey = status
+            for _ in range(batch):
+                key, subkey = backend.random_split(subkey)
+                r.append(perfect_sampling(key))
+
+            return r
+
+        if batch is None:
+            nbatch = 1
+        else:
+            nbatch = batch
+        s = self.state()
+        p = backend.abs(s) ** 2
+        if status 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
+            )
+        prob = backend.gather1d(p, ch)
+        confg = backend.mod(
+            backend.right_shift(
+                ch[..., None], backend.reverse(backend.arange(self._nqubits))
+            ),
+            2,
+        )
+        r = list(zip(confg, prob))
+        if batch is None:
+            r = r[0]
+        return r
 
     # TODO(@refraction-ray): more _before function like state_before? and better API?