From 1102d943a844540328c137a63ad5794682a58684 Mon Sep 17 00:00:00 2001
From: Zhaoyilunnn <dilved1014@gmail.com>
Date: Sat, 21 Oct 2023 07:16:06 +0000
Subject: [PATCH 1/6] feat: param shift autograd demo -- vjp computation

---
 quafu/algorithms/__init__.py                  |  2 +-
 quafu/algorithms/estimator.py                 | 22 +++--
 quafu/algorithms/gradients/__init__.py        | 15 ++++
 .../param_shift.py}                           | 18 ++--
 quafu/algorithms/hamiltonian.py               | 29 +++----
 quafu/algorithms/layers/__init__.py           |  1 +
 quafu/algorithms/layers/library.py            | 15 ++++
 quafu/algorithms/layers/qnode.py              | 84 +++++++++++++++++++
 quafu/algorithms/layers/torch.py              | 53 ++++++++++++
 quafu/simulators/torch.py                     | 15 ++++
 .../{optimizer_test.py => gradient_test.py}   |  2 +-
 tests/quafu/algorithms/layers_test.py         | 34 ++++++++
 12 files changed, 262 insertions(+), 28 deletions(-)
 create mode 100644 quafu/algorithms/gradients/__init__.py
 rename quafu/algorithms/{optimizer.py => gradients/param_shift.py} (83%)
 create mode 100644 quafu/algorithms/layers/__init__.py
 create mode 100644 quafu/algorithms/layers/library.py
 create mode 100644 quafu/algorithms/layers/qnode.py
 create mode 100644 quafu/algorithms/layers/torch.py
 create mode 100644 quafu/simulators/torch.py
 rename tests/quafu/algorithms/{optimizer_test.py => gradient_test.py} (96%)
 create mode 100644 tests/quafu/algorithms/layers_test.py

diff --git a/quafu/algorithms/__init__.py b/quafu/algorithms/__init__.py
index 2665fbf..efa3209 100644
--- a/quafu/algorithms/__init__.py
+++ b/quafu/algorithms/__init__.py
@@ -1,5 +1,5 @@
 """Algorithm module"""
 
-from .hamiltonian import *
+from .hamiltonian import Hamiltonian
 from .ansatz import *
 from .estimator import *
diff --git a/quafu/algorithms/estimator.py b/quafu/algorithms/estimator.py
index 6041ac8..cdba8c4 100644
--- a/quafu/algorithms/estimator.py
+++ b/quafu/algorithms/estimator.py
@@ -21,6 +21,17 @@
 from quafu.algorithms.hamiltonian import Hamiltonian
 
 
+def execute_circuit(circ: QuantumCircuit, observables: Hamiltonian):
+    """Execute circuit on quafu simulator"""
+    sim_state = simulate(circ, output="state_vector").get_statevector()
+
+    expectation = np.matmul(
+        np.matmul(sim_state.conj().T, observables.get_matrix()), sim_state
+    ).real
+
+    return expectation
+
+
 class Estimator:
     """Estimate expectation for quantum circuits and observables"""
 
@@ -62,11 +73,12 @@ def _run_real_machine(self, observables: Hamiltonian):
 
     def _run_simulation(self, observables: Hamiltonian):
         """Run using quafu simulator"""
