Merge branch 'master' into beta

Author: refraction-ray

CHANGELOG.md

@@ -8,6 +8,12 @@
 
 - Torch support is updraded to 2.0, and now support native vmap and native functional grad, and thus `vvag`. Still jit support is conflict with these functional transformations and be turned off by default
 
+- Add `torch_interfaces_kws` that support static keyword arguments when wrapping with the interface
+
+### Fixed
+
+- Add tests and fixed some missing methods for cupy backend, cupy backend is now ready to use (though still not guaranteed)
+
 ## 0.8.0
 
 ### Added

tensorcircuit/interfaces/__init__.py

@@ -12,5 +12,5 @@
 )
 from .numpy import numpy_interface, np_interface
 from .scipy import scipy_interface, scipy_optimize_interface
-from .torch import torch_interface, pytorch_interface
+from .torch import torch_interface, pytorch_interface, torch_interface_kws
 from .tensorflow import tensorflow_interface, tf_interface

tensorcircuit/interfaces/torch.py

@@ -2,7 +2,8 @@
 Interface wraps quantum function as a torch function
 """
 
-from typing import Any, Callable, Tuple
+from typing import Any, Callable, Dict, Tuple
+from functools import partial
 
 from ..cons import backend
 from ..utils import is_sequence
@@ -112,3 +113,51 @@ def backward(ctx: Any, *grad_y: Any) -> Any:
 
 
 pytorch_interface = torch_interface
+
+
+def torch_interface_kws(
+    f: Callable[..., Any], jit: bool = True, enable_dlpack: bool = False
+) -> Callable[..., Any]:
+    """
+    similar to py:meth:`tensorcircuit.interfaces.torch.torch_interface`,
+    but now the interface support static arguments for function ``f``,
+    which is not a tensor and can be used with keyword arguments
+
+    :Example:
+
+    .. code-block:: python
+
+        tc.set_backend("tensorflow")
+
+        def f(tensor, integer):
+            r = 0.
+            for i in range(integer):
+                r += tensor
+            return r
+
+        fnew = torch_interface_kws(f)
+
+        print(fnew(torch.ones([2]), integer=3))
+        print(fnew(torch.ones([2]), integer=4))
+
+    :param f: _description_
+    :type f: Callable[..., Any]
+    :param jit: _description_, defaults to True
+    :type jit: bool, optional
+    :param enable_dlpack: _description_, defaults to False
+    :type enable_dlpack: bool, optional
+    :return: _description_
+    :rtype: Callable[..., Any]
+    """
+    cache_dict: Dict[Tuple[Any, ...], Callable[..., Any]] = {}
+
+    def wrapper(*args: Any, **kws: Any) -> Any:
+        key = tuple([(k, v) for k, v in kws.items()])
+        if key not in cache_dict:
+            fnew = torch_interface(
+                partial(f, **kws), jit=jit, enable_dlpack=enable_dlpack
+            )
+            cache_dict[key] = fnew
+        return cache_dict[key](*args)
+
+    return wrapper

tests/test_interfaces.py

@@ -108,6 +108,74 @@ def f3(x):
     np.testing.assert_allclose(pg, 2 * np.ones([2]).astype(np.complex64), atol=1e-5)
 
 
+@pytest.mark.skipif(is_torch is False, reason="torch not installed")
+@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb")])
+def test_torch_interface_kws(backend):
+    def f(param, n):
+        c = tc.Circuit(n)
+        c = tc.templates.blocks.example_block(c, param)
+        loss = c.expectation(
+            [
+                tc.gates.x(),
+                [
+                    1,
+                ],
+            ]
+        )
+        return tc.backend.real(loss)
+
+    f_jit_torch = tc.interfaces.torch_interface_kws(f, jit=True, enable_dlpack=True)
+
+    param = torch.ones([4, 4], requires_grad=True)
+    l = f_jit_torch(param, n=4)
+    l = l**2
+    l.backward()
+
+    pg = param.grad
+    np.testing.assert_allclose(pg.shape, [4, 4])
+    np.testing.assert_allclose(pg[0, 1], -2.146e-3, atol=1e-5)
+
+    def f2(paramzz, paramx, n, nlayer):
+        c = tc.Circuit(n)
+        for i in range(n):
+            c.H(i)
+        for j in range(nlayer):  # 2
+            for i in range(n - 1):
+                c.exp1(i, i + 1, unitary=tc.gates._zz_matrix, theta=paramzz[j, i])
+            for i in range(n):
+                c.rx(i, theta=paramx[j, i])
+        loss1 = c.expectation(
+            [
+                tc.gates.x(),
+                [
+                    1,
+                ],
+            ]
+        )
+        loss2 = c.expectation(
+            [
+                tc.gates.x(),
+                [
+                    2,
+                ],
+            ]
+        )
+        return tc.backend.real(loss1), tc.backend.real(loss2)
+
+    f2_torch = tc.interfaces.torch_interface_kws(f2, jit=True, enable_dlpack=True)
+
+    paramzz = torch.ones([2, 4], requires_grad=True)
+    paramx = torch.ones([2, 4], requires_grad=True)
+
+    l1, l2 = f2_torch(paramzz, paramx, n=4, nlayer=2)
+    l = l1 - l2
+    l.backward()
+
+    pg = paramzz.grad
+    np.testing.assert_allclose(pg.shape, [2, 4])
+    np.testing.assert_allclose(pg[0, 0], -0.41609, atol=1e-5)
+
+
 @pytest.mark.skipif(is_torch is False, reason="torch not installed")
 @pytest.mark.xfail(
     (int(tf.__version__.split(".")[1]) < 9)