add pytree support

Author: refraction-ray

CHANGELOG.md

@@ -6,6 +6,18 @@
 
 - Add more type auto conversion for `tc.gates.Gate` as inputs
 
+- Add `tree_flatten` and `tree_unflatten` method on backends
+
+- Add torch optimizer to the backend agnostic optimizer abstraction
+
+### Changed
+
+- Refactor the tree utils, add native torch support for pytree utils
+
+### Fixed
+
+- grad in torch backend now support pytrees
+
 ## 0.1.2
 
 ### Added

docs/source/textbook/chap4.ipynb

@@ -743,7 +743,9 @@
     "\n",
     "\n",
     "\n",
-    "* > 关于Grover search的进一步说明：我们从以上的推导可以发现，Grover Search 应用于搜索某个问题的解。在经典计算中，对某些问题，最好的方法也只能通过暴力尝试所有$O(N)$个输入并计算$f(x)$来找到解，而Grover search 可以将尝试的次数减少到$O(\\sqrt{N})$。如果每一次经典的尝试和量子的尝试相差的时间不大的话，那么量子计算就会比经典计算快得多。值得注意的是，这种搜索与数据库的搜索并不一致，是无序无结构的搜索。Grover search原则上也可以进行数据库搜索，但其相对于经典算法的优越性则要看数据存储的结构。"
+    "* > 关于Grover search的进一步说明：我们从以上的推导可以发现，Grover Search 应用于搜索某个问题的解。在经典计算中，对某些问题，最好的方法也只能通过暴力尝试所有$O(N)$个输入并计算$f(x)$来找到解，而Grover search 可以将尝试的次数减少到$O(\\sqrt{N})$。如果每一次经典的尝试和量子的尝试相差的时间不大的话，那么量子计算就会比经典计算快得多。值得注意的是，这种搜索与数据库的搜索并不一致，是无序无结构的搜索。Grover search原则上也可以进行数据库搜索，但其相对于经典算法的优越性则要看数据存储的结构。\n",
+    "\n",
+    "更多关于基础 Grover 搜索算法的扩展和证明，可以参考【1】。"
    ]
   },
   {
@@ -838,6 +840,15 @@
    "source": [
     "c.sample()"
    ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## 参考文献\n",
+    "\n",
+    "【1】https://arxiv.org/pdf/quant-ph/9605034.pdf"
+   ]
   }
  ],
  "metadata": {

tensorcircuit/backends/abstract_backend.py

@@ -672,26 +672,52 @@ def tree_map(self: Any, f: Callable[..., Any], *pytrees: Any) -> Any:
         :rtype: Any
         """
         try:
-            import jax as libjax
+            import tensorflow as tf
 
-            has_jax = True
         except ImportError:
-            has_jax = False
-            try:
-                import tensorflow as tf
+            raise NotImplementedError("No installed ML backend for `tree_map`")
 
-                has_tf = True
-            except ImportError:
-                has_tf = False
+        return tf.nest.map_structure(f, *pytrees)
 
-        if has_jax:
-            r = libjax.tree_map(f, *pytrees)
-        elif has_tf:
-            r = tf.nest.map_structure(f, *pytrees)
-        else:
-            raise NotImplementedError("Only tensorflow and jax support `tree_map`")
+    def tree_flatten(self: Any, pytree: Any) -> Tuple[Any, Any]:
+        """
+        Flatten python structure to 1D list
 
-        return r
+        :param pytree: python structure to be flattened
+        :type pytree: Any
+        :return: The 1D list of flattened structure and treedef
+            which can be used for later unflatten
+        :rtype: Tuple[Any, Any]
+        """
+        try:
+            import tensorflow as tf
+
+        except ImportError:
+            raise NotImplementedError("No installed ML backend for `tree_flatten`")
+
+        leaves = tf.nest.flatten(pytree)
+        treedef = pytree
+
+        return leaves, treedef
+
+    def tree_unflatten(self: Any, treedef: Any, leaves: Any) -> Any:
+        """
+        Pack 1D list to pytree defined via ``treedef``
+
+        :param treedef: Def of pytree structure, the second return from ``tree_flatten``
+        :type treedef: Any
+        :param leaves: the 1D list of flattened data structure
+        :type leaves: Any
+        :return: Packed pytree
+        :rtype: Any
+        """
+        try:
+            import tensorflow as tf
+
+        except ImportError:
+            raise NotImplementedError("No installed ML backend for `tree_unflatten`")
+
+        return tf.nest.pack_sequence_as(treedef, leaves)
 
     def set_random_state(
         self: Any, seed: Optional[int] = None, get_only: bool = False

tensorcircuit/backends/jax_backend.py

@@ -378,6 +378,15 @@ def is_tensor(self, a: Any) -> bool:
     def solve(self, A: Tensor, b: Tensor, assume_a: str = "gen") -> Tensor:
         return jsp.linalg.solve(A, b, assume_a)
 
+    def tree_map(self, f: Callable[..., Any], *pytrees: Any) -> Any:
+        return libjax.tree_map(f, *pytrees)
+
+    def tree_flatten(self: Any, pytree: Any) -> Tuple[Any, Any]:
+        return libjax.tree_flatten(pytree)  # type: ignore
+
+    def tree_unflatten(self: Any, treedef: Any, leaves: Any) -> Any:
+        return libjax.tree_unflatten(treedef, leaves)
+
     def set_random_state(
         self, seed: Optional[Union[int, PRNGKeyArray]] = None, get_only: bool = False
     ) -> Any:

tensorcircuit/backends/pytorch_backend.py

@@ -23,6 +23,7 @@
 
 dtypestr: str
 Tensor = Any
+pytree = Any
 
 torchlib: Any
 
@@ -34,6 +35,28 @@
 # To be added once pytorch backend is ready
 
 
+class torch_optimizer:
+    def __init__(self, optimizer: Any) -> None:
+        self.optimizer = optimizer
+        self.is_init = False
+
+    def update(self, grads: pytree, params: pytree) -> pytree:
+        # flatten grad and param
+        params, treedef = PyTorchBackend.tree_flatten(None, params)
+        grads, _ = PyTorchBackend.tree_flatten(None, grads)
+        if self.is_init is False:
+            self.optimizer = self.optimizer(params)
+            self.is_init = True
+        with torchlib.no_grad():
+            for g, p in zip(grads, params):
+                p.grad = g
+        self.optimizer.step()
+        self.optimizer.zero_grad()
+        # reorg the param
+        params = PyTorchBackend.tree_unflatten(None, treedef, params)
+        return params
+
+
 def _conj_torch(self: Any, tensor: Tensor) -> Tensor:
     t = torchlib.conj(tensor)
     return t.resolve_conj()  # any side effect?
@@ -355,6 +378,16 @@ def cast(self, a: Tensor, dtype: str) -> Tensor:
     def solve(self, A: Tensor, b: Tensor, **kws: Any) -> Tensor:
         return torchlib.linalg.solve(A, b)
 
+    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)
+
+    def tree_flatten(self: Any, pytree: Any) -> Tuple[Any, Any]:
+        return torchlib.utils._pytree.tree_flatten(pytree)  # type: ignore
+
+    def tree_unflatten(self: Any, treedef: Any, leaves: Any) -> Any:
+        return torchlib.utils._pytree.tree_unflatten(leaves, treedef)
+
     def cond(
         self,
         pred: bool,
@@ -413,6 +446,13 @@ def value_and_grad(
         argnums: Union[int, Sequence[int]] = 0,
         has_aux: bool = False,
     ) -> Callable[..., Tuple[Any, Any]]:
+        def ask_require(t: Tensor) -> Any:
+            t.requires_grad_(True)
+            return t
+
+        def get_grad(t: Tensor) -> Tensor:
+            return t.grad
+
         def wrapper(*args: Any, **kws: Any) -> Any:
             x = []
             if isinstance(argnums, int):
@@ -423,15 +463,15 @@ def wrapper(*args: Any, **kws: Any) -> Any:
                 argnumsl = argnums  # type: ignore
             for i, arg in enumerate(args):
                 if i in argnumsl:
-                    x.append(arg.requires_grad_(True))
+                    x.append(self.tree_map(ask_require, arg))
                 else:
                     x.append(arg)
             y = f(*x, **kws)
             if has_aux:
                 y[0].backward()
             else:
                 y.backward()
-            gs = [x[i].grad for i in argnumsl]
+            gs = [self.tree_map(get_grad, x[i]) for i in argnumsl]
             if len(gs) == 1:
                 gs = gs[0]
             return y, gs
@@ -532,3 +572,5 @@ def vectorized_value_and_grad(
         return f
 
     vvag = vectorized_value_and_grad
+
+    optimizer = torch_optimizer

tests/conftest.py

@@ -40,14 +40,11 @@ def jaxb():
 def torchb():
     try:
         tc.set_backend("pytorch")
-        tc.set_dtype("float64")
         yield
         tc.set_backend("numpy")
-        tc.set_dtype("complex64")
     except ImportError as e:
         print(e)
         tc.set_backend("numpy")
-        tc.set_dtype("complex64")
         pytest.skip("****** No torch backend found, skipping test suit *******")
 
 

tests/test_backends.py

@@ -779,14 +779,29 @@ def test_solve(backend):
     np.testing.assert_allclose(xp, x[:, 0], atol=1e-5)
 
 
-@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
+@pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
+def test_treeutils(backend):
+    d0 = {"a": np.ones([2]), "b": [tc.backend.zeros([]), tc.backend.ones([1, 1])]}
+    leaves, treedef = tc.backend.tree_flatten(d0)
+    d1 = tc.backend.tree_unflatten(treedef, leaves)
+    d2 = tc.backend.tree_map(lambda x: 2 * x, d1)
+    np.testing.assert_allclose(2 * np.ones([1, 1]), d2["b"][1])
+
+
+@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("torchb")])
 def test_optimizers(backend):
     if tc.backend.name == "jax":
         try:
             import optax
         except ImportError:
             pytest.skip("optax is not installed")
 
+    if tc.backend.name == "pytorch":
+        try:
+            import torch
+        except ImportError:
+            pytest.skip("torch is not installed")
+
     def f(params, n):
         c = tc.Circuit(n)
         c = tc.templates.blocks.example_block(c, params["a"])
@@ -802,6 +817,9 @@ def get_opt():
         elif tc.backend.name == "jax":
             optimizer2 = optax.adam(5e-2)
             opt = tc.backend.optimizer(optimizer2)
+        elif tc.backend.name == "pytorch":
+            optimizer3 = partial(torch.optim.Adam, lr=5e-2)
+            opt = tc.backend.optimizer(optimizer3)
         else:
             raise ValueError("%s doesn't support optimizer interface" % tc.backend.name)
         return opt
@@ -810,8 +828,8 @@ def get_opt():
     opt = get_opt()
 
     params = {
-        "a": tc.backend.implicit_randn([4, n]),
-        "b": tc.backend.implicit_randn([4, n]),
+        "a": tc.backend.ones([4, n], dtype="float32"),
+        "b": tc.backend.ones([4, n], dtype="float32"),
     }
 
     for _ in range(20):
@@ -828,12 +846,11 @@ def f2(params, n):
 
     vgs2 = tc.backend.jit(tc.backend.value_and_grad(f2, argnums=0), static_argnums=1)
 
-    params = tc.backend.implicit_randn([4, n])
+    params = tc.backend.ones([4, n], dtype="float32")
     opt = get_opt()
 
     for _ in range(20):
         loss, grads = vgs2(params, n)
-        print(grads, params)
         params = opt.update(grads, params)
         print(loss)