Add VJP support to PL-Lightning #181

maliasadi · 2021-11-25T06:55:20Z

Context:
Implement and Integrate Vector Jacobian Product (VJP) methods in PL-Lightning along with adding tests following PL implementation here and using pybind11.

Description of the Change:
Add vecMatrixProd to Util.hpp and class VectorJacobianProduct to AdjointDiff.cpp with two public methods:

tensorDot: Calculates tensor-dot of a given vector of gradient outputs and a Jacobian.
vectorJacobianProduct: Calculates the vjp for the selected set of parametric gates utilizing the AdjointJacobian method.

These methods are fully tested at Test_VectorJacobianProduct.cpp.

Python bindings are added to Bindings.cpp and lightning.qubit is now offering two methods to compute vjps:

vector_jacobian_product: Generate the the vector-Jacobian products of a tape.
batch_vjp: Generate the the vector-Jacobian products of a batch of tapes.

These methods are fully tested at test_vector_jacobian_product.

…htning into lightning-add-vjp

codecov · 2021-11-25T06:59:47Z

Codecov Report

Merging #181 (fb80f55) into master (82c197b) will increase coverage by 1.26%.
The diff coverage is 98.41%.

❗ Current head fb80f55 differs from pull request most recent head 9df0c3a. Consider uploading reports for the commit 9df0c3a to get more accurate results

@@            Coverage Diff             @@
##           master     #181      +/-   ##
==========================================
+ Coverage   92.97%   94.23%   +1.26%     
==========================================
  Files           4        4              
  Lines         185      243      +58     
==========================================
+ Hits          172      229      +57     
- Misses         13       14       +1

Impacted Files	Coverage Δ
pennylane_lightning/lightning_qubit.py	`92.16% <98.41%> (+3.27%)`	⬆️

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 82c197b...9df0c3a. Read the comment docs.

github-actions · 2021-11-25T07:01:04Z

Test Report (C++) on Ubuntu

      1 files ±    0       1 suites ±0 0s ⏱️ ±0s
  368 tests +  23   368 ✔️ +  23 0 💤 ±0 0 ❌ ±0
2 205 runs +286 2 205 ✔️ +286 0 💤 ±0 0 ❌ ±0

Results for commit 9df0c3a. ± Comparison against base commit 82c197b.

♻️ This comment has been updated with latest results.

…nto lightning-add-vjp

…lightning into lightning-add-vjp

mlxd

Phenomenal work @maliasadi !
I had a few small comments. Also, just wondering if we can add tests that trigger the conditions CodeFactor/CodeCov complain about?

mlxd · 2021-11-29T09:28:16Z

pennylane_lightning/src/algorithms/AdjointDiff.cpp

+template class Pennylane::Algorithms::VectorJacobianProduct<float>;
+template class Pennylane::Algorithms::VectorJacobianProduct<double>;


Do you think this should be included in the Adjoint file, or separated into its own file in the algorithms directory?

it's definitely better to be added into its own file namely JacobianProd.cpp in the same directory.

mlxd · 2021-11-29T09:30:11Z

pennylane_lightning/src/tests/CMakeLists.txt

                                Test_StateVector_Nonparam.cpp 
                                Test_StateVector_Param.cpp 
                                Test_StateVectorManaged_Nonparam.cpp 
                                Test_StateVectorManaged_Param.cpp 
-                                Test_Util.cpp
+                                Test_Bindings.cpp


Just curious about the reordering of these? If ordering doesn't matter, I think keeping alphabetical may be a good idea.

Nice catch! It wasn't intentional 🙂

pennylane_lightning/src/util/Util.hpp

josh146

Thanks @maliasadi! This is great, I've left some comments throughout.

Correct me if I am wrong, but currently the built-in Jacobian method only supports gradients technically, since only tapes with a single, scalar, output is supported?

As we iterate the gradient support in PL and lightning, I think it makes sense to 'decouple' the VJP and Jacobian methods at some point, so that a full Jacobian is not reconstructed when the VJP is requested

pennylane_lightning/lightning_qubit.py

josh146 · 2021-11-29T11:08:06Z

tests/test_vector_jacobian_product.py

+
+        vjps = dev.batch_vjp(tapes, dys, vjp_pybind=vjp_pybind)
+
+        assert sum(len(t) for t in vjps) == sum(len(t.trainable_params) for t in tapes)


really nice testing!!

maliasadi · 2021-11-30T21:58:54Z

just wondering if we can add tests that trigger the conditions CodeFactor/CodeCov complain about?

Thank you @mlxd! I added adjoint_diff_support_check method that raises QuantumFunctionError in case of finding not supported measurements, observables, or operations by the Lightning adjoint differentiation method for a given tape, and used this method in both adjoint_jacobian and vector_jacobian_product to fix the CodeFactor complain. Please, Let me know if you think this is not the best approach to tackle this issue.

maliasadi · 2021-12-01T17:38:42Z

Thank you @josh146! 🙂

As we iterate the gradient support in PL and lightning, I think it makes sense to 'decouple' the VJP and Jacobian methods at some point, so that a full Jacobian is not reconstructed when the VJP is requested

You're right! However, one can use compute_vjp to compute VJP from a pre-calculated Jacobian. I hope this addresses the issue with reconstructing Jacobian in the vector_jacobian_product method. Moreover, I updated vector_jacobian_product and batch_vjp to also return the adjoint_jacobian results. Please let me know what you think about these new changes.

mlxd

Great job putting this together @maliasadi

