add torchnn module

Author: refraction-ray

CHANGELOG.md

@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Added
+
+- Add PyTorch nn Module wrapper in `torchnn`
+
 ## 0.1.3
 
 ### Added
@@ -32,7 +36,7 @@
 
 - Fix installation issue with tensorflow requirements on MACOS with M1 chip
 
-- Improve M1 macOS compayibility with unjit tensorflow ops
+- Improve M1 macOS compatibility with unjit tensorflow ops
 
 - Fixed SVD backprop bug on jax backend of wide matrix
 

docs/source/infras.rst

@@ -34,6 +34,8 @@ Overview of Modules
 
 - :py:mod:`tensorcircuit.keras`: Provide TensorFlow Keras layers, as well as wrappers of jitted function, save/load from tf side.
 
+- :py:mod:`tensorcircuit.torchnn`: Provide PyTorch nn Modules.
+
 **MPS and MPO Utiliy Modules:**
 
 - :py:mod:`tensorcircuit.quantum`: Provide definition and classes for Matrix Product States as well as Matrix Product Operators, we also include various quantum physics and quantum information quantities in this module.

tensorcircuit/__init__.py

@@ -31,5 +31,10 @@
 except ModuleNotFoundError:
     pass  # in case tf is not installed
 
+try:
+    from . import torchnn
+except ModuleNotFoundError:
+    pass  # in case torch is not installed
+
 # just for fun
 from .asciiart import set_ascii

tensorcircuit/cons.py

@@ -31,6 +31,7 @@
     "tensorcircuit.densitymatrix2",
     "tensorcircuit.channels",
     "tensorcircuit.keras",
+    "tensorcircuit.torchnn",
     "tensorcircuit.quantum",
     "tensorcircuit.simplify",
     "tensorcircuit.interfaces",

tensorcircuit/torchnn.py

@@ -0,0 +1,43 @@
+"""
+PyTorch nn Module wrapper for quantum function
+"""
+
+from typing import Any, Callable, Sequence, Tuple, Union
+
+import torch
+
+from .cons import backend
+from .interfaces import torch_interface, is_sequence
+
+Tensor = Any
+
+
+class QuantumNet(torch.nn.Module):  # type: ignore
+    def __init__(
+        self,
+        f: Callable[..., Any],
+        weights_shape: Sequence[Tuple[int, ...]],
+        initializer: Union[Any, Sequence[Any]] = None,
+        use_vmap: bool = True,
+        use_interface: bool = True,
+        use_jit: bool = True,
+    ):
+        super().__init__()
+        if use_vmap:
+            f = backend.vmap(f, vectorized_argnums=0)
+        if use_interface:
+            f = torch_interface(f, jit=use_jit)
+        self.f = f
+        self.q_weights = []
+        if isinstance(weights_shape[0], int):
+            weights_shape = [weights_shape]
+        if not is_sequence(initializer):
+            initializer = [initializer]
+        for ws, initf in zip(weights_shape, initializer):
+            if initf is None:
+                initf = torch.randn
+            self.q_weights.append(torch.nn.Parameter(initf(ws)))  # type: ignore
+
+    def forward(self, inputs: Tensor) -> Tensor:
+        ypred = self.f(inputs, *self.q_weights)
+        return ypred

tests/test_torchnn.py

@@ -0,0 +1,51 @@
+import os
+import sys
+import numpy as np
+import pytest
+from pytest_lazyfixture import lazy_fixture as lf
+
+thisfile = os.path.abspath(__file__)
+modulepath = os.path.dirname(os.path.dirname(thisfile))
+
+sys.path.insert(0, modulepath)
+
+import tensorcircuit as tc
+
+try:
+    import torch
+except ImportError:
+    pytest.skip("torch is not installed")
+
+
+@pytest.mark.parametrize("backend", [lf("tfb"), lf("jaxb"), lf("torchb")])
+def test_quantumnet(backend):
+
+    n = 6
+    nlayers = 2
+
+    def qpred(x, weights):
+        c = tc.Circuit(n)
+        for i in range(n):
+            c.rx(i, theta=x[i])
+        for j in range(nlayers):
+            for i in range(n - 1):
+                c.cnot(i, i + 1)
+            for i in range(n):
+                c.rx(i, theta=weights[2 * j, i])
+                c.ry(i, theta=weights[2 * j + 1, i])
+        ypred = tc.backend.stack([c.expectation_ps(x=[i]) for i in range(n)])
+        return tc.backend.real(ypred)
+
+    if tc.backend.name == "pytorch":
+        use_interface = False
+    else:
+        use_interface = True
+
+    ql = tc.torchnn.QuantumNet(
+        qpred, weights_shape=[2 * nlayers, n], use_interface=use_interface
+    )
+
+    yp = ql(torch.ones([3, n]))
+    print(yp)
+
+    np.testing.assert_allclose(yp.shape, np.array([3, n]))