-        sim_state = simulate(self._circ, output="state_vector").get_statevector()
-        expectation = np.matmul(
-            np.matmul(sim_state.conj().T, observables.get_matrix()), sim_state
-        ).real
-        return expectation
+        # sim_state = simulate(self._circ, output="state_vector").get_statevector()
+        # expectation = np.matmul(
+        #     np.matmul(sim_state.conj().T, observables.get_matrix()), sim_state
+        # ).real
+        # return expectation
+        return execute_circuit(self._circ, observables)
 
     def run(self, observables: Hamiltonian, params: List[float]):
         """Calculate estimation for given observables
diff --git a/quafu/algorithms/gradients/__init__.py b/quafu/algorithms/gradients/__init__.py
new file mode 100644
index 0000000..566181f
--- /dev/null
+++ b/quafu/algorithms/gradients/__init__.py
@@ -0,0 +1,15 @@
+# (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 .param_shift import ParamShift
diff --git a/quafu/algorithms/optimizer.py b/quafu/algorithms/gradients/param_shift.py
similarity index 83%
rename from quafu/algorithms/optimizer.py
rename to quafu/algorithms/gradients/param_shift.py
index 69fa52f..5eba3bc 100644
--- a/quafu/algorithms/optimizer.py
+++ b/quafu/algorithms/gradients/param_shift.py
@@ -11,9 +11,11 @@
 # 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.
+"""Quafu parameter shift"""
 
-import numpy as np
 from typing import List
+import numpy as np
+
 from quafu.algorithms import Estimator
 from quafu.algorithms.hamiltonian import Hamiltonian
 
@@ -31,7 +33,7 @@ def __call__(self, obs: Hamiltonian, params: List[float]):
             estimator (Estimator): estimator to calculate expectation values
             params (List[float]): params to optimize
         """
-        return self._grad(obs, params)
+        return self.grad(obs, params)
 
     def _gen_param_shift_vals(self, params):
         """Given a param list with n values, replicate to 2*n param list"""
@@ -42,13 +44,19 @@ def _gen_param_shift_vals(self, params):
         minus_params = params - offsets * np.pi / 2
         return plus_params.tolist() + minus_params.tolist()
 
-    def _grad(self, obs: Hamiltonian, params: List[float]):
+    def grad(self, obs: Hamiltonian, params: List[float]):
+        """grad.
+
+        Args:
+            obs (Hamiltonian): obs
+            params (List[float]): params
+        """
         shifted_params_lists = self._gen_param_shift_vals(params)
 
         res = np.zeros(len(shifted_params_lists))
         for i, shifted_params in enumerate(shifted_params_lists):
             res[i] = self._est.run(obs, shifted_params)
 