As discussed, I think we can safely merge this and build upon it later for the optimal implementation.

josh146

Agree with Lee, great work @maliasadi!

maliasadi added 6 commits November 24, 2021 11:21

Add vecMatrixProd

b83d920

Merge branch 'master' of https://github.com/PennyLaneAI/pennylane-lig…

9d6c7da

…htning into lightning-add-vjp

Add vjp to lightning_qubit.py and tests

c3c7525

Update formatting

77fc34a

Add class VectorJacobianProduct

ac39aaa

Update formatting

1210aac

maliasadi added 15 commits November 25, 2021 02:04

Update C++ class

90ea7ed

Update clang-tidy

94ced72

Update VectorJacobianProduct public methods

b219561

Add Test_VJP.cpp

9c163b8

Update cpp tests

516d306

Update cpp tests

d833c44

Complete adding cpp tests

74aeeac

Merge branch 'master' of github.com:PennyLaneAI/pennylane-lightning i…

a3896b5

…nto lightning-add-vjp

Update formatting w/ clang-tidy-12

e2f8bd6

Merge branch 'master' into lightning-add-vjp

b30d626

Merge branch 'lightning-add-vjp' of github.com:PennyLaneAI/pennylane-…

d85f20b

…lightning into lightning-add-vjp

Add Bindings

6d0b610

Update Bindings

d249d5e

Fix issue with vjp pybinds and Add more tests

09cbfcd

Update #181

ed76b80

maliasadi changed the title ~~[WIP] Add VJP support to PL-Lightning~~ Add VJP support to PL-Lightning Nov 26, 2021

maliasadi requested review from mlxd and josh146 November 26, 2021 19:54

Apply codecov suggestions

99e7928

maliasadi marked this pull request as ready for review November 26, 2021 20:49

mlxd requested changes Nov 29, 2021

View reviewed changes

josh146 requested changes Nov 29, 2021

View reviewed changes

maliasadi added 3 commits November 30, 2021 15:30

Apply code review suggestions

205f874

Fix rendering math formulas in docs

25b93e8

Apply code factor suggestions

b9416e6

maliasadi closed this Nov 30, 2021

maliasadi reopened this Nov 30, 2021

maliasadi added 2 commits November 30, 2021 16:32

Update python tests

30534c0

Update vector_jacobian_product method

20b1e12

maliasadi added 2 commits November 30, 2021 17:52

Add adjoint_diff_support_check method

c0f4202

Add more tests for batch_vjp

3d0bf52

maliasadi requested a review from mlxd November 30, 2021 23:19

maliasadi added 2 commits December 1, 2021 12:10

Update VJP Python bindings

6bca830

Update tests

0742302

maliasadi requested review from josh146 and AmintorDusko December 1, 2021 17:39

Merge branch 'master' into lightning-add-vjp

8cfb4ab

mlxd approved these changes Dec 3, 2021

View reviewed changes

maliasadi requested review from AmintorDusko and removed request for AmintorDusko December 3, 2021 17:58

josh146 approved these changes Dec 6, 2021

View reviewed changes

Merge branch 'master' into lightning-add-vjp

9df0c3a

mlxd merged commit be1599f into master Dec 6, 2021

mlxd deleted the lightning-add-vjp branch December 6, 2021 11:39

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add VJP support to PL-Lightning #181

Add VJP support to PL-Lightning #181

maliasadi commented Nov 25, 2021 •

edited

Loading

codecov bot commented Nov 25, 2021 •

edited

Loading

github-actions bot commented Nov 25, 2021 •

edited

Loading

mlxd left a comment

mlxd Nov 29, 2021

maliasadi Nov 30, 2021

mlxd Nov 29, 2021

maliasadi Nov 30, 2021

josh146 left a comment

josh146 Nov 29, 2021

maliasadi commented Nov 30, 2021 •

edited

Loading

maliasadi commented Dec 1, 2021

mlxd left a comment

josh146 left a comment

		template class Pennylane::Algorithms::VectorJacobianProduct<float>;
		template class Pennylane::Algorithms::VectorJacobianProduct<double>;


		vjps = dev.batch_vjp(tapes, dys, vjp_pybind=vjp_pybind)

		assert sum(len(t) for t in vjps) == sum(len(t.trainable_params) for t in tapes)

Add VJP support to PL-Lightning #181

Add VJP support to PL-Lightning #181

Conversation

maliasadi commented Nov 25, 2021 • edited Loading

codecov bot commented Nov 25, 2021 • edited Loading

Codecov Report

github-actions bot commented Nov 25, 2021 • edited Loading

Test Report (C++) on Ubuntu

mlxd left a comment

Choose a reason for hiding this comment

mlxd Nov 29, 2021

Choose a reason for hiding this comment

maliasadi Nov 30, 2021

Choose a reason for hiding this comment

mlxd Nov 29, 2021

Choose a reason for hiding this comment

maliasadi Nov 30, 2021

Choose a reason for hiding this comment

josh146 left a comment

Choose a reason for hiding this comment

josh146 Nov 29, 2021

Choose a reason for hiding this comment

maliasadi commented Nov 30, 2021 • edited Loading

maliasadi commented Dec 1, 2021

mlxd left a comment

Choose a reason for hiding this comment

josh146 left a comment

Choose a reason for hiding this comment

maliasadi commented Nov 25, 2021 •

edited

Loading

codecov bot commented Nov 25, 2021 •

edited

Loading

github-actions bot commented Nov 25, 2021 •

edited

Loading

maliasadi commented Nov 30, 2021 •

edited

Loading