Register `Operator` classes and pytrees #4458

albi3ro · 2023-08-10T21:34:52Z

This PR adds a pennylane/pytrees.py module that contains the code for registering a custom object with a pytree framework. The only framework is currently jax, but we can extend that in the future.

The dynamic typing in Adjoint.__new__ caused some issues, so that was slightly rewritten, and is no longer causing issues.

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

…LaneAI/pennylane into operator-flatten-unflatten

codecov · 2023-08-10T22:07:00Z

Codecov Report

Merging #4458 (98f1260) into master (af1eba4) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #4458   +/-   ##
=======================================
  Coverage   99.71%   99.71%           
=======================================
  Files         377      378    +1     
  Lines       34180    34197   +17     
=======================================
+ Hits        34084    34101   +17     
  Misses         96       96

Files Changed	Coverage Δ
pennylane/operation.py	`97.41% <100.00%> (+0.01%)`	⬆️
pennylane/ops/op_math/adjoint.py	`99.19% <100.00%> (-0.03%)`	⬇️
pennylane/pytrees.py	`100.00% <100.00%> (ø)`

vincentmr

Nicely done. Thanks @albi3ro .

BorjaRequena

Looking good @albi3ro! Thanks for taking care of this pytreefication, it will allow us to do a lot of cool things with operators 🚀 Which makes me wonder if we should find a compact way to test those, like testing jax.jit, jax.grad and jax.vmap on some operator functions, if that even makes sense 🤔

Also, could you add this line that we could not include in the previous PR? 😄

tests/ops/qubit/test_parametric_ops.py

pennylane/operation.py

pennylane/pytrees.py

tests/test_operation.py

timmysilv

very exciting, looks good! do we want to add any details about ops-as-pytrees to the docs, perhaps just in the internal section? I'm unsure myself.

Also unrelated but I realized while checking out the Adjoint docs - the contents of the "Developer Details" section of the docs are not indented, so they don't hide in the HTML drop-down when rendered. Could you indent them here?

Co-authored-by: BorjaRequena <59647767+BorjaRequena@users.noreply.github.com>

josh146 · 2023-08-16T13:56:40Z

@timmysilv -- could be a good section to add to the https://docs.pennylane.ai/en/stable/development/adding_operators.html page?

albi3ro · 2023-08-16T14:04:16Z

very exciting, looks good! do we want to add any details about ops-as-pytrees to the docs, perhaps just in the internal section? I'm unsure myself.

Also unrelated but I realized while checking out the Adjoint docs - the contents of the "Developer Details" section of the docs are not indented, so they don't hide in the HTML drop-down when rendered. Could you indent them here?

@josh146 @timmysilv

https://docs.pennylane.ai/en/latest/code/api/pennylane.operation.Operator.html#serialization

josh146 · 2023-08-16T14:23:21Z

Looks good @albi3ro! Maybe out of scope for this PR, but it looks like in the sentence

See the Operator._flatten and Operator._unflatten methods for more information.

the links don't work 🤔 Is it because they are private methods/are not showing in the docs?

albi3ro · 2023-08-16T14:33:11Z

Looks good @albi3ro! Maybe out of scope for this PR, but it looks like in the sentence

See the Operator._flatten and Operator._unflatten methods for more information.

the links don't work 🤔 Is it because they are private methods/are not showing in the docs?

Private methods.

josh146 · 2023-08-16T14:57:57Z

Got it, thanks @albi3ro. In that case, we may need to revisit that docstring, to make it easier to direct people looking for more information (or, if we do want Operator._flatten etc. to be documented, make them public methods)

albi3ro · 2023-08-16T15:03:44Z

Got it, thanks @albi3ro. In that case, we may need to revisit that docstring, to make it easier to direct people looking for more information (or, if we do want Operator._flatten etc. to be documented, make them public methods)

I wonder if we could consider making private methods on Operator show up in the documentation. They are private because nothing outside of the class should touch them, but they are relevant to anyone trying to implement the class.

trbromley · 2023-08-16T18:57:50Z

pennylane/pytrees.py

+
+has_jax = True
+try:
+    import jax.tree_util as jax_tree_util


Is there a way we can avoid importing JAX with PennyLane?

No.

And i saw no import time difference.

This is a shame, it locks us in to having to import JAX at the same time as PennyLane.

* more custom definitions * more templates * more tests * more tests for templates * Update doc/development/deprecations.rst * more tests * final round of template tests * remove qml.equal change * Update pennylane/templates/embeddings/angle.py * fix op_math tests * Update pennylane/templates/layers/random.py Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> * Update doc/development/deprecations.rst Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> * finishing touches * some more tests and fixes * fix remaining tests, lint * fix test lint * register operator with jax * register with torch too * tidying * register operations as pytrees * merge * adding tests * Removing changes we didn't need to make to test_adjoint * sphinx, black, and pylint * adjoint changes and pylint [skip ci] * test for adjoint pickling * black and pylint * Update tests/ops/qubit/test_parametric_ops.py Co-authored-by: BorjaRequena <59647767+BorjaRequena@users.noreply.github.com> * indent developer details section * black * add multiprocessing with adjoint test --------- Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com> Co-authored-by: BorjaRequena <59647767+BorjaRequena@users.noreply.github.com>