-        n = len(res)
-        grads = (res[: n // 2] - res[n // 2 :]) / 2
+        num_shift_params = len(res)
+        grads = (res[: num_shift_params // 2] - res[num_shift_params // 2 :]) / 2
         return grads
diff --git a/quafu/algorithms/hamiltonian.py b/quafu/algorithms/hamiltonian.py
index 61fbf56..071875b 100644
--- a/quafu/algorithms/hamiltonian.py
+++ b/quafu/algorithms/hamiltonian.py
@@ -66,7 +66,6 @@ def from_pauli_list(pauli_list: Iterable[tuple[str, complex]]) -> Hamiltonian:
         if size == 0:
             raise QuafuError("Pauli list cannot be empty.")
 
-        num_qubits = len(pauli_list[0][0])
         coeffs = np.zeros(size, dtype=complex)
 
         pauli_str_list = []
@@ -81,31 +80,29 @@ def to_legacy_quafu_pauli_list(self):
         """Transform to legacy quafu pauli list format,
         this is a temperal function and should be deleted later"""
         res = []
-        for pauli in self._pauli_str_list:
-            for i, p in enumerate(pauli[::-1]):
-                if p in ["X", "Y", "Z"]:
-                    res.append([p, [i]])
+        for pauli_str in self._pauli_str_list:
+            for i, pauli in enumerate(pauli_str[::-1]):
+                if pauli in ["X", "Y", "Z"]:
+                    res.append([pauli, [i]])
         return res
 
     def _get_pauli_mat(self, pauli_str: str):
         """Calculate the matrix of a pauli string"""
         mat = None
-        for p in pauli_str[::-1]:
-            mat = PAULI_MAT[p] if mat is None else np.kron(PAULI_MAT[p], mat)
+        for pauli in pauli_str[::-1]:
+            mat = PAULI_MAT[pauli] if mat is None else np.kron(PAULI_MAT[pauli], mat)
         return mat
 
     def matrix_generator(self):
         """Generating matrix for each Pauli str"""
-        for i, p in enumerate(self._pauli_str_list):
-            yield self._coeffs[i] * self._get_pauli_mat(p)
+        for i, pauli_str in enumerate(self._pauli_str_list):
+            yield self._coeffs[i] * self._get_pauli_mat(pauli_str)
 
     def get_matrix(self):
         """Generate matrix of Hamiltonian"""
 
-        mat = None
-        for m in self.matrix_generator():
-            if mat is None:
-                mat = m
-                continue
-            mat += m
-        return mat
+        dim = 2**self.num_qubits
+        matrix = np.zeros((dim, dim), dtype=complex)
+        for mat in self.matrix_generator():
+            matrix += mat
+        return matrix
diff --git a/quafu/algorithms/layers/__init__.py b/quafu/algorithms/layers/__init__.py
new file mode 100644
index 0000000..90a2e83
--- /dev/null
+++ b/quafu/algorithms/layers/__init__.py
@@ -0,0 +1 @@
+from .qnode import compute_vjp, jacobian
diff --git a/quafu/algorithms/layers/library.py b/quafu/algorithms/layers/library.py
new file mode 100644
index 0000000..769fe20
--- /dev/null
+++ b/quafu/algorithms/layers/library.py
@@ -0,0 +1,15 @@
+# (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.
+
+"""Pre-built quafu circuit blocks"""
diff --git a/quafu/algorithms/layers/qnode.py b/quafu/algorithms/layers/qnode.py
new file mode 100644
index 0000000..a1cdf61
--- /dev/null
+++ b/quafu/algorithms/layers/qnode.py
@@ -0,0 +1,84 @@
+# (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.
+
+
+import numpy as np
+from quafu import QuantumCircuit
+from quafu.algorithms import Hamiltonian
+from quafu.algorithms.estimator import Estimator
+from quafu.algorithms.gradients import ParamShift
+
+
+def _generate_expval_z(num_qubits: int):
+    obs_list = []
+    base_pauli = "I" * num_qubits
+    for i in range(num_qubits):
+        pauli = base_pauli[:i] + "Z" + base_pauli[i + 1 :]
+        obs_list.append(Hamiltonian.from_pauli_list([(pauli, 1)]))
+    return obs_list
+
+
+# TODO(zhaoyilun): support more measurement types
+def execute(circ: QuantumCircuit):
+    """execute.
+
+    Args:
+        circ (QuantumCircuit): circ
+    """
+    obs_list = _generate_expval_z(circ.num)
+    estimator = Estimator(circ)
+    params = [g.paras for g in circ.parameterized_gates]
+    output = [estimator.run(obs, params) for obs in obs_list]
+    return np.array(output)
+
+
+# TODO(zhaoyilun): support more gradient methods
+def jacobian(circ: QuantumCircuit, params_input: np.ndarray):
+    """Calculate Jacobian matrix
+
+    Args:
+        circ (QuantumCircuit): circ
+        params_input (np.ndarray): params_input, with shape [batch_size, num_params]
+    """
+    batch_size, num_params = params_input.shape
+    obs_list = _generate_expval_z(circ.num)
+    num_expvals = len(obs_list)
+    estimator = Estimator(circ)
+    calc_grad = ParamShift(estimator)
+    output = np.zeros((batch_size, num_expvals, num_params))
+    for i in range(batch_size):
+        grad_list = [
+            np.array(calc_grad(obs, params_input[i, :].tolist())) for obs in obs_list
+        ]
+        output[i, :, :] = np.stack(grad_list)
+    return output
+
+
+def compute_vjp(jac: np.ndarray, dy: np.ndarray):
+    """compute vector-jacobian product
+
+    Args:
+        jac (np.ndarray): jac with shape (batch_size, num_expvals, num_params)
+        dy (np.ndarray): dy
+    """
+    batch_size, num_expvals, num_params = jac.shape
+    assert dy.shape[0] == batch_size and dy.shape[1] == num_expvals
+
+    vjp = np.zeros((batch_size, num_params))
+
+    for i in range(batch_size):
+        vjp[i] = dy[i, :].T @ jac[i, :, :]
+
+    vjp = vjp.sum(0)
+    return vjp
diff --git a/quafu/algorithms/layers/torch.py b/quafu/algorithms/layers/torch.py
new file mode 100644
index 0000000..c886de5
--- /dev/null
+++ b/quafu/algorithms/layers/torch.py
@@ -0,0 +1,53 @@
+# (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.
+
+"""quafu PyTorch quantum layer"""
+
+from typing import Any
+import torch
+from quafu import QuantumCircuit
+from quafu import QuantumCircuit, simulate
+from quafu.elements.element_gates.matrices import ZMatrix
+
+
+class ExecuteCircuits(torch.autograd.Function):
+    """TODO(zhaoyilun): document"""
+
+    @staticmethod
+    def forward(ctx, parameters, **kwargs) -> Any:
+        ctx.run_fn = kwargs["run_fn"]
+        ctx.circ = kwargs["circ"]
+        out = ctx.run_fn(ctx.circ, parameters)
+        return out
+
+    @staticmethod
+    def backward(ctx: Any, grad_out) -> Any:
+        circ, grad_fn = ctx.saved_tensors
+        grad = grad_fn(circ, grad_out)
+        return grad
+
+
+# TODO(zhaoyilun): doc
+def execute(circ: QuantumCircuit, run_fn, grad_fn, parameters: torch.Tensor):
+    """execute.
+
+    Args:
+        circ:
+        run_fn:
+        grad_fn:
+    """
+
+    kwargs = {"circ": circ, "run_fn": run_fn, "grad_fn": grad_fn}
+
+    return ExecuteCircuits.apply(parameters, **kwargs)
diff --git a/quafu/simulators/torch.py b/quafu/simulators/torch.py
new file mode 100644
index 0000000..57c9f07
--- /dev/null
+++ b/quafu/simulators/torch.py
@@ -0,0 +1,15 @@
+# (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.
+
+"""Simulate the execution of a quantum circuit using pytorch"""
diff --git a/tests/quafu/algorithms/optimizer_test.py b/tests/quafu/algorithms/gradient_test.py
similarity index 96%
rename from tests/quafu/algorithms/optimizer_test.py
rename to tests/quafu/algorithms/gradient_test.py
index b8ede14..685f113 100644
--- a/tests/quafu/algorithms/optimizer_test.py
+++ b/tests/quafu/algorithms/gradient_test.py
@@ -16,7 +16,7 @@
 import sys
 from quafu.algorithms.estimator import Estimator
 from quafu.algorithms.hamiltonian import Hamiltonian
-from quafu.algorithms.optimizer import ParamShift
+from quafu.algorithms.gradients import ParamShift
 from quafu.circuits.quantum_circuit import QuantumCircuit
 
 
diff --git a/tests/quafu/algorithms/layers_test.py b/tests/quafu/algorithms/layers_test.py
new file mode 100644
index 0000000..6a9c5bf
--- /dev/null
+++ b/tests/quafu/algorithms/layers_test.py
@@ -0,0 +1,34 @@
+# (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.
+import numpy as np
+
+from quafu.algorithms.layers import jacobian, compute_vjp
+from quafu.circuits.quantum_circuit import QuantumCircuit
+
+
+def test_compute_vjp():
+    circ = QuantumCircuit(2)
+    circ.x(0)
+    circ.rx(0, 0.1)
+    circ.ry(1, 0.5)
+    circ.ry(0, 0.1)
+
+    params_input = np.random.randn(4, 3)
+    jac = jacobian(circ, params_input)
+
+    dy = np.random.randn(4, 2)
+    vjp = compute_vjp(jac, dy)
+
+    assert len(vjp.shape) == 1
+    assert vjp.shape[0] == 3

From 31c438654ed7389901c2fce32f6e30d9616ac857 Mon Sep 17 00:00:00 2001
From: Zhaoyilunnn <dilved1014@gmail.com>
Date: Sat, 28 Oct 2023 15:45:30 +0000
Subject: [PATCH 2/6] fix(workflow): directly install wheel for ut

---
 .github/workflows/wheels.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
index dace014..3186a1f 100644
--- a/.github/workflows/wheels.yml
+++ b/.github/workflows/wheels.yml
@@ -49,7 +49,7 @@ jobs:
         run: python -m cibuildwheel --output-dir dist
 
       - name: Run unit tests
-        run: pip install . && pytest tests/
+        run: pip install dist/*.whl && pytest tests/
 
       - name: Publish package
         run: python -m twine upload dist/*.whl

From 5b902a30f957141c3857efca41a40ddbcfc394c7 Mon Sep 17 00:00:00 2001
From: Zhaoyilunnn <dilved1014@gmail.com>
Date: Sat, 28 Oct 2023 15:59:57 +0000
Subject: [PATCH 3/6] ci: add unit test

---
 .github/workflows/unittest.yml | 21 +++++++++++++++++++++
 .github/workflows/wheels.yml   |  3 ---
 2 files changed, 21 insertions(+), 3 deletions(-)
 create mode 100644 .github/workflows/unittest.yml

diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
new file mode 100644
index 0000000..64b4a81
--- /dev/null
+++ b/.github/workflows/unittest.yml
@@ -0,0 +1,21 @@
+name: Unit test
+on:
+  push:
+  pull_request:
+    branches: ['main']
+jobs:
+  unnittest:
+    name: Unit test
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - uses: actions/setup-python@v4
+        name: Install Python
+        with:
+          python-version: '3.10'
+
+      - name: Install dependency
+        run: python -m pip install -r requirements.txt && python -m pip install pytest
+
+      - name: Run unit test
+        run: python -m pip install . && pytest tests/
diff --git a/.github/workflows/wheels.yml b/.github/workflows/wheels.yml
index 3186a1f..4a58b18 100644
--- a/.github/workflows/wheels.yml
+++ b/.github/workflows/wheels.yml
@@ -48,9 +48,6 @@ jobs:
       - name: Build wheels
         run: python -m cibuildwheel --output-dir dist
 
-      - name: Run unit tests
-        run: pip install dist/*.whl && pytest tests/
-
       - name: Publish package
         run: python -m twine upload dist/*.whl
         if: ${{ contains(github.ref, '/tags/') }}

From 77b4485f3b2756c86f032eceee7e9297867ae666 Mon Sep 17 00:00:00 2001
From: Zhaoyilunnn <dilved1014@gmail.com>
Date: Sun, 29 Oct 2023 08:26:04 +0000
Subject: [PATCH 4/6] feat: torch autograd with param shift demo

---
 quafu/algorithms/layers/qnode.py      | 22 +++++++-------
 quafu/algorithms/layers/torch.py      | 30 ++++++++++++-------
 tests/quafu/algorithms/layers_test.py | 42 ++++++++++++++++++++++-----
 3 files changed, 66 insertions(+), 28 deletions(-)

diff --git a/quafu/algorithms/layers/qnode.py b/quafu/algorithms/layers/qnode.py
index a1cdf61..60192af 100644
--- a/quafu/algorithms/layers/qnode.py
+++ b/quafu/algorithms/layers/qnode.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 
+from typing import List, Optional
 import numpy as np
 from quafu import QuantumCircuit
 from quafu.algorithms import Hamiltonian
@@ -30,15 +31,17 @@ def _generate_expval_z(num_qubits: int):
 
 
 # TODO(zhaoyilun): support more measurement types
-def execute(circ: QuantumCircuit):
-    """execute.
+def run_circ(circ: QuantumCircuit, params: Optional[List[float]] = None):
+    """Execute a circuit
 
     Args:
         circ (QuantumCircuit): circ
+        params (Optional[List[float]]): params
     """
     obs_list = _generate_expval_z(circ.num)
     estimator = Estimator(circ)
-    params = [g.paras for g in circ.parameterized_gates]
+    if params is None:
+        params = [g.paras for g in circ.parameterized_gates]
     output = [estimator.run(obs, params) for obs in obs_list]
     return np.array(output)
 
@@ -53,10 +56,10 @@ def jacobian(circ: QuantumCircuit, params_input: np.ndarray):
     """
     batch_size, num_params = params_input.shape
     obs_list = _generate_expval_z(circ.num)
-    num_expvals = len(obs_list)
+    num_outputs = len(obs_list)
     estimator = Estimator(circ)
     calc_grad = ParamShift(estimator)
-    output = np.zeros((batch_size, num_expvals, num_params))
+    output = np.zeros((batch_size, num_outputs, num_params))
     for i in range(batch_size):
         grad_list = [
             np.array(calc_grad(obs, params_input[i, :].tolist())) for obs in obs_list
@@ -69,16 +72,15 @@ def compute_vjp(jac: np.ndarray, dy: np.ndarray):
     """compute vector-jacobian product
 
     Args:
-        jac (np.ndarray): jac with shape (batch_size, num_expvals, num_params)
-        dy (np.ndarray): dy
+        jac (np.ndarray): jac with shape (batch_size, num_outputs, num_params)
+        dy (np.ndarray): dy with shape (batch_size, num_outputs)
     """
-    batch_size, num_expvals, num_params = jac.shape
-    assert dy.shape[0] == batch_size and dy.shape[1] == num_expvals
+    batch_size, num_outputs, num_params = jac.shape
+    assert dy.shape[0] == batch_size and dy.shape[1] == num_outputs
 
     vjp = np.zeros((batch_size, num_params))
 
     for i in range(batch_size):
         vjp[i] = dy[i, :].T @ jac[i, :, :]
 
-    vjp = vjp.sum(0)
     return vjp
diff --git a/quafu/algorithms/layers/torch.py b/quafu/algorithms/layers/torch.py
index c886de5..6ddf95b 100644
--- a/quafu/algorithms/layers/torch.py
+++ b/quafu/algorithms/layers/torch.py
@@ -14,28 +14,36 @@
 
 """quafu PyTorch quantum layer"""
 
-from typing import Any
 import torch
+import numpy as np
 from quafu import QuantumCircuit
-from quafu import QuantumCircuit, simulate
-from quafu.elements.element_gates.matrices import ZMatrix
+from quafu.algorithms.layers.qnode import compute_vjp, jacobian
 
 
 class ExecuteCircuits(torch.autograd.Function):
     """TODO(zhaoyilun): document"""
 
     @staticmethod
-    def forward(ctx, parameters, **kwargs) -> Any:
+    def forward(ctx, parameters, kwargs):
         ctx.run_fn = kwargs["run_fn"]
         ctx.circ = kwargs["circ"]
-        out = ctx.run_fn(ctx.circ, parameters)
-        return out
+        ctx.save_for_backward(parameters)
+        parameters = parameters.numpy().tolist()
+        outputs = []
+        for para in parameters:
+            out = ctx.run_fn(ctx.circ, para)
+            outputs.append(out)
+        outputs = np.stack(outputs)
+        outputs = torch.from_numpy(outputs)
+        return outputs
 
     @staticmethod
-    def backward(ctx: Any, grad_out) -> Any:
-        circ, grad_fn = ctx.saved_tensors
-        grad = grad_fn(circ, grad_out)
-        return grad
+    def backward(ctx, grad_out):
+        (parameters,) = ctx.saved_tensors
+        jac = jacobian(ctx.circ, parameters.numpy())
+        vjp = compute_vjp(jac, grad_out.numpy())
+        vjp = torch.from_numpy(vjp)
+        return vjp, None
 
 
 # TODO(zhaoyilun): doc
@@ -50,4 +58,4 @@ def execute(circ: QuantumCircuit, run_fn, grad_fn, parameters: torch.Tensor):
 
     kwargs = {"circ": circ, "run_fn": run_fn, "grad_fn": grad_fn}
 
-    return ExecuteCircuits.apply(parameters, **kwargs)
+    return ExecuteCircuits.apply(parameters, kwargs)
diff --git a/tests/quafu/algorithms/layers_test.py b/tests/quafu/algorithms/layers_test.py
index 6a9c5bf..6f0b9e9 100644
--- a/tests/quafu/algorithms/layers_test.py
+++ b/tests/quafu/algorithms/layers_test.py
@@ -13,22 +13,50 @@
 # limitations under the License.
 import numpy as np
 
+import torch.nn
 from quafu.algorithms.layers import jacobian, compute_vjp
+from quafu.algorithms.layers.torch import execute
+from quafu.algorithms.layers.qnode import run_circ
 from quafu.circuits.quantum_circuit import QuantumCircuit
 
 
-def test_compute_vjp():
+class MyModel(torch.nn.Module):
+    def __init__(self, circ: QuantumCircuit):
+        super().__init__()
+        self.circ = circ
+        self.linear = torch.nn.Linear(3, 3, dtype=torch.double)
+
+    def forward(self, features):
+        out = self.linear(features)
+        out = execute(self.circ, run_circ, None, out)
+        return out
+
+
+class TestLayers:
     circ = QuantumCircuit(2)
     circ.x(0)
     circ.rx(0, 0.1)
     circ.ry(1, 0.5)
     circ.ry(0, 0.1)
 
-    params_input = np.random.randn(4, 3)
-    jac = jacobian(circ, params_input)
+    def test_compute_vjp(self):
+        params_input = np.random.randn(4, 3)
+        jac = jacobian(self.circ, params_input)
+
+        dy = np.random.randn(4, 2)
+        vjp = compute_vjp(jac, dy)
 
-    dy = np.random.randn(4, 2)
-    vjp = compute_vjp(jac, dy)
+        assert len(vjp.shape) == 2
+        assert vjp.shape[0] == 4
 
-    assert len(vjp.shape) == 1
-    assert vjp.shape[0] == 3
+    def test_torch_layer(self):
+        batch_size = 1
+        model = MyModel(self.circ)
+        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()

From 9c9415387978faed6047b2e3680a8817a69eed53 Mon Sep 17 00:00:00 2001
From: Zhaoyilunnn <dilved1014@gmail.com>
Date: Sun, 29 Oct 2023 08:32:28 +0000
Subject: [PATCH 5/6] ci: fix unit test torch dependency

---
 .github/workflows/unittest.yml | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
index 64b4a81..e1f1761 100644
--- a/.github/workflows/unittest.yml
+++ b/.github/workflows/unittest.yml
@@ -17,5 +17,12 @@ jobs:
       - name: Install dependency
         run: python -m pip install -r requirements.txt && python -m pip install pytest
 
-      - name: Run unit test
-        run: python -m pip install . && pytest tests/
+        # TODO(zhaoyilun): Build seperate package for pyquafu-torch
+      - name: Install torch
+      - run: python -m pip install torch torchvision torchaudio
+
+      - name: Install pyquafu
+        run: python -m pip install
+
+      - name: Run unit tests
+      - run: pytest tests/

From c4049887dfc08707da6cfe5e0c65b23363da0c29 Mon Sep 17 00:00:00 2001
From: Zhaoyilunnn <dilved1014@gmail.com>
Date: Sun, 29 Oct 2023 08:35:32 +0000
Subject: [PATCH 6/6] ci: fix unit test typo

---
 .github/workflows/unittest.yml | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/unittest.yml b/.github/workflows/unittest.yml
index e1f1761..625872a 100644
--- a/.github/workflows/unittest.yml
+++ b/.github/workflows/unittest.yml
@@ -17,12 +17,12 @@ jobs:
       - name: Install dependency
         run: python -m pip install -r requirements.txt && python -m pip install pytest
 
-        # TODO(zhaoyilun): Build seperate package for pyquafu-torch
+      # TODO(zhaoyilun): Build seperate package for pyquafu-torch
       - name: Install torch
-      - run: python -m pip install torch torchvision torchaudio
+        run: python -m pip install torch torchvision torchaudio
 
       - name: Install pyquafu
-        run: python -m pip install
+        run: python -m pip install .
 
       - name: Run unit tests
-      - run: pytest tests/
+        run: pytest tests/