Add JVP/VJP support #98

grwlf · 2023-04-21T13:37:30Z

[sc-37338]

In this PR we add support for JVP and VJP operations to the Grad MLIR dialect.

Add 1-part JVP and VJP operations at the MLIR level.
Implement lowering to grad.grad and linalg.generic operations.
Implement einsumLinalgGeneric helper function in MLIR.
Provide the required Python JAX API extensions.

Note that this PR includes a workaround for #111 . Otherwise it would fail some tests.

Review process (summary on the time-demanding work):

Move MLIR patterns to separate files, adjust headers, etc.
Make tests non-parameterized
Split the JVP/VJP callee arguments from the (co-)tangent arguments
Split the JVP/VJP callee results from the jvp/vjp results
Unify the MLIR symbol usage verifiers for grad/jvp/vjp
Rebase and test the PR against the Temporary fix for memory issues. #121 workaround

codecov · 2023-04-21T13:45:18Z

Codecov Report

Merging #98 (34477ab) into main (e40a94e) will increase coverage by 0.08%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main      #98      +/-   ##
==========================================
+ Coverage   98.85%   98.93%   +0.08%     
==========================================
  Files          37       37              
  Lines        6215     6310      +95     
  Branches      317      328      +11     
==========================================
+ Hits         6144     6243      +99     
+ Misses         40       38       -2     
+ Partials       31       29       -2

Impacted Files	Coverage Δ
frontend/catalyst/__init__.py	`90.00% <100.00%> (ø)`
frontend/catalyst/jax_primitives.py	`96.42% <100.00%> (+0.30%)`	⬆️
frontend/catalyst/pennylane_extensions.py	`97.42% <100.00%> (+1.31%)`	⬆️
frontend/catalyst/utils/calculate_grad_shape.py	`92.10% <100.00%> (ø)`

erick-xanadu

Looks good. Some changes, some comments. I haven't yet looked to deeply into the JVP/VJP algorithm. That will take me a bit longer.

frontend/catalyst/compiler.py

frontend/catalyst/jax_primitives.py

frontend/catalyst/pennylane_extensions.py

frontend/test/pytest/test_jvpvjp.py

mlir/lib/Gradient/Transforms/jvp_vjp_lowering.cpp

mlir/test/Gradient/JVPTest.mlir

frontend/catalyst/pennylane_extensions.py

dime10

Hey @grwlf, the PR looks great so far! I haven't gone through all of it yet, but I have a few suggestions around the MLIR ops for you.

It might be easier to review once all the debugging code is cleaned up.

frontend/test/pytest/test_jvpvjp.py

mlir/include/Gradient/IR/GradientOps.td

mlir/lib/Gradient/IR/GradientOps.cpp

mlir/lib/Gradient/Transforms/jvp_vjp_lowering.cpp

mlir/include/Gradient/IR/GradientOps.td

mlir/lib/Gradient/Transforms/jvp_vjp_lowering.cpp

grwlf · 2023-05-11T13:42:14Z

FYI: I found problems with the handling of scalar arguments, now I'm working on the correction. The main problem here is how to complete the einsum definition if some or all arguments are scalars.

UPDATE 1: I fixed the scalar problem quickly, but after that I faced another problem with numeric differences between PL and C quantum program gradients.

UPDATE 2: The numeric problem was due to the incorrect usage of PL/JAX interface. Should be fine now.

frontend/test/pytest/test_jvpvjp.py

frontend/catalyst/compiler.py

runtime/lib/capi/RuntimeCAPI.cpp

mlir/include/Gradient/IR/GradientOps.td

dime10

Had some time to look at the rest of the code, nice contribution @grwlf ! 🙂

mlir/lib/Gradient/Transforms/GradMethods/JVPVJPPatterns.hpp

mlir/lib/Gradient/Transforms/jvpvjp_lowering.cpp

mlir/lib/Gradient/Utils/EinsumLinalgGeneric.cpp

frontend/catalyst/pennylane_extensions.py

dime10

Hi @grwlf, this PR should be ready to merge as soon as the last remaining comments are resolved and all debug/commented out code is removed :)

mlir/lib/Gradient/IR/GradientOps.cpp

mlir/include/Gradient/IR/GradientOps.td

mlir/lib/Gradient/Transforms/gradient_lowering.cpp

doc/changelog.md

dime10

Great work @grwlf 💯

Looks good to me!

mlir/lib/Gradient/Transforms/GradMethods/JVPVJPPatterns.cpp

doc/changelog.md

frontend/catalyst/pennylane_extensions.py

Co-authored-by: David Ittah <dime10@users.noreply.github.com>

This reverts commit 26b2fdb.

…it usage

…tiable

…jvp/vjp results

This reverts commit a1ec311.

…ove some debug messages

Co-authored-by: David Ittah <dime10@users.noreply.github.com> Co-authored-by: Josh Izaac <josh146@gmail.com>

…cussion)

…llable in comments

grwlf changed the title ~~Add JVP jax extension~~ Add JVP/VJP support Apr 21, 2023

grwlf force-pushed the jvp_vjp branch from f75781e to 9f12683 Compare April 24, 2023 19:23

grwlf force-pushed the jvp_vjp branch 2 times, most recently from 6e80e6b to d2c5813 Compare May 10, 2023 08:45

grwlf requested review from dime10 and erick-xanadu May 10, 2023 10:40

