infra for random of backends; randn

Author: refraction-ray

README.md

@@ -98,6 +98,14 @@ pylint tensorcircuit tests
 mypy tensorcircuit
 ```
 
+### Integrated script
+
+For now, we introduce one for all checker for development:
+
+```bash
+./check_all.sh
+```
+
 ### CI
 
 We currently use GitHub Action for test CI, but it has limited quota for free private repo.

check_all.sh

@@ -0,0 +1,13 @@
+#! /bin/sh
+set -e 
+echo "black check"
+black . --check
+echo "mypy check"
+mypy tensorcircuit
+echo "pylint check"
+pylint tensorcircuit tests
+echo "pytest check"
+pytest --cov=tensorcircuit -vv
+echo "sphinx check"
+cd docs && make html
+echo "all checks passed, congratulates!"

mypy.ini

@@ -3,8 +3,20 @@ python_version = 3.6
 ignore_missing_imports = True
 strict = True
 warn_unused_ignores = False
+disallow_untyped_calls = False
+
 
 [mypy-tensorcircuit.applications.van]
 ;;mypy simply cannot ignore files with wildcard patterns...
 ;;only module level * works...
 ignore_errors = True
+
+
+;;[mypy-numpy.*]
+;;ignore_errors = True
+
+;; doesn't work due to https://github.com/python/mypy/issues/10757
+;; mypy + numpy is currently a disaster, never use mypy in your next project
+;; unless you enjoy writting sth worse than C
+;; both the establish status of mypy and support from other packages are just wasting your time
+;; GET AWAY MYPY AND TYPE ANNOTATION !!! WRITTING PYTHON AS IT IS !!!

tensorcircuit/backends.py

@@ -29,7 +29,8 @@
 
 
 Tensor = Any
-
+PRNGKeyArray = Any  # libjax.random.PRNGKeyArray
+RGenerator = Any  # tf.random.Generator
 
 libjax: Any
 jnp: Any
@@ -356,6 +357,77 @@ def tree_map(  # pylint: disable=unused-variable
 
         return r
 
+    def set_random_state(  # pylint: disable=unused-variable
+        self: Any, seed: Optional[int] = None
+    ) -> None:
+        """
+        set random state attached in the backend
+
+        :param seed: int, defaults to None
+        :type seed: Optional[int], optional
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `set_random_state`.".format(self.name)
+        )
+
+    def implicit_randn(  # pylint: disable=unused-variable
+        self: Any,
+        shape: Union[int, Sequence[int]] = 1,
+        mean: float = 0,
+        stddev: float = 1,
+        dtype: str = "32",
+    ) -> Tensor:
+        """
+        call random normal function with the random state management behind the scene
+
+        :param shape: [description], defaults to 1
+        :type shape: Union[int, Sequence[int]], optional
+        :param mean: [description], defaults to 0
+        :type mean: float, optional
+        :param stddev: [description], defaults to 1
+        :type stddev: float, optional
+        :param dtype: [description], defaults to "32"
+        :type dtype: str, optional
+        :return: [description]
+        :rtype: Tensor
+        """
+        g = getattr(self, "g", None)
+        if g is None:
+            self.set_random_state()
+            g = getattr(self, "g", None)
+        r = self.stateful_randn(g, shape, mean, stddev, dtype)
+        return r
+
+    def stateful_randn(  # pylint: disable=unused-variable
+        self: Any,
+        g: Any,
+        shape: Union[int, Sequence[int]] = 1,
+        mean: float = 0,
+        stddev: float = 1,
+        dtype: str = "32",
+    ) -> Tensor:
+        """
+        [summary]
+
+        :param self: [description]
+        :type self: Any
+        :param g: stateful register for each package
+        :type g: Any
+        :param shape: shape of output sampling tensor
+        :type shape: Union[int, Sequence[int]]
+        :param mean: [description], defaults to 0
+        :type mean: float, optional
+        :param stddev: [description], defaults to 1
+        :type stddev: float, optional
+        :param dtype: only real data type is supported, "32" or "64", defaults to "32"
+        :type dtype: str, optional
+        :return: [description]
+        :rtype: Tensor
+        """
+        raise NotImplementedError(
+            "Backend '{}' has not implemented `stateful_randn`.".format(self.name)
+        )
+
     def grad(  # pylint: disable=unused-variable
         self: Any, f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0
     ) -> Callable[..., Any]:
@@ -580,6 +652,33 @@ def real(self, a: Tensor) -> Tensor:
     def cast(self, a: Tensor, dtype: str) -> Tensor:
         return a.astype(getattr(np, dtype))
 
+    def set_random_state(self, seed: Optional[int] = None) -> None:
+        g = np.random.default_rng(seed)  # None auto supported
+        self.g = g
+
+    def stateful_randn(
+        self,
+        g: np.random.Generator,
+        shape: Union[int, Sequence[int]] = 1,
+        mean: float = 0,
+        stddev: float = 1,
+        dtype: str = "32",
+    ) -> Tensor:
+        if isinstance(dtype, str):
+            dtype = dtype[-2:]
+        if isinstance(shape, int):
+            shape = (shape,)
+        r = g.normal(loc=mean, scale=stddev, size=shape)  # type: ignore
+        if dtype == "32":
+            r = r.astype(np.float32)
+        elif dtype == "64":
+            r = r.astype(np.float64)
+        elif not isinstance(dtype, str):
+            r = r.astype(dtype)
+        else:
+            raise ValueError("unspported `dtype` %s" % dtype)
+        return r
+
     def grad(
         self, f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0
     ) -> Callable[..., Any]:
@@ -735,6 +834,49 @@ def is_tensor(self, a: Any) -> bool:
             return True
         return False
 
+    def set_random_state(self, seed: Optional[int] = None) -> None:
+        if seed is None:
+            seed = np.random.randint(42)
+        g = libjax.random.PRNGKey(seed)
+        self.g = g
+
+    def implicit_randn(
+        self,
+        shape: Union[int, Sequence[int]] = 1,
+        mean: float = 0,
+        stddev: float = 1,
+        dtype: str = "32",
+    ) -> Tensor:
+        g = getattr(self, "g", None)
+        if g is None:
+            self.set_random_state()
+            g = getattr(self, "g", None)
+        key, subkey = libjax.random.split(g)
+        r = self.stateful_randn(subkey, shape, mean, stddev, dtype)
+        self.g = key
+        return r
+
+    def stateful_randn(
+        self,
+        g: PRNGKeyArray,
+        shape: Union[int, Sequence[int]] = 1,
+        mean: float = 0,
+        stddev: float = 1,
+        dtype: str = "32",
+    ) -> Tensor:
+        if isinstance(dtype, str):
+            dtype = dtype[-2:]
+        if isinstance(shape, int):
+            shape = (shape,)
+        if dtype == "32":
+            dtyper = jnp.float32
+        elif dtype == "64":
+            dtyper = jnp.float64
+        elif not isinstance(dtype, str):
+            dtyper = dtype
+        r = libjax.random.normal(g, shape=shape, dtype=dtyper) * stddev + mean
+        return r
+
     def grad(
         self, f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0
     ) -> Any:
@@ -824,6 +966,38 @@ def _outer_product_tf(self: Any, tensor1: Tensor, tensor2: Tensor) -> Tensor:
     return tf.tensordot(tensor1, tensor2, 0)
 
 
+def _random_choice_tf(
+    g: RGenerator,
+    a: Union[int, Sequence[int], Tensor],
+    shape: Union[int, Sequence[int]],
+    p: Optional[Union[Sequence[float], Tensor]] = None,
+) -> Tensor:
+    # only replace=True support, replace=False is not implemented
+    # for stateless random module, tf has corresponding categorical function similar to choice
+    # however, such utility is not implemented with ``tf.random.Generator``
+    if isinstance(a, int):
+        assert a > 0
+        a = tf.range(a)
+    assert len(a.shape) == 1
+    if isinstance(shape, int):
+        shape = (shape,)
+    if p is None:
+        dtype = tf.float32
+        p = tf.ones_like(a)
+        p = tf.cast(p, dtype=dtype)
+        p /= tf.reduce_sum(p)
+    else:
+        if not isinstance(p, tf.Tensor):
+            p = tf.constant(p)
+        dtype = p.dtype
+    shape1 = reduce(mul, shape)
+    p_cuml = tf.cumsum(p)
+    r = p_cuml[-1] * (1 - g.uniform([shape1], dtype=dtype))
+    ind = tf.searchsorted(p_cuml, r)
+    res = tf.gather(a, ind)
+    return tf.reshape(res, shape)
+
+
 # temporary hot replace until new version of tensorflow is released,
 # see issue: https://github.com/google/TensorNetwork/issues/940
 # avoid buggy tensordot2 in tensornetwork
@@ -910,6 +1084,35 @@ def real(self, a: Tensor) -> Tensor:
     def cast(self, a: Tensor, dtype: str) -> Tensor:
         return tf.cast(a, dtype=getattr(tf, dtype))
 
+    def set_random_state(self, seed: Optional[Union[int, RGenerator]] = None) -> None:
+        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
+
+    def stateful_randn(
+        self,
+        g: RGenerator,
+        shape: Union[int, Sequence[int]] = 1,
+        mean: float = 0,
+        stddev: float = 1,
+        dtype: str = "32",
+    ) -> Tensor:
+        if isinstance(dtype, str):
+            dtype = dtype[-2:]
+        if isinstance(shape, int):
+            shape = (shape,)
+        if dtype == "32":
+            dtyper = tf.float32
+        elif dtype == "64":
+            dtyper = tf.float64
+        elif not isinstance(dtype, str):
+            dtyper = dtype
+        return g.normal(shape, mean, stddev, dtype=dtyper)
+
     def grad(
         self, f: Callable[..., Any], argnums: Union[int, Sequence[int]] = 0
     ) -> Callable[..., Any]:
@@ -1090,6 +1293,8 @@ def wrapper(
     vvag = vectorized_value_and_grad
 
 
+# TODO(@refraction-ray): lack stateful random methods implementation for now
+# To be added once pytorch backend is ready
 class PyTorchBackend(pytorch_backend.PyTorchBackend):  # type: ignore
     def __init__(self) -> None:
         super(PyTorchBackend, self).__init__()

tensorcircuit/gates.py

@@ -17,6 +17,7 @@
 thismodule = sys.modules[__name__]
 
 Tensor = Any
+Array = Any
 
 # Common single qubit states as np.ndarray objects
 zero_state = np.array([1.0, 0.0], dtype=npdtype)
@@ -107,7 +108,7 @@ class Gate(tn.Node):  # type: ignore
     pass
 
 
-def num_to_tensor(*num: float, dtype: Optional[str] = None) -> Any:
+def num_to_tensor(*num: Union[float, Tensor], dtype: Optional[str] = None) -> Any:
     l = []
     if not dtype:
         dtype = dtypestr
@@ -124,7 +125,7 @@ def num_to_tensor(*num: float, dtype: Optional[str] = None) -> Any:
 array_to_tensor = num_to_tensor
 
 
-def gate_wrapper(m: np.array, n: Optional[str] = None) -> Gate:
+def gate_wrapper(m: Tensor, n: Optional[str] = None) -> Gate:
     if not n:
         n = "unknowngate"
     m = m.astype(npdtype)
@@ -161,7 +162,7 @@ def matrix_for_gate(gate: Gate) -> Tensor:
     return t
 
 
-def bmatrix(a: np.array) -> str:
+def bmatrix(a: Array) -> str:
     """
     Returns a LaTeX bmatrix
 
@@ -261,11 +262,10 @@ def random_single_qubit_gate() -> Gate:
     """
     Returns the random single qubit gate described in https://arxiv.org/abs/2002.07730.
     """
-
     # Get the random parameters
-    theta, alpha, phi = np.random.rand(3) * 2 * np.pi
+    theta, alpha, phi = np.random.rand(3) * 2 * np.pi  # type: ignore
 
-    return rgate(theta, alpha, phi)
+    return rgate(theta, alpha, phi)  # type: ignore
 
 
 rs = random_single_qubit_gate

tensorcircuit/quantum.py

@@ -13,7 +13,6 @@
     Tuple,
     Set,
     List,
-    Type,
 )
 
 import numpy as np
@@ -66,7 +65,7 @@ def quantum_constructor(
 
 def identity(
     space: Sequence[int],
-    dtype: Type[np.number] = np.float64,
+    dtype: Any = np.float64,
 ) -> "QuOperator":
     """Construct a `QuOperator` representing the identity on a given space.
     Internally, this is done by constructing `CopyNode`s for each edge, with