albi3ro and others added 28 commits June 23, 2023 10:59

more custom definitions

acb772f

more templates

4a9d2ec

Merge branch 'master' into operator-flatten-unflatten

b164f18

Merge branch 'master' into operator-flatten-unflatten

472e2bf

more tests

b547114

more tests for templates

9797ef5

Merge branch 'master' into operator-flatten-unflatten

e0238b4

Update doc/development/deprecations.rst

17fbb14

more tests

6ec0a79

merge

ab4df86

final round of template tests

858c7be

Merge branch 'master' into operator-flatten-unflatten

87038f0

remove qml.equal change

ba0a53b

Update pennylane/templates/embeddings/angle.py

d720d74

fix op_math tests

61d3527

Update pennylane/templates/layers/random.py

ebced5e

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

Update doc/development/deprecations.rst

7c0a342

Co-authored-by: Tom Bromley <49409390+trbromley@users.noreply.github.com>

Merge branch 'operator-flatten-unflatten' of https://github.com/Penny…

da54ab2

…LaneAI/pennylane into operator-flatten-unflatten

finishing touches

2f24da9

some more tests and fixes

781a294

fix remaining tests, lint

2ba0d05

fix test lint

98cb842

register operator with jax

1d0f1f7

register with torch too

6519cf9

tidying

6bb9763

register operations as pytrees

ccb41ac

Merge branch 'master' into register-operator-pytree

194285d

merge

0130ead

adding tests

501c8bc

vincentmr approved these changes Aug 14, 2023

View reviewed changes

mudit2812 added this to the v0.32 milestone Aug 14, 2023

BorjaRequena mentioned this pull request Aug 14, 2023

qml.Projector compatibility with new default qubit #4452

Merged

albi3ro and others added 4 commits August 15, 2023 13:35

adjoint changes and pylint [skip ci]

2e1d359

test for adjoint pickling

39a7cd6

black and pylint

04bea2f

Merge branch 'master' into register-operator-pytree

498cfb5

albi3ro requested review from timmysilv and BorjaRequena August 15, 2023 21:17

BorjaRequena reviewed Aug 16, 2023

View reviewed changes

tests/ops/qubit/test_parametric_ops.py Show resolved Hide resolved

pennylane/operation.py Show resolved Hide resolved

pennylane/pytrees.py Show resolved Hide resolved

pennylane/pytrees.py Show resolved Hide resolved

tests/test_operation.py Show resolved Hide resolved

timmysilv approved these changes Aug 16, 2023

View reviewed changes

Update tests/ops/qubit/test_parametric_ops.py

a374721

Co-authored-by: BorjaRequena <59647767+BorjaRequena@users.noreply.github.com>

Merge branch 'master' into register-operator-pytree

fc415c5

indent developer details section

844ba37

albi3ro enabled auto-merge (squash) August 16, 2023 14:07

black

a81bd76

albi3ro added the merge-ready ✔️ All tests pass and the PR is ready to be merged. label Aug 16, 2023

albi3ro disabled auto-merge August 16, 2023 15:42

add multiprocessing with adjoint test

98f1260

albi3ro enabled auto-merge (squash) August 16, 2023 15:57

albi3ro merged commit 9aab699 into master Aug 16, 2023
39 checks passed

albi3ro deleted the register-operator-pytree branch August 16, 2023 17:02

trbromley reviewed Aug 16, 2023

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Register `Operator` classes and pytrees #4458

Register `Operator` classes and pytrees #4458

albi3ro commented Aug 10, 2023 •

edited

Loading

codecov bot commented Aug 10, 2023 •

edited

Loading

vincentmr left a comment

BorjaRequena left a comment

timmysilv left a comment

josh146 commented Aug 16, 2023

albi3ro commented Aug 16, 2023 •

edited

Loading

josh146 commented Aug 16, 2023

albi3ro commented Aug 16, 2023

josh146 commented Aug 16, 2023

albi3ro commented Aug 16, 2023

trbromley Aug 16, 2023

albi3ro Aug 16, 2023

trbromley Aug 17, 2023

Register Operator classes and pytrees #4458

Register Operator classes and pytrees #4458

Conversation

albi3ro commented Aug 10, 2023 • edited Loading

codecov bot commented Aug 10, 2023 • edited Loading

Codecov Report

vincentmr left a comment

Choose a reason for hiding this comment

BorjaRequena left a comment

Choose a reason for hiding this comment

timmysilv left a comment

Choose a reason for hiding this comment

josh146 commented Aug 16, 2023

albi3ro commented Aug 16, 2023 • edited Loading

josh146 commented Aug 16, 2023

albi3ro commented Aug 16, 2023

josh146 commented Aug 16, 2023

albi3ro commented Aug 16, 2023

trbromley Aug 16, 2023

Choose a reason for hiding this comment

albi3ro Aug 16, 2023

Choose a reason for hiding this comment

trbromley Aug 17, 2023

Choose a reason for hiding this comment

Register `Operator` classes and pytrees #4458

Register `Operator` classes and pytrees #4458

albi3ro commented Aug 10, 2023 •

edited

Loading

codecov bot commented Aug 10, 2023 •

edited

Loading

albi3ro commented Aug 16, 2023 •

edited

Loading