Skip to content

Commit bf6e121

Browse files
vincentmrgithub-actions[bot]albi3roAmintorDusko
authored
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]>
1 parent 7442701 commit bf6e121

20 files changed

+347
-38
lines changed

‎.github/CHANGELOG.md

+12
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,19 @@
22

33
### New features since last release
44

5+
* `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.
6+
[(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)
7+
58
### Breaking changes
69

10+
* Overload `applyOperation` with a fifth `matrix` argument to all state vector classes to support arbitrary operations in `AdjointJacobianBase`.
11+
[(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)
12+
713
### Improvements
814

15+
* Modify `setup.py` to use backend-specific build directory (`f"build_{backend}"`) to accelerate rebuilding backends in alternance.
16+
[(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)
17+
918
* Update Dockerfile and rewrite the `build-wheel-lightning-gpu` stage to build Lightning-GPU from the `pennylane-lightning` monorepo.
1019
[(#539)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/539)
1120

@@ -25,6 +34,9 @@
2534

2635
### Bug fixes
2736

37+
* Move deprecated `stateprep` `QuantumScript` argument into the operation list in `mpitests/test_adjoint_jacobian.py`.
38+
[(#540)] (https://github.com/PennyLaneAI/pennylane-lightning/pull/540)
39+
2840
* Fix MPI Python unit tests for the adjoint method.
2941
[(#538)](https://github.com/PennyLaneAI/pennylane-lightning/pull/538)
3042

‎.github/workflows/tests_gpu_cu11.yml

+2-1
Original file line numberDiff line numberDiff line change
@@ -319,7 +319,8 @@ jobs:
319319
with:
320320
name: ubuntu-codecov-results-python
321321
path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
322-
322+
if-no-files-found: error
323+
323324
upload-to-codecov-linux-python:
324325
needs: [pythontestswithLGPU]
325326
name: Upload coverage data to codecov

‎.github/workflows/tests_gpu_kokkos.yml

+2-1
Original file line numberDiff line numberDiff line change
@@ -313,7 +313,7 @@ jobs:
313313
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
314314
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
315315
PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS 2> /dev/null || echo Something went wrong with Pytest
316-
mv coverage.xml coverage-${{ github.job }}.xml
316+
mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
317317
318318
- name: Install all backend devices
319319
if: ${{ matrix.pl_backend == 'all' }}
@@ -345,3 +345,4 @@ jobs:
345345
with:
346346
name: ubuntu-codecov-results-python
347347
path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
348+
if-no-files-found: error

‎.github/workflows/tests_linux.yml

+4-1
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,7 @@ jobs:
174174
with:
175175
name: ubuntu-codecov-results-python
176176
path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
177+
if-no-files-found: error
177178

178179
cpptestswithOpenBLAS:
179180
if: ${{ !contains(fromJSON('["workflow_call"]'), github.event_name) }}
@@ -323,6 +324,7 @@ jobs:
323324
with:
324325
name: ubuntu-codecov-results-python
325326
path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
327+
if-no-files-found: error
326328

327329
build_and_cache_Kokkos:
328330
name: "Build and cache Kokkos"
@@ -512,7 +514,7 @@ jobs:
512514
PL_DEVICE=${DEVICENAME} python -m pytest tests/ $COVERAGE_FLAGS
513515
pl-device-test --device ${DEVICENAME} --skip-ops --shots=20000 $COVERAGE_FLAGS --cov-append
514516
pl-device-test --device ${DEVICENAME} --shots=None --skip-ops $COVERAGE_FLAGS --cov-append
515-
mv coverage.xml coverage-${{ github.job }}.xml
517+
mv coverage.xml coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
516518
517519
- name: Install all backend devices
518520
if: ${{ matrix.pl_backend == 'all' }}
@@ -541,6 +543,7 @@ jobs:
541543
with:
542544
name: ubuntu-codecov-results-python
543545
path: ./main/coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
546+
if-no-files-found: error
544547

545548
upload-to-codecov-linux-python:
546549
needs: [pythontests, pythontestswithOpenBLAS, pythontestswithKokkos]

‎.github/workflows/tests_linux_x86_mpi_gpu.yml

+2-2
Original file line numberDiff line numberDiff line change
@@ -138,16 +138,16 @@ jobs:
138138
uses: actions/upload-artifact@v3
139139
if: always()
140140
with:
141-
if-no-files-found: error
142141
name: ubuntu-tests-reports
143142
path: ./Build/tests/results/
143+
if-no-files-found: error
144144

145145
- name: Upload code coverage results
146146
uses: actions/upload-artifact@v3
147147
with:
148-
if-no-files-found: error
149148
name: ubuntu-codecov-results-cpp
150149
path: ./Build/coverage-${{ github.job }}-lightning_gpu_${{ matrix.mpilib }}.info
150+
if-no-files-found: error
151151

152152
- name: Cleanup
153153
if: always()

‎.github/workflows/tests_windows.yml

+5-2
Original file line numberDiff line numberDiff line change
@@ -61,13 +61,14 @@ jobs:
6161
with:
6262
name: windows-test-report-${{ github.job }}-${{ matrix.pl_backend }}
6363
path: .\Build\tests\results\
64-
64+
if-no-files-found: error
65+
6566
- name: Upload coverage results
6667
uses: actions/upload-artifact@v3
6768
with:
6869
name: windows-coverage-report
6970
path: .\coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
70-
71+
if-no-files-found: error
7172

7273
win-set-matrix-x86:
7374
name: Set builder matrix
@@ -218,12 +219,14 @@ jobs:
218219
with:
219220
name: windows-test-report-${{ github.job }}-${{ matrix.pl_backend }}
220221
path: .\Build\tests\results\
222+
if-no-files-found: error
221223

222224
- name: Upload coverage results
223225
uses: actions/upload-artifact@v3
224226
with:
225227
name: windows-coverage-report
226228
path: .\coverage-${{ github.job }}-${{ matrix.pl_backend }}.xml
229+
if-no-files-found: error
227230

228231
upload-to-codecov-windows:
229232
needs: [cpptests, cpptestswithkokkos]

‎mpitests/test_adjoint_jacobian.py

+116-12
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,18 @@
4242
)
4343

4444

45+
@pytest.fixture(name="dev", params=fixture_params)
46+
def fixture_dev(request):
47+
"""Returns a PennyLane device."""
48+
return qml.device(
49+
device_name,
50+
wires=8,
51+
mpi=True,
52+
c_dtype=request.param[0],
53+
batch_obs=request.param[1],
54+
)
55+
56+
4557
def Rx(theta):
4658
r"""One-qubit rotation about the x axis.
4759
@@ -78,17 +90,6 @@ def Rz(theta):
7890
class TestAdjointJacobian: # pylint: disable=too-many-public-methods
7991
"""Tests for the adjoint_jacobian method"""
8092

81-
@pytest.fixture(params=fixture_params)
82-
def dev(self, request):
83-
"""Returns a PennyLane device."""
84-
return qml.device(
85-
device_name,
86-
wires=8,
87-
mpi=True,
88-
c_dtype=request.param[0],
89-
batch_obs=request.param[1],
90-
)
91-
9293
def test_not_expval(self, dev):
9394
"""Test if a QuantumFunctionError is raised for a tape with measurements that are not
9495
expectation values"""
@@ -189,7 +190,7 @@ def test_pauli_rotation_gradient(self, stateprep, G, theta, dev):
189190
)
190191

191192
tape = qml.tape.QuantumScript(
192-
[G(theta, 0)], [qml.expval(qml.PauliZ(0))], [stateprep(random_state, 0)]
193+
[stateprep(random_state, 0), G(theta, 0)], [qml.expval(qml.PauliZ(0))]
193194
)
194195

195196
tape.trainable_params = {1}
@@ -1381,3 +1382,106 @@ def circuit(params):
13811382
comm.Barrier()
13821383

13831384
assert np.allclose(j_cpu, j_gpu)
1385+
1386+
1387+
@pytest.mark.parametrize("n_targets", range(1, 5))
1388+
def test_qubit_unitary(dev, n_targets):
1389+
"""Tests that ``qml.QubitUnitary`` can be included in circuits differentiated with the adjoint method."""
1390+
n_wires = len(dev.wires)
1391+
dev_def = qml.device("default.qubit.legacy", wires=n_wires)
1392+
h = 1e-3 if dev.R_DTYPE == np.float32 else 1e-7
1393+
c_dtype = np.complex64 if dev.R_DTYPE == np.float32 else np.complex128
1394+
1395+
np.random.seed(1337)
1396+
par = 2 * np.pi * np.random.rand(n_wires)
1397+
U = np.random.rand(2**n_targets, 2**n_targets) + 1j * np.random.rand(
1398+
2**n_targets, 2**n_targets
1399+
)
1400+
U, _ = np.linalg.qr(U)
1401+
init_state = np.random.rand(2**n_wires) + 1j * np.random.rand(2**n_wires)
1402+
init_state /= np.sqrt(np.dot(np.conj(init_state), init_state))
1403+
1404+
comm = MPI.COMM_WORLD
1405+
par = comm.bcast(par, root=0)
1406+
U = comm.bcast(U, root=0)
1407+
init_state = comm.bcast(init_state, root=0)
1408+
1409+
init_state = np.array(init_state, requires_grad=False, dtype=c_dtype)
1410+
U = np.array(U, requires_grad=False, dtype=c_dtype)
1411+
obs = qml.operation.Tensor(*(qml.PauliZ(i) for i in range(n_wires)))
1412+
1413+
def circuit(x):
1414+
qml.StatePrep(init_state, wires=range(n_wires))
1415+
for i in range(n_wires // 2):
1416+
qml.RY(x[i], wires=i)
1417+
qml.QubitUnitary(U, wires=range(n_targets))
1418+
for i in range(n_wires // 2, n_wires):
1419+
qml.RY(x[i], wires=i)
1420+
return qml.expval(obs)
1421+
1422+
circ = qml.QNode(circuit, dev, diff_method="adjoint")
1423+
circ_ps = qml.QNode(circuit, dev, diff_method="parameter-shift")
1424+
circ_def = qml.QNode(circuit, dev_def, diff_method="adjoint")
1425+
jac = qml.jacobian(circ)(par)
1426+
jac_ps = qml.jacobian(circ_ps)(par)
1427+
jac_def = qml.jacobian(circ_def)(par)
1428+
1429+
comm.Barrier()
1430+
1431+
assert len(jac) == n_wires
1432+
assert not np.allclose(jac, 0.0)
1433+
assert np.allclose(jac, jac_ps, atol=h, rtol=0)
1434+
assert np.allclose(jac, jac_def, atol=h, rtol=0)
1435+
1436+
1437+
@pytest.mark.parametrize("n_targets", [1, 2])
1438+
def test_diff_qubit_unitary(dev, n_targets):
1439+
"""Tests that ``qml.QubitUnitary`` can be differentiated with the adjoint method."""
1440+
n_wires = len(dev.wires)
1441+
dev_def = qml.device("default.qubit", wires=n_wires)
1442+
h = 1e-3 if dev.R_DTYPE == np.float32 else 1e-7
1443+
c_dtype = np.complex64 if dev.R_DTYPE == np.float32 else np.complex128
1444+
1445+
np.random.seed(1337)
1446+
par = 2 * np.pi * np.random.rand(n_wires)
1447+
U = np.random.rand(2**n_targets, 2**n_targets) + 1j * np.random.rand(
1448+
2**n_targets, 2**n_targets
1449+
)
1450+
U, _ = np.linalg.qr(U)
1451+
init_state = np.random.rand(2**n_wires) + 1j * np.random.rand(2**n_wires)
1452+
init_state /= np.sqrt(np.dot(np.conj(init_state), init_state))
1453+
1454+
comm = MPI.COMM_WORLD
1455+
par = comm.bcast(par, root=0)
1456+
U = comm.bcast(U, root=0)
1457+
init_state = comm.bcast(init_state, root=0)
1458+
1459+
init_state = np.array(init_state, requires_grad=False, dtype=c_dtype)
1460+
U = np.array(U, requires_grad=False, dtype=c_dtype)
1461+
obs = qml.operation.Tensor(*(qml.PauliZ(i) for i in range(n_wires)))
1462+
1463+
def circuit(x, u_mat):
1464+
qml.StatePrep(init_state, wires=range(n_wires))
1465+
for i in range(n_wires // 2):
1466+
qml.RY(x[i], wires=i)
1467+
qml.QubitUnitary(u_mat, wires=range(n_targets))
1468+
for i in range(n_wires // 2, n_wires):
1469+
qml.RY(x[i], wires=i)
1470+
return qml.expval(obs)
1471+
1472+
circ = qml.QNode(circuit, dev, diff_method="adjoint")
1473+
circ_def = qml.QNode(circuit, dev_def, diff_method="adjoint")
1474+
circ_fd = qml.QNode(circuit, dev, diff_method="finite-diff", h=h)
1475+
circ_ps = qml.QNode(circuit, dev, diff_method="parameter-shift")
1476+
jacs = qml.jacobian(circ)(par, U)
1477+
jacs_def = qml.jacobian(circ_def)(par, U)
1478+
jacs_fd = qml.jacobian(circ_fd)(par, U)
1479+
jacs_ps = qml.jacobian(circ_ps)(par, U)
1480+
1481+
comm.Barrier()
1482+
1483+
for jac, jac_def, jac_fd, jac_ps in zip(jacs, jacs_def, jacs_fd, jacs_ps):
1484+
assert not np.allclose(jac, 0.0)
1485+
assert np.allclose(jac, jac_fd, atol=h, rtol=0)
1486+
assert np.allclose(jac, jac_ps, atol=h, rtol=0)
1487+
assert np.allclose(jac, jac_def, atol=h, rtol=0)

‎pennylane_lightning/core/_serialize.py

+6-3
Original file line numberDiff line numberDiff line change
@@ -307,11 +307,14 @@ def serialize_ops(
307307
for single_op in op_list:
308308
name = single_op.name
309309
names.append(name)
310-
311-
if not hasattr(self.sv_type, name):
310+
# QubitUnitary is a special case, it has a parameter which is not differentiable.
311+
# We thus pass a dummy 0.0 parameter which will not be referenced
312+
if name == "QubitUnitary":
313+
params.append([0.0])
314+
mats.append(matrix(single_op))
315+
elif not hasattr(self.sv_type, name):
312316
params.append([])
313317
mats.append(matrix(single_op))
314-
315318
else:
316319
params.append(single_op.parameters)
317320
mats.append([])

‎pennylane_lightning/core/_version.py

+1-1
Original file line numberDiff line numberDiff line change
@@ -16,4 +16,4 @@
1616
Version number (major.minor.patch[-label])
1717
"""
1818

19-
__version__ = "0.34.0-dev6"
19+
__version__ = "0.34.0-dev7"

‎pennylane_lightning/core/src/algorithms/AdjointJacobianBase.hpp

+4-2
Original file line numberDiff line numberDiff line change
@@ -59,7 +59,8 @@ template <class StateVectorT, class Derived> class AdjointJacobianBase {
5959
state.applyOperation(operations.getOpsName()[op_idx],
6060
operations.getOpsWires()[op_idx],
6161
operations.getOpsInverses()[op_idx] ^ adj,
62-
operations.getOpsParams()[op_idx]);
62+
operations.getOpsParams()[op_idx],
63+
operations.getOpsMatrices()[op_idx]);
6364
}
6465
}
6566

@@ -79,7 +80,8 @@ template <class StateVectorT, class Derived> class AdjointJacobianBase {
7980
state.applyOperation(operations.getOpsName()[op_idx],
8081
operations.getOpsWires()[op_idx],
8182
!operations.getOpsInverses()[op_idx],
82-
operations.getOpsParams()[op_idx]);
83+
operations.getOpsParams()[op_idx],
84+
operations.getOpsMatrices()[op_idx]);
8385
}
8486

8587
/**

‎pennylane_lightning/core/src/algorithms/tests/Test_AdjointJacobian.cpp

+6-2
Original file line numberDiff line numberDiff line change
@@ -351,8 +351,10 @@ template <typename TypeList> void testAdjointJacobian() {
351351
"PauliX", std::vector<size_t>{1}),
352352
std::make_shared<NamedObs<StateVectorT>>(
353353
"PauliX", std::vector<size_t>{2}));
354+
std::vector<ComplexT> cnot{1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0,
355+
0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0};
354356
auto ops = OpsData<StateVectorT>(
355-
{"RZ", "RY", "RZ", "CNOT", "CNOT", "RZ", "RY", "RZ"},
357+
{"RZ", "RY", "RZ", "QubitUnitary", "CNOT", "RZ", "RY", "RZ"},
356358
{{param[0]},
357359
{param[1]},
358360
{param[2]},
@@ -362,7 +364,9 @@ template <typename TypeList> void testAdjointJacobian() {
362364
{param[1]},
363365
{param[2]}},
364366
{{0}, {0}, {0}, {0, 1}, {1, 2}, {1}, {1}, {1}},
365-
{false, false, false, false, false, false, false, false});
367+
{false, false, false, false, false, false, false, false},
368+
std::vector<std::vector<ComplexT>>{
369+
{}, {}, {}, cnot, {}, {}, {}, {}});
366370

367371
JacobianData<StateVectorT> tape{
368372
num_params, psi.getLength(), psi.getData(), {obs}, ops, tp};

‎pennylane_lightning/core/src/algorithms/tests/mpi/Test_AdjointJacobianMPI.cpp

+6-2
Original file line numberDiff line numberDiff line change
@@ -249,8 +249,10 @@ template <typename TypeList> void testAdjointJacobian() {
249249
"PauliX", std::vector<size_t>{1}),
250250
std::make_shared<NamedObsMPI<StateVectorT>>(
251251
"PauliX", std::vector<size_t>{2}));
252+
std::vector<ComplexT> cnot{1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0,
253+
0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 1.0, 0.0};
252254
auto ops = OpsData<StateVectorT>(
253-
{"RZ", "RY", "RZ", "CNOT", "CNOT", "RZ", "RY", "RZ"},
255+
{"RZ", "RY", "RZ", "QubitUnitary", "CNOT", "RZ", "RY", "RZ"},
254256
{{param[0]},
255257
{param[1]},
256258
{param[2]},
@@ -260,7 +262,9 @@ template <typename TypeList> void testAdjointJacobian() {
260262
{param[1]},
261263
{param[2]}},
262264
{{0}, {0}, {0}, {0, 1}, {1, 2}, {1}, {1}, {1}},
263-
{false, false, false, false, false, false, false, false});
265+
{false, false, false, false, false, false, false, false},
266+
std::vector<std::vector<ComplexT>>{
267+
{}, {}, {}, cnot, {}, {}, {}, {}});
264268

265269
JacobianDataMPI<StateVectorT> tape{num_params, psi, {obs}, ops, tp};
266270

0 commit comments

Comments
 (0)