more random methods in backends;more stable sort of nodes

Author: refraction-ray

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Added
+
+- add `get_random_state` and `random_split` methods to backends
+
 ### Fixed
 
 - avoid error on watch non `tf.Tensor` in tensorflow backend grad method

examples/mcnoise_boost.py

@@ -0,0 +1,104 @@
+"""
+boosting the monte carlo noise simulation on general error with circuit layerwise slicing
+"""
+
+from functools import partial
+import time
+import jax
+
+
+import tensorcircuit as tc
+
+tc.set_backend("jax")
+
+n = 10
+nlayer = 4
+
+
+@partial(tc.backend.jit, static_argnums=(2, 3))
+def f1(key, param, n, nlayer):
+    if key is not None:
+        tc.backend.set_random_state(key)
+    c = tc.Circuit(n)
+    for i in range(n):
+        c.H(i)
+    for j in range(nlayer):
+        for i in range(n - 1):
+            c.cnot(i, i + 1)
+            c.apply_general_kraus(tc.channels.phasedampingchannel(0.15), i)
+            c.apply_general_kraus(tc.channels.phasedampingchannel(0.15), i + 1)
+        for i in range(n):
+            c.rx(i, theta=param[j, i])
+    return tc.backend.real(c.expectation((tc.gates.z(), [int(n / 2)])))
+
+
+@partial(tc.backend.jit, static_argnums=(2))
+def templatecnot(s, param, i):
+    c = tc.Circuit(n, inputs=s)
+    c.cnot(i, i + 1)
+    return c.state()
+
+
+@partial(tc.backend.jit, static_argnums=(3))
+def templatenoise(key, s, param, i):
+    c = tc.Circuit(n, inputs=s)
+    status = tc.backend.stateful_randu(key)[0]
+    c.apply_general_kraus(tc.channels.phasedampingchannel(0.15), i, status=status)
+    return c.state()
+
+
+@partial(tc.backend.jit, static_argnums=(2))
+def templaterz(s, param, j):
+    c = tc.Circuit(n, inputs=s)
+    for i in range(n):
+        c.rx(i, theta=param[j, i])
+    return c.state()
+
+
+@partial(tc.backend.jit, static_argnums=(2, 3))
+def f2(key, param, n, nlayer):
+    c = tc.Circuit(n)
+    for i in range(n):
+        c.H(i)
+    s = c.state()
+    for j in range(nlayer):
+        for i in range(n - 1):
+            s = templatecnot(s, param, i)
+            key, subkey = tc.backend.random_split(key)
+            s = templatenoise(subkey, s, param, i)
+            key, subkey = tc.backend.random_split(key)
+            s = templatenoise(subkey, s, param, i + 1)
+        s = templaterz(s, param, j)
+    return tc.backend.real(tc.expectation((tc.gates.z(), [int(n / 2)]), ket=s))
+
+
+vagf1 = tc.backend.jit(tc.backend.value_and_grad(f1, argnums=1), static_argnums=(2, 3))
+
+vagf2 = tc.backend.jit(tc.backend.value_and_grad(f2, argnums=1), static_argnums=(2, 3))
+
+param = tc.backend.ones([nlayer, n])
+
+
+def benchmark(f, tries=3):
+    time0 = time.time()
+    key = tc.backend.get_random_state(42)
+    print(f(key, param, n, nlayer)[0])
+    time1 = time.time()
+    for _ in range(tries):
+        print(f(key, param, n, nlayer)[0])
+    time2 = time.time()
+    print(
+        "staging time: ",
+        time1 - time0,
+        "running time: ",
+        (time2 - time1) / tries,
+    )
+
+
+print("without layerwise slicing jit")
+benchmark(vagf1)
+print("=============================")
+print("with layerwise slicing jit")
+benchmark(vagf2)
+
+# 235/0.36 vs. 26/0.04

tensorcircuit/backends/abstract_backend.py

@@ -436,18 +436,50 @@ def tree_map(  # pylint: disable=unused-variable
         return r
 
     def set_random_state(  # pylint: disable=unused-variable
-        self: Any, seed: Optional[int] = None
-    ) -> None:
+        self: Any, seed: Optional[int] = None, get_only: bool = False
+    ) -> Any:
         """
         set random state attached in the backend
 
         :param seed: int, defaults to None
         :type seed: Optional[int], optional
+        :param get_only:
+        :type get_only: bool
         """
         raise NotImplementedError(
             "Backend '{}' has not implemented `set_random_state`.".format(self.name)
         )
 