grwlf marked this pull request as ready for review May 10, 2023 10:40

erick-xanadu reviewed May 10, 2023

View reviewed changes

josh146 reviewed May 10, 2023

View reviewed changes

frontend/catalyst/pennylane_extensions.py Outdated Show resolved Hide resolved

dime10 reviewed May 10, 2023

View reviewed changes

grwlf force-pushed the jvp_vjp branch from ff179fb to d19e55f Compare May 11, 2023 20:16

dime10 reviewed May 12, 2023

View reviewed changes

PennyLaneAI deleted a comment from grwlf May 15, 2023

dime10 reviewed May 16, 2023

View reviewed changes

grwlf force-pushed the jvp_vjp branch 3 times, most recently from 30e592c to 8bf4a38 Compare May 17, 2023 10:59

grwlf requested review from dime10 and josh146 May 17, 2023 11:26

dime10 reviewed May 17, 2023

View reviewed changes

doc/changelog.md Outdated Show resolved Hide resolved

dime10 approved these changes May 17, 2023

View reviewed changes

mlir/lib/Gradient/Transforms/GradMethods/JVPVJPPatterns.cpp Outdated Show resolved Hide resolved

grwlf mentioned this pull request May 17, 2023

Update the project documentation with the testing guidelines #133

Open

dime10 reviewed May 17, 2023

View reviewed changes

doc/changelog.md Outdated Show resolved Hide resolved

dime10 reviewed May 17, 2023

View reviewed changes

doc/changelog.md Outdated Show resolved Hide resolved

josh146 reviewed May 18, 2023

View reviewed changes

frontend/catalyst/pennylane_extensions.py Outdated Show resolved Hide resolved

josh146 reviewed May 18, 2023

View reviewed changes

frontend/catalyst/pennylane_extensions.py Outdated Show resolved Hide resolved

josh146 reviewed May 18, 2023

View reviewed changes

frontend/catalyst/pennylane_extensions.py Outdated Show resolved Hide resolved

grwlf and others added 27 commits May 18, 2023 15:55

Apply suggestions from code review

80e9f6e

Co-authored-by: David Ittah <dime10@users.noreply.github.com>

Revert "Update frontend/catalyst/jax_primitives.py"

fb33557

This reverts commit 26b2fdb.

Revert review suggestion on lowering function argument names

171a009

Fix a typo

930acef

Address the formatting

e8c16be

Review suggestions: move jvp pattern registration code

2914f96

Review suggestions: move compDiffGradIndices to Gradient/Utils

da30424

Review suggestions: use built-in operand accessors; fix incorrect tra…

9bc37a9

…it usage

Address formatting issues

586acb9

Review suggestions: use old name for calculate_grad_shape

394b8cd

Review suggestions: support the lowerOnly == "jp" case

7f1d07d

Review suggestions: hopefully add a better name for make_jax_differen…

693cd23

…tiable

Rebase on top of #121; Fix conflicts; Make tests accept iterables as …

2663c01

…jvp/vjp results

Revert "[HACK] Disabling CSE and mlir_free to workaround #111"

83f241b

This reverts commit a1ec311.

Address format issues

37ed9a0

Remove unused function declaration

9f4c781

Update the changelog

57c6728

Add missing argnum test cases

d98fc16

Review suggestions: update the changelog

2155266

Review suggestions: add missing diffArgIndices vs tangents check; rem…

4c65ae9

…ove some debug messages

Address format issues

846fcba

Review suggestions: add changelog.md example

14d6d54

Fix grammar mistakes

1de33bc

Co-authored-by: David Ittah <dime10@users.noreply.github.com> Co-authored-by: Josh Izaac <josh146@gmail.com>

Address CodeFactor issues

72aed53

Revert the comment change made by accident (it is currently under dis…

fb98b1e

…cussion)

Review suggestions: make DifferentiableLike to be a type hint, use Ca…

652d5d4

…llable in comments

Fix rebase artifacts

34477ab

grwlf force-pushed the jvp_vjp branch from 22f5ce3 to 34477ab Compare May 18, 2023 16:14

grwlf merged commit 4721af3 into main May 18, 2023
15 checks passed

grwlf deleted the jvp_vjp branch May 18, 2023 16:30

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add JVP/VJP support #98

Add JVP/VJP support #98

grwlf commented Apr 21, 2023 •

edited

codecov bot commented Apr 21, 2023 •

edited

erick-xanadu left a comment

dime10 left a comment •

edited

grwlf commented May 11, 2023 •

edited

dime10 left a comment

dime10 left a comment

dime10 left a comment

Add JVP/VJP support #98

Add JVP/VJP support #98

Conversation

grwlf commented Apr 21, 2023 • edited

codecov bot commented Apr 21, 2023 • edited

Codecov Report

erick-xanadu left a comment

Choose a reason for hiding this comment

dime10 left a comment • edited

Choose a reason for hiding this comment

grwlf commented May 11, 2023 • edited

dime10 left a comment

Choose a reason for hiding this comment

dime10 left a comment

Choose a reason for hiding this comment

dime10 left a comment

Choose a reason for hiding this comment

grwlf commented Apr 21, 2023 •

edited

codecov bot commented Apr 21, 2023 •

edited

dime10 left a comment •

edited

grwlf commented May 11, 2023 •

edited