The adjoint method in lightning supports qubit unitaries (#540)

* Initial commit to add adjoint diff support for circuits with qubit-unitaries. * Auto update version * Try nick-fields/retry for clang-tidy check. * Fix LQubit c++ tests. * Generalize unitary addiff tests to L-Kokkos. * Revert condition for branching in applyOperation. * Fix pip install and format * remove prep argument * Auto update version * Trigger CI * Generalize fix and tests to L-GPU. * Set timeout_minutes for retry. * Add L-GPU-MPI [skip ci]. * trigger CI * Auto update version * Apply prep fix to mpitests. * Add C++ coverage for QubitUnitary in addiff. * Add test_qubit_unitary docstring [skip ci]. * Update changelog [skip ci] * Add diff of unitary tests. * Try removing gcc-13. * Remove annoying headers. * Revert changes to format.yml. * Update pennylane_lightning/core/src/simulators/lightning_gpu/StateVectorCudaManaged.hpp [skip ci] Co-authored-by: Amintor Dusko <[email protected]> * Update changelog [skip ci] * Remove unused tol variable [skip ci] * Revert changes to the opsdata init in adjoint tests. * Auto update version * nuni => n_targets param is less confusing. * Revert changes to std::complex * Increase tests_windows timeouts. * Revert to 30 min. * Try removing {} init. * if-no-files-found: error in uploads. * Fix coverage.xml name * Auto update version * Fix parameter counting for circuits with QubitUnitaries. * Auto update version * trigger ci * Update .github/CHANGELOG.md [skip ci] Co-authored-by: Amintor Dusko <[email protected]> * Use name instead of class to serialize QubitUnitary. * Update CHANGELOG.md [skip ci] * Update CHANGELOG.md [skip ci] * Add float32 mpitests * Rename fixture function. * trigger ci --------- Co-authored-by: Dev version update bot <github-actions[bot]@users.noreply.github.com> Co-authored-by: albi3ro <[email protected]> Co-authored-by: AmintorDusko <[email protected]> Co-authored-by: Amintor Dusko <[email protected]>
PennyLaneAI · Nov 15, 2023 · bf6e121 · bf6e121
1 parent 7442701
commit bf6e121
Show file tree

Hide file tree

Showing 20 changed files with 347 additions and 38 deletions.
diff --git a/.github/CHANGELOG.md b/.github/CHANGELOG.md
@@ -2,10 +2,19 @@
 
 ### New features since last release
 
+* `qml.QubitUnitary` operators can be included in a circuit differentiated with the adjoint method. Lightning handles circuits with arbitrary non-differentiable `qml.QubitUnitary` operators. 1,2-qubit `qml.QubitUnitary` operators with differentiable parameters can be differentiated using decomposition.
+  [(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)
+
 ### Breaking changes
 
+* Overload `applyOperation` with a fifth `matrix` argument to all state vector classes to support arbitrary operations in `AdjointJacobianBase`.
+  [(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)
+
 ### Improvements
 
+* Modify `setup.py` to use backend-specific build directory (`f"build_{backend}"`) to accelerate rebuilding backends in alternance.
+  [(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)
+
 * Update Dockerfile and rewrite the `build-wheel-lightning-gpu` stage to build Lightning-GPU from the `pennylane-lightning` monorepo.
   [(#539)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/539)
 
@@ -25,6 +34,9 @@
 
 ### Bug fixes
 
+* Move deprecated `stateprep` `QuantumScript` argument into the operation list in `mpitests/test_adjoint_jacobian.py`.
+  [(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)
+
 * Fix MPI Python unit tests for the adjoint method.
   [(#538)](https://github.com/PennyLaneAI/pennylane-lightning/pull/538)
 

diff --git a/.github/workflows/tests_gpu_cu11.yml b/.github/workflows/tests_gpu_cu11.yml
@@ -319,7 +319,8 @@ jobs:
         with:
           name: ubuntu-codecov-results-python
           path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
-
+          if-no-files-found: error
+
   upload-to-codecov-linux-python:
     needs: [pythontestswithLGPU]
     name: Upload coverage data to codecov

diff --git a/.github/workflows/tests_gpu_kokkos.yml b/.github/workflows/tests_gpu_kokkos.yml
@@ -313,7 +313,7 @@ jobs:
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append 2> /dev/null || echo Something went wrong with pl-device-test shot=20000
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append 2> /dev/null || echo Something went wrong with pl-device-test shot=None
           PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS 2> /dev/null || echo Something went wrong with Pytest
-          mv coverage.xml coverage-${{ github.job }}.xml
+          mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
 
       - name: Install all backend devices
         if: ${{ matrix.pl_backend == 'all' }}
@@ -345,3 +345,4 @@ jobs:
         with:
           name: ubuntu-codecov-results-python
           path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
+          if-no-files-found: error
diff --git a/.github/workflows/tests_linux.yml b/.github/workflows/tests_linux.yml
@@ -174,6 +174,7 @@ jobs:
         with:
           name: ubuntu-codecov-results-python
           path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
+          if-no-files-found: error
 
   cpptestswithOpenBLAS:
     if: ${{ !contains(fromJSON('["workflow_call"]'), github.event_name) }}
@@ -323,6 +324,7 @@ jobs:
         with:
           name: ubuntu-codecov-results-python
           path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
+          if-no-files-found: error
 
   build_and_cache_Kokkos:
     name: "Build and cache Kokkos"
@@ -512,7 +514,7 @@ jobs:
           PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
           pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
           pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
-          mv coverage.xml coverage-${{ github.job }}.xml
+          mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
 
       - name: Install all backend devices
         if: ${{ matrix.pl_backend == 'all' }}
@@ -541,6 +543,7 @@ jobs:
         with:
           name: ubuntu-codecov-results-python
           path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
+          if-no-files-found: error
 
   upload-to-codecov-linux-python:
     needs: [pythontests, pythontestswithOpenBLAS, pythontestswithKokkos]

diff --git a/.github/workflows/tests_linux_x86_mpi_gpu.yml b/.github/workflows/tests_linux_x86_mpi_gpu.yml
@@ -138,16 +138,16 @@ jobs:
         uses: actions/upload-artifact@v3
         if: always()
         with:
-          if-no-files-found: error
           name: ubuntu-tests-reports
           path: ./Build/tests/results/
+          if-no-files-found: error
 
       - name: Upload code coverage results
         uses: actions/upload-artifact@v3
         with:
-          if-no-files-found: error
           name: ubuntu-codecov-results-cpp
           path: ./Build/coverage-${{ github.job }}-lightning_gpu_${{ matrix.mpilib }}.info
+          if-no-files-found: error
 
       - name: Cleanup
         if: always()

diff --git a/.github/workflows/tests_windows.yml b/.github/workflows/tests_windows.yml
@@ -61,13 +61,14 @@ jobs:
         with:
           name: windows-test-report-${{ github.job }}-${{ matrix.pl_backend }}
           path: .\Build\tests\results\
-
+          if-no-files-found: error
+
       - name: Upload coverage results
         uses: actions/upload-artifact@v3
         with:
           name: windows-coverage-report
           path: .\coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
-
+          if-no-files-found: error
 
   win-set-matrix-x86:
     name: Set builder matrix
@@ -218,12 +219,14 @@ jobs:
         with:
           name: windows-test-report-${{ github.job }}-${{ matrix.pl_backend }}
           path: .\Build\tests\results\
+          if-no-files-found: error
 
       - name: Upload coverage results
         uses: actions/upload-artifact@v3
         with:
           name: windows-coverage-report
           path: .\coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
+          if-no-files-found: error
 
   upload-to-codecov-windows:
     needs: [cpptests, cpptestswithkokkos]

diff --git a/mpitests/test_adjoint_jacobian.py b/mpitests/test_adjoint_jacobian.py
@@ -42,6 +42,18 @@
 )
 
 
+@pytest.fixture(name="dev", params=fixture_params)
+def fixture_dev(request):
+    """Returns a PennyLane device."""
+    return qml.device(
+        device_name,
+        wires=8,
+        mpi=True,
+        c_dtype=request.param[0],
+        batch_obs=request.param[1],
+    )
+
+
 def Rx(theta):
     r"""One-qubit rotation about the x axis.
 
@@ -78,17 +90,6 @@ def Rz(theta):
 class TestAdjointJacobian:  # pylint: disable=too-many-public-methods
     """Tests for the adjoint_jacobian method"""
 
-    @pytest.fixture(params=fixture_params)
-    def dev(self, request):
-        """Returns a PennyLane device."""
-        return qml.device(
-            device_name,
-            wires=8,
-            mpi=True,
-            c_dtype=request.param[0],
-            batch_obs=request.param[1],
-        )
-
     def test_not_expval(self, dev):
         """Test if a QuantumFunctionError is raised for a tape with measurements that are not
         expectation values"""
@@ -189,7 +190,7 @@ def test_pauli_rotation_gradient(self, stateprep, G, theta, dev):
         )
 
         tape = qml.tape.QuantumScript(
-            [G(theta, 0)], [qml.expval(qml.PauliZ(0))], [stateprep(random_state, 0)]
+            [stateprep(random_state, 0), G(theta, 0)], [qml.expval(qml.PauliZ(0))]
         )
 
         tape.trainable_params = {1}
@@ -1381,3 +1382,106 @@ def circuit(params):
     comm.Barrier()
 
     assert np.allclose(j_cpu, j_gpu)
+
+
+@pytest.mark.parametrize("n_targets", range(1, 5))
+def test_qubit_unitary(dev, n_targets):
+    """Tests that ``qml.QubitUnitary`` can be included in circuits differentiated with the adjoint method."""
+    n_wires = len(dev.wires)
+    dev_def = qml.device("default.qubit.legacy", wires=n_wires)
+    h = 1e-3 if dev.R_DTYPE == np.float32 else 1e-7
+    c_dtype = np.complex64 if dev.R_DTYPE == np.float32 else np.complex128
+
+    np.random.seed(1337)
+    par = 2 * np.pi * np.random.rand(n_wires)
+    U = np.random.rand(2**n_targets, 2**n_targets) + 1j * np.random.rand(
+        2**n_targets, 2**n_targets
+    )
+    U, _ = np.linalg.qr(U)
+    init_state = np.random.rand(2**n_wires) + 1j * np.random.rand(2**n_wires)
+    init_state /= np.sqrt(np.dot(np.conj(init_state), init_state))
+
+    comm = MPI.COMM_WORLD
+    par = comm.bcast(par, root=0)
+    U = comm.bcast(U, root=0)
+    init_state = comm.bcast(init_state, root=0)
+
+    init_state = np.array(init_state, requires_grad=False, dtype=c_dtype)
+    U = np.array(U, requires_grad=False, dtype=c_dtype)
+    obs = qml.operation.Tensor(*(qml.PauliZ(i) for i in range(n_wires)))
+
+    def circuit(x):
+        qml.StatePrep(init_state, wires=range(n_wires))
+        for i in range(n_wires // 2):
+            qml.RY(x[i], wires=i)
+        qml.QubitUnitary(U, wires=range(n_targets))
+        for i in range(n_wires // 2, n_wires):
+            qml.RY(x[i], wires=i)
+        return qml.expval(obs)
+
+    circ = qml.QNode(circuit, dev, diff_method="adjoint")
+    circ_ps = qml.QNode(circuit, dev, diff_method="parameter-shift")
+    circ_def = qml.QNode(circuit, dev_def, diff_method="adjoint")
+    jac = qml.jacobian(circ)(par)
+    jac_ps = qml.jacobian(circ_ps)(par)
+    jac_def = qml.jacobian(circ_def)(par)
+
+    comm.Barrier()
+
+    assert len(jac) == n_wires
+    assert not np.allclose(jac, 0.0)
+    assert np.allclose(jac, jac_ps, atol=h, rtol=0)
+    assert np.allclose(jac, jac_def, atol=h, rtol=0)
+
+
+@pytest.mark.parametrize("n_targets", [1, 2])
+def test_diff_qubit_unitary(dev, n_targets):
+    """Tests that ``qml.QubitUnitary`` can be differentiated with the adjoint method."""
+    n_wires = len(dev.wires)
+    dev_def = qml.device("default.qubit", wires=n_wires)
+    h = 1e-3 if dev.R_DTYPE == np.float32 else 1e-7
+    c_dtype = np.complex64 if dev.R_DTYPE == np.float32 else np.complex128
+
+    np.random.seed(1337)
+    par = 2 * np.pi * np.random.rand(n_wires)
+    U = np.random.rand(2**n_targets, 2**n_targets) + 1j * np.random.rand(
+        2**n_targets, 2**n_targets
+    )
+    U, _ = np.linalg.qr(U)
+    init_state = np.random.rand(2**n_wires) + 1j * np.random.rand(2**n_wires)
+    init_state /= np.sqrt(np.dot(np.conj(init_state), init_state))
+
+    comm = MPI.COMM_WORLD
+    par = comm.bcast(par, root=0)
+    U = comm.bcast(U, root=0)
+    init_state = comm.bcast(init_state, root=0)
+
+    init_state = np.array(init_state, requires_grad=False, dtype=c_dtype)
+    U = np.array(U, requires_grad=False, dtype=c_dtype)
+    obs = qml.operation.Tensor(*(qml.PauliZ(i) for i in range(n_wires)))
+
+    def circuit(x, u_mat):
+        qml.StatePrep(init_state, wires=range(n_wires))
+        for i in range(n_wires // 2):
+            qml.RY(x[i], wires=i)
+        qml.QubitUnitary(u_mat, wires=range(n_targets))
+        for i in range(n_wires // 2, n_wires):
+            qml.RY(x[i], wires=i)
+        return qml.expval(obs)
+
+    circ = qml.QNode(circuit, dev, diff_method="adjoint")
+    circ_def = qml.QNode(circuit, dev_def, diff_method="adjoint")
+    circ_fd = qml.QNode(circuit, dev, diff_method="finite-diff", h=h)
+    circ_ps = qml.QNode(circuit, dev, diff_method="parameter-shift")
+    jacs = qml.jacobian(circ)(par, U)
+    jacs_def = qml.jacobian(circ_def)(par, U)
+    jacs_fd = qml.jacobian(circ_fd)(par, U)
+    jacs_ps = qml.jacobian(circ_ps)(par, U)
+
+    comm.Barrier()
+
+    for jac, jac_def, jac_fd, jac_ps in zip(jacs, jacs_def, jacs_fd, jacs_ps):
+        assert not np.allclose(jac, 0.0)
+        assert np.allclose(jac, jac_fd, atol=h, rtol=0)
+        assert np.allclose(jac, jac_ps, atol=h, rtol=0)
+        assert np.allclose(jac, jac_def, atol=h, rtol=0)
diff --git a/pennylane_lightning/core/_serialize.py b/pennylane_lightning/core/_serialize.py
@@ -307,11 +307,14 @@ def serialize_ops(
             for single_op in op_list:
                 name = single_op.name
                 names.append(name)
-
-                if not hasattr(self.sv_type, name):
+                # QubitUnitary is a special case, it has a parameter which is not differentiable.
+                # We thus pass a dummy 0.0 parameter which will not be referenced
+                if name == "QubitUnitary":
+                    params.append([0.0])
+                    mats.append(matrix(single_op))
+                elif not hasattr(self.sv_type, name):
                     params.append([])
                     mats.append(matrix(single_op))
-
                 else:
                     params.append(single_op.parameters)
                     mats.append([])

diff --git a/pennylane_lightning/core/_version.py b/pennylane_lightning/core/_version.py
@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.34.0-dev6"
+__version__ = "0.34.0-dev7"
diff --git a/pennylane_lightning/core/src/algorithms/AdjointJacobianBase.hpp b/pennylane_lightning/core/src/algorithms/AdjointJacobianBase.hpp
@@ -59,7 +59,8 @@ template <class StateVectorT, class Derived> class AdjointJacobianBase {
             state.applyOperation(operations.getOpsName()[op_idx],
                                  operations.getOpsWires()[op_idx],
                                  operations.getOpsInverses()[op_idx] ^ adj,
-                                 operations.getOpsParams()[op_idx]);
+                                 operations.getOpsParams()[op_idx],
+                                 operations.getOpsMatrices()[op_idx]);
         }
     }
 
@@ -79,7 +80,8 @@ template <class StateVectorT, class Derived> class AdjointJacobianBase {
         state.applyOperation(operations.getOpsName()[op_idx],
                              operations.getOpsWires()[op_idx],
                              !operations.getOpsInverses()[op_idx],
-                             operations.getOpsParams()[op_idx]);
+                             operations.getOpsParams()[op_idx],
+                             operations.getOpsMatrices()[op_idx]);
     }
 
     /**

diff --git a/pennylane_lightning/core/src/algorithms/tests/Test_AdjointJacobian.cpp b/pennylane_lightning/core/src/algorithms/tests/Test_AdjointJacobian.cpp
@@ -351,8 +351,10 @@ template <typename TypeList> void testAdjointJacobian() {
                     "PauliX", std::vector<size_t>{1}),
                 std::make_shared<NamedObs<StateVectorT>>(
                     "PauliX", std::vector<size_t>{2}));
+            std::vector<ComplexT> cnot{1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0,
+                                       0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0};
             auto ops = OpsData<StateVectorT>(
-                {"RZ", "RY", "RZ", "CNOT", "CNOT", "RZ", "RY", "RZ"},
+                {"RZ", "RY", "RZ", "QubitUnitary", "CNOT", "RZ", "RY", "RZ"},
                 {{param[0]},
                  {param[1]},
                  {param[2]},
@@ -362,7 +364,9 @@ template <typename TypeList> void testAdjointJacobian() {
                  {param[1]},
                  {param[2]}},
                 {{0}, {0}, {0}, {0, 1}, {1, 2}, {1}, {1}, {1}},
-                {false, false, false, false, false, false, false, false});
+                {false, false, false, false, false, false, false, false},
+                std::vector<std::vector<ComplexT>>{
+                    {}, {}, {}, cnot, {}, {}, {}, {}});
 
             JacobianData<StateVectorT> tape{
                 num_params, psi.getLength(), psi.getData(), {obs}, ops, tp};

diff --git a/pennylane_lightning/core/src/algorithms/tests/mpi/Test_AdjointJacobianMPI.cpp b/pennylane_lightning/core/src/algorithms/tests/mpi/Test_AdjointJacobianMPI.cpp
@@ -249,8 +249,10 @@ template <typename TypeList> void testAdjointJacobian() {
                     "PauliX", std::vector<size_t>{1}),
                 std::make_shared<NamedObsMPI<StateVectorT>>(
                     "PauliX", std::vector<size_t>{2}));
+            std::vector<ComplexT> cnot{1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0,
+                                       0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0};
             auto ops = OpsData<StateVectorT>(
-                {"RZ", "RY", "RZ", "CNOT", "CNOT", "RZ", "RY", "RZ"},
+                {"RZ", "RY", "RZ", "QubitUnitary", "CNOT", "RZ", "RY", "RZ"},
                 {{param[0]},
                  {param[1]},
                  {param[2]},
@@ -260,7 +262,9 @@ template <typename TypeList> void testAdjointJacobian() {
                  {param[1]},
                  {param[2]}},
                 {{0}, {0}, {0}, {0, 1}, {1, 2}, {1}, {1}, {1}},
-                {false, false, false, false, false, false, false, false});
+                {false, false, false, false, false, false, false, false},
+                std::vector<std::vector<ComplexT>>{
+                    {}, {}, {}, cnot, {}, {}, {}, {}});
 
             JacobianDataMPI<StateVectorT> tape{num_params, psi, {obs}, ops, tp};