+    def get_random_state(  # pylint: disable=unused-variable
+        self: Any, seed: Optional[int] = None
+    ) -> Any:
+        """
+        get backend specific random state object
+
+        :param seed: [description], defaults to None
+        :type seed: Optional[int], optional
+        :return: [description]
+        :rtype: Any
+        """
+        return self.set_random_state(seed, True)
+
+    def random_split(  # pylint: disable=unused-variable
+        self: Any, key: Any
+    ) -> Tuple[Any, Any]:
+        """
+        a jax like split API, but does't split the key generator for other backends.
+        just for a consistent interface of random code, be careful that you know what the function actually does.
+
+        :param key: [description]
+        :type key: Any
+        :return: [description]
+        :rtype: Tuple[Any, Any]
+        """
+        return key, key
+
+    # Though try hard, the current random API abstraction may not be perfect when with nested jit or vmap
+    # so keep every random function a direct status parameters in case.
+
     def implicit_randn(  # pylint: disable=unused-variable
         self: Any,
         shape: Union[int, Sequence[int]] = 1,

tensorcircuit/backends/jax_backend.py

@@ -223,14 +223,21 @@ def is_tensor(self, a: Any) -> bool:
             return True
         return False
 
-    def set_random_state(self, seed: Optional[Union[int, PRNGKeyArray]] = None) -> None:
+    def set_random_state(
+        self, seed: Optional[Union[int, PRNGKeyArray]] = None, get_only: bool = False
+    ) -> Any:
         if seed is None:
             seed = np.random.randint(42)
         if isinstance(seed, int):
             g = libjax.random.PRNGKey(seed)
         else:
             g = seed
-        self.g = g
+        if get_only is False:
+            self.g = g
+        return g
+
+    def random_split(self, key: Any) -> Tuple[Any, Any]:
+        return libjax.random.split(key)  # type: ignore
 
     def implicit_randn(
         self,

tensorcircuit/backends/numpy_backend.py

@@ -136,9 +136,13 @@ def cast(self, a: Tensor, dtype: str) -> Tensor:
             return a.astype(getattr(np, dtype))
         return a.astype(dtype)
 
-    def set_random_state(self, seed: Optional[int] = None) -> None:
+    def set_random_state(
+        self, seed: Optional[int] = None, get_only: bool = False
+    ) -> Any:
         g = np.random.default_rng(seed)  # None auto supported
-        self.g = g
+        if get_only is False:
+            self.g = g
+        return g
 
     def stateful_randn(
         self,

tensorcircuit/backends/tensorflow_backend.py

@@ -204,14 +204,18 @@ def cast(self, a: Tensor, dtype: str) -> Tensor:
             return tf.cast(a, dtype=getattr(tf, dtype))
         return tf.cast(a, dtype=dtype)
 
-    def set_random_state(self, seed: Optional[Union[int, RGenerator]] = None) -> None:
+    def set_random_state(
+        self, seed: Optional[Union[int, RGenerator]] = None, get_only: bool = False
+    ) -> Any:
         if seed is None:
             g = tf.random.Generator.from_non_deterministic_state()
         elif isinstance(seed, int):
             g = tf.random.Generator.from_seed(seed)
         else:
             g = seed
-        self.g = g
+        if get_only is False:
+            self.g = g
+        return g
 
     def stateful_randn(
         self,

tensorcircuit/circuit.py

@@ -673,7 +673,7 @@ def _general_kraus_2(
         # building for jax+GPU ~100s 12 qubit * 5 layers
         # 370s 14 qubit * 7 layers, 0.35s running on vT4
         # vmap, grad, vvag are all fine for this function
-        # layerwise jit technique can greatly boost the staging time
+        # layerwise jit technique can greatly boost the staging time, see in /examples
         sites = len(index)
         kraus_tensor = [k.tensor for k in kraus]
 

tensorcircuit/cons.py

@@ -336,14 +336,8 @@ def _get_path_cache_friendly(
                 i += 1
     # TODO(@refraction-ray): may be not that cache friendly, since the edge id correspondence is not that fixed?
     input_sets = [set([mapping_dict[id(e)] for e in node.edges]) for node in nodes]
-    placeholder = [1e10]
-    for s in input_sets:
-        if len(s) > 1:
-            break
-    else:
-        placeholder = [1e10, 1e10]
-    order = np.argsort(np.array(list(map(sorted, input_sets)) + [placeholder], dtype=object))[:-1]  # type: ignore
-    # TODO(@refraction-ray): more stable and unwarning arg sorting here
+    placeholder = [[1e20 for _ in range(100)]]
+    order = np.argsort(np.array(list(map(sorted, input_sets)) + placeholder, dtype=object))[:-1]  # type: ignore
     nodes_new = [nodes[i] for i in order]
     input_sets = [set([mapping_dict[id(e)] for e in node.edges]) for node in nodes_new]
     output_set = set([mapping_dict[id(e)] for e in tn.get_subgraph_dangling(nodes_new)])

tests/test_backends.py

@@ -348,3 +348,17 @@ def test_sparse_methods(backend):
         np.array([[1], [2], [0], [0]], dtype=np.complex64),
         atol=1e-5,
     )
+
+
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb")])
+def test_backend_randoms_v2(backend):
+    g = tc.backend.get_random_state(42)
+    for t in tc.backend.stateful_randc(g, 3, [3]):
+        assert t >= 0
+        assert t < 3
+    key = tc.backend.get_random_state(42)
+    r = []
+    for _ in range(2):
+        key, subkey = tc.backend.random_split(key)
+        r.append(tc.backend.stateful_randc(subkey, 3, [5]))
+    assert tuple(r[0]) != tuple(r[1])