add numpy interface

Author: refraction-ray

Diff for: tensorcircuit/backends/abstract_backend.py

@@ -268,8 +268,6 @@ def abs(self: Any, a: Tensor) -> Tensor:
             "Backend '{}' has not implemented `abs`.".format(self.name)
         )
 
-    # TODO(@refraction-ray): abs docstring doesn't get registered in the doc
-
     def kron(self: Any, a: Tensor, b: Tensor) -> Tensor:
         """
         Return the kronecker product of two matrices ``a`` and ``b``.

Diff for: tensorcircuit/circuit.py

@@ -334,49 +334,6 @@ def _meta_apply(cls) -> None:
             for alias_gate in gate_alias[1:]:
                 setattr(cls, alias_gate, getattr(cls, present_gate))
 
-    # @classmethod
-    # def from_qcode(
-    #     cls, qcode: str
-    # ) -> "Circuit":  # forward reference, see https://github.com/python/mypy/issues/3661
-    #     """
-    #     [WIP], make circuit object from non universal simple assembly quantum language
-
-    #     :param qcode:
-    #     :type qcode: str
-    #     :return: :py:class:`Circuit` object
-    #     """
-    #     # TODO(@refraction-ray): change to OpenQASM IO
-    #     lines = [s for s in qcode.split("\n") if s.strip()]
-    #     nqubits = int(lines[0])
-    #     c = cls(nqubits)
-    #     for l in lines[1:]:
-    #         ls = [s for s in l.split(" ") if s.strip()]
-    #         g = ls[0]
-    #         index = []
-    #         errloc = 0
-    #         for i, s in enumerate(ls[1:]):
-    #             try:
-    #                 si = int(s)
-    #                 index.append(si)
-    #             except ValueError:
-    #                 errloc = i + 1
-    #                 break
-    #         kwdict = {}
-    #         if errloc > 0:
-    #             for j, s in enumerate(ls[errloc::2]):
-    #                 kwdict[s] = float(ls[2 * j + 1 + errloc])
-    #         getattr(c, g)(*index, **kwdict)
-    #     return c
-
-    # def to_qcode(self) -> str:
-    #     """
-    #     [WIP]
-
-    #     :return: qcode str of corresponding circuit
-    #     :rtype: str
-    #     """
-    #     return self._qcode
-
     def apply_single_gate(self, gate: Gate, index: int) -> None:
         """
         Apply the gate to the bit with the given index.

Diff for: tensorcircuit/interfaces/__init__.py

@@ -8,6 +8,7 @@
     general_args_to_numpy,
     general_args_to_backend,
 )
+from .numpy import numpy_interface, np_interface
 from .scipy import scipy_interface, scipy_optimize_interface
 from .torch import torch_interface, pytorch_interface
 from .tensorflow import tensorflow_interface, tf_interface

Diff for: tensorcircuit/interfaces/numpy.py

@@ -0,0 +1,61 @@
+"""
+Interface wraps quantum function as a numpy function
+"""
+
+from typing import Any, Callable
+from functools import wraps
+
+from ..cons import backend
+from .tensortrans import general_args_to_numpy, numpy_args_to_backend
+
+Tensor = Any
+
+
+def numpy_interface(
+    fun: Callable[..., Any],
+    jit: bool = True,
+) -> Callable[..., Any]:
+    """
+    Convert ``fun`` on ML backend into a numpy function
+
+    :Example:
+
+    .. code-block:: python
+
+        K = tc.set_backend("tensorflow")
+
+        def f(params, n):
+            c = tc.Circuit(n)
+            for i in range(n):
+                c.rx(i, theta=params[i])
+            for i in range(n-1):
+                c.cnot(i, i+1)
+            r = K.real(c.expectation_ps(z=[n-1]))
+            return r
+
+        n = 3
+        f_np = tc.interfaces.numpy_interface(f, jit=True)
+        f_np(np.ones([n]), n)  # 0.1577285
+
+
+    :param fun: The quantum function
+    :type fun: Callable[..., Any]
+    :param jit: whether to jit ``fun``, defaults to True
+    :type jit: bool, optional
+    :return: The numpy interface compatible version of ``fun``
+    :rtype: Callable[..., Any]
+    """
+    if jit:
+        fun = backend.jit(fun)
+
+    @wraps(fun)
+    def numpy_fun(*args: Any, **kws: Any) -> Any:
+        backend_args = numpy_args_to_backend(args)
+        r = fun(*backend_args, **kws)
+        np_r = general_args_to_numpy(r)
+        return np_r
+
+    return numpy_fun
+
+
+np_interface = numpy_interface

Diff for: tensorcircuit/interfaces/tensortrans.py

@@ -3,6 +3,7 @@
 """
 
 from typing import Any
