pennylane icon indicating copy to clipboard operation
pennylane copied to clipboard

Published 20 hours ago • verified publisher icon PennyLaneAI

Use native pytrees with jax interface

Open albi3ro opened this issue 1 year ago • 6 comments

[sc-42956]

Part of the motivation between binding pennylane objects (operators, measurements, and tapes) as pytrees was the opportunity to improve how we bind derivatives to jax.

Now that all pennylane operators are registered as valid pytrees, we can use that to alter how we bind derivatives to jax.

The benefits of this change are:

  • Simplified code

The downsides are:

Sometimes operators are not valid pytrees for some reason, or are not able to store tangents in place of the original data. Either of those conditions will break the machine learning interface.

albi3ro avatar Dec 22 '23 00:12 albi3ro

Hello. You may have forgotten to update the changelog! Please edit doc/releases/changelog-dev.md with:

  • A one-to-two sentence description of the change. You may include a small working example for new features.
  • A link back to this PR.
  • Your name (or GitHub username) in the contributors section.

github-actions[bot] avatar Dec 22 '23 00:12 github-actions[bot]

@albi3ro, what is the status of this PR? The story in Shortcut is in the 0.36 backlog epic.

trbromley avatar Apr 15 '24 13:04 trbromley

@albi3ro, what is the status of this PR? The story in Shortcut is in the 0.36 backlog epic.

Now unblocked since the projector PR is in.

albi3ro avatar May 16 '24 18:05 albi3ro

Nice! Is there more to do here (aside from resolving merge conflicts)? Can we open it up for review?

trbromley avatar May 24 '24 12:05 trbromley

And this PR isn't replaced by https://github.com/PennyLaneAI/pennylane/pull/5701?

trbromley avatar May 24 '24 13:05 trbromley

@albi3ro do we still need this PR?

trbromley avatar Jul 22 '24 12:07 trbromley