Merge branch 'beta' into cloud

Author: refraction-ray

check_all.sh

@@ -8,6 +8,8 @@ echo "pylint check"
 pylint tensorcircuit tests examples/*.py
 echo "pytest check"
 pytest -n auto --cov=tensorcircuit -vv -W ignore::DeprecationWarning
+# for test on gpu machine, please set `export TF_FORCE_GPU_ALLOW_GROWTH=true` for tf
+# and `export XLA_PYTHON_CLIENT_PREALLOCATE=false` for jax to avoid OOM in testing
 echo "sphinx check"
 cd docs && sphinx-build source build/html && sphinx-build source -D language="zh" build/html_cn
 echo "all checks passed, congratulation! 💐"

examples/variational_dynamics_circuit.py

@@ -184,10 +184,10 @@ def numdiff(i):
             f.append(-0.5j)
             door.append([0, i, numdiff(i)])
         for i in range(N - 1):
-            f.append(-1j)
+            f.append(-0.5j)
             door.append([5, i, i + 1])
         for i in range(N - 1):
-            f.append(-1j)
+            f.append(-0.5j)
             door.append([3, i, i + 1])
     for i in range(N):
         h.append(1)
@@ -235,54 +235,60 @@ def numdiff(i):
 
     # variation realize
     ODE_theta = tf.zeros(len(door), dtype="float64")
+
+    a_batch_theta = []
+    a_batch_is_k = []
+    a_batch_is_q = []
+    for k in range(len(door)):
+        for q in range(len(door)):
+            is_k = [0 for _ in range(len(door))]
+            is_k[k] = 1
+            is_q = [0 for _ in range(len(door))]
+            is_q[q] = 1
+            if how_variation == 0:
+                a_batch_theta.append(np.angle(f[q]) - np.angle(f[k]))
+            else:
+                a_batch_theta.append(np.angle(f[q]) - np.angle(f[k]) - math.pi / 2)
+            a_batch_is_k.append(is_k)
+            a_batch_is_q.append(is_q)
+    a_batch_theta = tc.array_to_tensor(a_batch_theta)
+    a_batch_is_k = tf.constant(a_batch_is_k)
+    a_batch_is_q = tf.constant(a_batch_is_q)
+
+    c_batch_theta = []
+    c_batch_is_k = []
+    c_batch_is_q = []
+    for k in range(len(door)):
+        for q in range(len(h_door)):
+            is_k = [0 for _ in range(len(door))]
+            is_k[k] = 1
+            is_q = [0 for _ in range(len(door))]
+            is_q[q] = 1
+            c_batch_is_k.append(is_k)
+            c_batch_is_q.append(is_q)
+            if how_variation == 0:
+                c_batch_theta.append(np.angle(h[q]) - np.angle(f[k]) - math.pi / 2)
+            else:
+                c_batch_theta.append(np.angle(h[q]) - np.angle(f[k]) + math.pi)
+    c_batch_theta = tc.array_to_tensor(c_batch_theta)
+    c_batch_is_k = tf.constant(c_batch_is_k)
+    c_batch_is_q = tf.constant(c_batch_is_q)
+
     for T in range(int(t / dt)):
         # calculate coefficient in paper
-        A = np.zeros((len(door), len(door)))
-        C = np.zeros(len(door))
-        batch_theta = []
-        batch_is_k = []
-        batch_is_q = []
-        for k in range(len(door)):
-            for q in range(len(door)):
-                is_k = [0 for _ in range(len(door))]
-                is_k[k] = 1
-                is_q = [0 for _ in range(len(door))]
-                is_q[q] = 1
-                if how_variation == 0:
-                    batch_theta.append(np.angle(f[q]) - np.angle(f[k]))
-                else:
-                    batch_theta.append(np.angle(f[q]) - np.angle(f[k]) - math.pi / 2)
-                batch_is_k.append(is_k)
-                batch_is_q.append(is_q)
-        batch_theta = tc.array_to_tensor(batch_theta)
-        batch_is_k = tf.constant(batch_is_k)
-        batch_is_q = tf.constant(batch_is_q)
-        vmap_result = Calculation_A_vmap(batch_theta, batch_is_k, batch_is_q, ODE_theta)
+
+        vmap_result = Calculation_A_vmap(
+            a_batch_theta, a_batch_is_k, a_batch_is_q, ODE_theta
+        )
         A = tf.cast(
             tf.tensordot(tf.abs(f), tf.abs(f), 0), dtype="float64"
         ) * tf.reshape(
             tc.backend.cast(vmap_result, dtype="float64"), [len(door), len(door)]
         )
 
-        batch_theta = []
-        batch_is_k = []
-        batch_is_q = []
-        for k in range(len(door)):
-            for q in range(len(h_door)):
-                is_k = [0 for _ in range(len(door))]
-                is_k[k] = 1
-                is_q = [0 for _ in range(len(door))]
-                is_q[q] = 1
-                if how_variation == 0:
-                    batch_theta.append(np.angle(h[q]) - np.angle(f[k]) - math.pi / 2)
-                else:
-                    batch_theta.append(np.angle(h[q]) - np.angle(f[k]) + math.pi)
-                batch_is_k.append(is_k)
-                batch_is_q.append(is_q)
-        batch_theta = tc.array_to_tensor(batch_theta)
-        batch_is_k = tf.constant(batch_is_k)
-        batch_is_q = tf.constant(batch_is_q)
-        vmap_result = Calculation_C_vmap(batch_theta, batch_is_k, batch_is_q, ODE_theta)
+        vmap_result = Calculation_C_vmap(
+            c_batch_theta, c_batch_is_k, c_batch_is_q, ODE_theta
+        )
         C = tf.reduce_sum(
             tf.cast(tf.tensordot(tf.abs(f), tf.abs(h), 0), dtype="float64")
             * tf.reshape(

tensorcircuit/backends/cupy_backend.py

@@ -8,6 +8,7 @@
 from typing import Any, Callable, Optional, Sequence, Tuple, Union
 
 import numpy as np
+import scipy
 
 try:  # old version tn compatiblity
     from tensornetwork.backends import base_backend
@@ -62,6 +63,18 @@ def sum(
     ) -> Tensor:
         return cp.sum(a, axis=axis, keepdims=keepdims)
 
+    def conj(self, tensor: Tensor) -> Tensor:
+        return tensor.conj()
+
+    def sign(self, tensor: Tensor) -> Tensor:
+        return cp.sign(tensor)
+
+    def multiply(self, tensor1: Tensor, tensor2: Tensor) -> Tensor:
+        return tensor1 * tensor2
+
+    def norm(self, tensor: Tensor) -> Tensor:
+        return cp.linalg.norm(tensor)
+
     def shape_tuple(self, tensor: Tensor) -> Tuple[int]:
         return tensor.shape  # type:ignore
 
@@ -100,7 +113,7 @@ def copy(self, a: Tensor) -> Tensor:
         return a.copy()
 
     def expm(self, a: Tensor) -> Tensor:
-        raise NotImplementedError
+        return self.convert_to_tensor(scipy.linalg.expm(self.numpy(a)))
 
     def abs(self, a: Tensor) -> Tensor:
         return cp.abs(a)
@@ -307,7 +320,7 @@ def stateful_randu(
             dtype = dtype[-2:]
         if isinstance(shape, int):
             shape = (shape,)
-        r = g.uniform(low=low, high=high, size=shape)
+        r = g.random(shape) * (high - low) + low
         if dtype == "32":
             r = r.astype(np.float32)
         elif dtype == "64":
@@ -368,22 +381,22 @@ def switch(self, index: Tensor, branches: Sequence[Callable[[], Tensor]]) -> Ten
         return branches[index]()
 
     def device(self, a: Tensor) -> str:
-        return "gpu"
+        return self._dev2str(a.device)
 
     def device_move(self, a: Tensor, dev: Any) -> Tensor:
-        if dev == "gpu":
-            return a
-        raise ValueError("CuPy backend only support GPU device")
+        if isinstance(dev, str):
+            dev = self._str2dev(dev)
+        with dev:
+            return cp.asarray(a)
 
     def _dev2str(self, dev: Any) -> str:
-        if dev == "gpu":
-            return "gpu"
-        raise ValueError("CuPy backend only support GPU device")
+        return f"gpu:{dev.id}"
 
     def _str2dev(self, str_: str) -> Any:
-        if str_ == "gpu":
-            return "gpu"
-        raise ValueError("CuPy backend only support GPU device")
+        if str_ == "cpu":
+            raise ValueError("CuPy backend only support GPU device")
+        else:
+            return cp.cuda.Device(int(str_.split(":")[-1]))
 
     def stop_gradient(self, a: Tensor) -> Tensor:
         raise NotImplementedError("CuPy backend doesn't support AD")
@@ -409,7 +422,7 @@ def jit(
         f: Callable[..., Any],
         static_argnums: Optional[Union[int, Sequence[int]]] = None,
         jit_compile: Optional[bool] = None,
-        **kws: Any
+        **kws: Any,
     ) -> Callable[..., Any]:
         logger.warning("CuPy backend has no jit interface, just do nothing")
         return f

tensorcircuit/gates.py

@@ -812,7 +812,7 @@ def exponential_gate(unitary: Tensor, theta: float, name: str = "none") -> Gate:
     :return: Exponential Gate
     :rtype: Gate
     """
-    theta = num_to_tensor(theta)
+    theta, unitary = num_to_tensor(theta, unitary)
     mat = backend.expm(-backend.i() * theta * unitary)
     dimension = reduce(mul, mat.shape)
     nolegs = int(np.log(dimension) / np.log(2))

tests/conftest.py

@@ -48,6 +48,18 @@ def torchb():
         pytest.skip("****** No torch backend found, skipping test suit *******")
 
 
+@pytest.fixture(scope="function")
+def cpb():
+    try:
+        tc.set_backend("cupy")
+        yield
+        tc.set_backend("numpy")
+    except ImportError as e:
+        print(e)
+        tc.set_backend("numpy")
+        pytest.skip("****** No cupy backend found, skipping test suit *******")
+
+
 @pytest.fixture(scope="function")
 def highp():
     tc.set_dtype("complex128")

tests/test_backends.py

@@ -300,7 +300,7 @@ def test_backend_methods_2(backend):
 def test_device_cpu_only(backend):
     a = tc.backend.ones([])
     dev_str = tc.backend.device(a)
-    assert dev_str == "cpu"
+    assert dev_str in ["cpu", "gpu:0"]
     tc.backend.device_move(a, dev_str)