+from functools import partial
 
 from ..cons import backend
 from ..backends import get_backend  # type: ignore
@@ -39,9 +40,17 @@ def which_backend(a: Tensor, return_backend: bool = True) -> Any:
 
 
 def tensor_to_numpy(t: Tensor) -> Array:
+    if isinstance(t, int) or isinstance(t, float):
+        return t
     return which_backend(t).numpy(t)
 
 
+def numpy_to_tensor(t: Array, backend: Any) -> Tensor:
+    if isinstance(t, int) or isinstance(t, float):
+        return t
+    return backend.convert_to_tensor(t)
+
+
 def tensor_to_dtype(t: Tensor) -> str:
     return which_backend(t).dtype(t)  # type: ignore
 
@@ -85,13 +94,13 @@ def numpy_args_to_backend(
         target_backend = get_backend(target_backend)
 
     if dtype is None:
-        return backend.tree_map(target_backend.convert_to_tensor, args)
+        return backend.tree_map(partial(numpy_to_tensor, backend=target_backend), args)
     else:
         if isinstance(dtype, str):
             leaves, treedef = backend.tree_flatten(args)
             dtype = [dtype for _ in range(len(leaves))]
             dtype = backend.tree_unflatten(treedef, dtype)
-        t = backend.tree_map(target_backend.convert_to_tensor, args)
+        t = backend.tree_map(partial(numpy_to_tensor, backend=target_backend), args)
         t = backend.tree_map(target_backend.cast, t, dtype)
         return t
 
@@ -100,7 +109,6 @@ def general_args_to_backend(
     args: Any, dtype: Any = None, target_backend: Any = None, enable_dlpack: bool = True
 ) -> Any:
     if not enable_dlpack:
-        # TODO(@refraction-ray): add device shift for numpy mediate transformation
         args = general_args_to_numpy(args)
         args = numpy_args_to_backend(args, dtype, target_backend)
         return args

Diff for: tensorcircuit/simplify.py

@@ -280,6 +280,3 @@ def _full_light_cone_cancel(nodes: List[Any]) -> List[Any]:
     while is_changed:
         nodes, is_changed = _light_cone_cancel(nodes)
     return nodes
-
-
-# TODO(@refraction-ray): utilize more simplification method in contractor preprocessing

Diff for: tests/test_interfaces.py

@@ -243,6 +243,23 @@ def f(param):
     np.testing.assert_allclose(r["fun"], -1.0, atol=1e-5)
 
 
+@pytest.mark.parametrize("backend", [lf("torchb"), lf("tfb"), lf("jaxb")])
+def test_numpy_interface(backend):
+    def f(params, n):
+        c = tc.Circuit(n)
+        for i in range(n):
+            c.rx(i, theta=params[i])
+        for i in range(n - 1):
+            c.cnot(i, i + 1)
+        r = tc.backend.real(c.expectation_ps(z=[n - 1]))
+        return r
+
+    n = 3
+    f_np = tc.interfaces.numpy_interface(f, jit=False)
+    r = f_np(np.ones([n]), n)
+    np.testing.assert_allclose(r, 0.1577285, atol=1e-5)
+
+
 @pytest.mark.parametrize("backend", [lf("npb"), lf("tfb"), lf("jaxb"), lf("torchb")])
 def test_args_transformation(backend):
     ans = tc.interfaces.general_args_to_numpy(

Diff for: tests/test_mpscircuit.py

@@ -15,10 +15,6 @@
 sys.path.insert(0, modulepath)
 import tensorcircuit as tc
 
-# TODO(@refraction-ray): mps circuit test: grad & jit, differentiable? jittable?
-# AD on jax backend may have issues for now, see
-# https://gist.github.com/refraction-ray/cc48c0b31984e6a04ee00050c0b36758
-# for a minimal demo
 
 N = 16
 D = 100