feat: restructure torch transformer to support __call__() of QNN

quafu/algorithms/ansatz.py

@@ -21,6 +21,7 @@
 from quafu.synthesis.evolution import ProductFormula
 from .hamiltonian import Hamiltonian
 from .interface_provider import InterfaceProvider
+from .templates import AngleEmbedding
 
 
 class Ansatz(QuantumCircuit, ABC):
@@ -155,6 +156,10 @@ def __init__(self, num_qubits: int, layers: List[Any], interface="torch"):
         self._weights = np.empty((1, 1))
         super().__init__(num_qubits)
 
+    def __call__(self, features):
+        """Compute outputs of QNN given input features"""
+        return self._transformer.execute(self, features)
+
     def _build(self):
         """Essentially initialize weights using transformer"""
         for layer in self._layers:

quafu/algorithms/interface/torch.py

@@ -20,12 +20,38 @@
 from ..gradients import compute_vjp, jacobian, run_circ
 
 
+# TODO(zhaoyilun): impl a ABC for transformers
 class TorchTransformer:
     @staticmethod
     def init_weights(shape):
         """Return torch gradient tensor with specified shape"""
         return torch.randn(*shape, requires_grad=True, dtype=torch.double)
 
+    # TODO(zhaoyilun): doc
+    @staticmethod
+    def execute(
+        circ: QuantumCircuit,
+        parameters: torch.Tensor,
+        run_fn=run_circ,
+        grad_fn=None,
+        method="internal",
+    ):
+        """execute.
+
+        Args:
+            circ:
+            run_fn:
+            grad_fn:
+        """
+
+        kwargs = {"circ": circ, "run_fn": run_fn, "grad_fn": grad_fn}
+
+        if method == "external":
+            return ExecuteCircuits.apply(parameters, kwargs)
+        if method == "internal":
+            return ExecuteCircuits.apply(circ.weights, kwargs)
+        raise NotImplementedError(f"Unsupported execution method: {method}")
+
 
 class ExecuteCircuits(torch.autograd.Function):
     """Parameters are input from previous layers"""
@@ -51,29 +77,3 @@ def backward(ctx, grad_out):
         vjp = compute_vjp(jac, grad_out.numpy())
         vjp = torch.from_numpy(vjp)
         return vjp, None
-
-
-# TODO(zhaoyilun): doc
-def execute(
-    circ: QuantumCircuit,
-    parameters: torch.Tensor,
-    run_fn=run_circ,
-    grad_fn=None,
-    method="internal",
-):
-    """execute.
-
-    Args:
-        circ:
-        run_fn:
-        grad_fn:
-    """
-
-    kwargs = {"circ": circ, "run_fn": run_fn, "grad_fn": grad_fn}
-
-    if method == "external":
-        return ExecuteCircuits.apply(parameters, kwargs)
-    elif method == "internal":
-        return ExecuteCircuits.apply(circ.weights, kwargs)
-    else:
-        raise NotImplementedError(f"Unsupported execution method: {method}")

quafu/algorithms/templates/__init__.py

@@ -0,0 +1,16 @@
+# (C) Copyright 2023 Beijing Academy of Quantum Information Sciences
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .angle import AngleEmbedding
+from .basic_entangle import BasicEntangleLayers

tests/quafu/algorithms/integration_test.py

@@ -61,7 +61,7 @@ def test_run(self):
         """
         num_layers = 2
         print("The test for ansatz.")
-        
+
         # test the zero qubit evolution
         hamiltonian__ = Hamiltonian.from_pauli_list(
             [("IIIII", 1), ("IIIII", 1), ("IIIII", 1), ("IIIII", 1)]

tests/quafu/algorithms/qnn_test.py

@@ -15,7 +15,7 @@
 
 import torch.nn
 from quafu.algorithms.gradients import jacobian, compute_vjp
-from quafu.algorithms.interface.torch import execute
+from quafu.algorithms.interface.torch import TorchTransformer
 from quafu.algorithms.templates.basic_entangle import BasicEntangleLayers
 from quafu.circuits.quantum_circuit import QuantumCircuit
 from quafu.algorithms import QuantumNeuralNetwork
@@ -29,7 +29,7 @@ def __init__(self, circ: QuantumCircuit):
 
     def forward(self, features):
         out = self.linear(features)
-        out = execute(self.circ, out, method="external")
+        out = TorchTransformer.execute(self.circ, out, method="external")
         return out
 
 
@@ -39,7 +39,19 @@ def __init__(self, circ: QuantumNeuralNetwork):
         self.circ = circ
 
     def forward(self, features):
-        out = execute(self.circ, features)
+        out = TorchTransformer.execute(self.circ, features)
+        return out
+
+
+class ModelQuantumNeuralNetworkNative(torch.nn.Module):
+    """Test execution of qnn()"""
+
+    def __init__(self, qnn: QuantumNeuralNetwork):
+        super().__init__()
+        self.qnn = qnn
+
+    def forward(self, features):
+        out = self.qnn(features)
         return out
 
 
@@ -50,6 +62,19 @@ class TestLayers:
     circ.ry(1, 0.5)
     circ.ry(0, 0.1)
 
+    def _model_grad(self, model, batch_size):
+        """Test one forward pass and gradient calculation of a model"""
+
+        # TODO(zhaoyilun): Make out dimension configurable
+        features = torch.randn(
+            batch_size, 3, requires_grad=True, dtype=torch.double
+        )  # batch_size=4, num_params=3
+        outputs = model(features)
+        targets = torch.randn(batch_size, 2, dtype=torch.double)
+        criterion = torch.nn.MSELoss()
+        loss = criterion(outputs, targets)
+        loss.backward()
+
     def test_compute_vjp(self):
         params_input = np.random.randn(4, 3)
         jac = jacobian(self.circ, params_input)
@@ -78,12 +103,11 @@ def test_torch_layer_qnn(self):
         entangle_layer = BasicEntangleLayers(weights, 2)
         qnn = QuantumNeuralNetwork(2, [entangle_layer])
         batch_size = 1
+
+        # Legacy invokation style
         model = ModelQuantumNeuralNetwork(qnn)
-        features = torch.randn(
-            batch_size, 3, requires_grad=True, dtype=torch.double
-        )  # batch_size=4, num_params=3
-        outputs = model(features)
-        targets = torch.randn(batch_size, 2, dtype=torch.double)
-        criterion = torch.nn.MSELoss()
-        loss = criterion(outputs, targets)
-        loss.backward()
+        self._model_grad(model, batch_size)
+
+        # New invokation style
+        model = ModelQuantumNeuralNetworkNative(qnn)
+        self._model_grad(model, batch_size)