Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add bind_nested_plxpr meta-transform #5520

Closed
wants to merge 30 commits into from
Closed

Conversation

albi3ro
Copy link
Contributor

@albi3ro albi3ro commented Apr 16, 2024

HIGHLY EXPERIMENTAL

Context:

This is part of an ongoing project to improve the ability to capture pennylane code into a hybrid intermediate representation.

Description of the Change:

This PR adds a qml.capture.bind_nested_plxpr meta-transform that automatically integrates qfunc transforms with
PLXPR. It:

  • Creates the associated primitive
  • Defines concrete and abstract implementations for the primitive
  • Hijacks calling the quantum function to instead create jaxpr and bind it to the primitive

Benefits:

Easy binding of qfunc transforms into PLXPR via a straightforward decorator/ meta-transform.

Possible Drawbacks:

The handling of keyword arguments is a bit shaky.

This level of nesting in function definitions can be really hard to read.

When I was messing around with it, I think I saw a memory leak, and not sure if its still there or not.

Related GitHub Issues:

[sc-59448]

Copy link
Contributor

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.

Base automatically changed from plxpr-capture-operations to master May 10, 2024 22:14
@trbromley
Copy link
Contributor

@albi3ro what is the status of this PR, do we still need it?

@albi3ro albi3ro closed this Aug 20, 2024
@albi3ro albi3ro deleted the bind-nested-jaxpr branch August 20, 2024 14:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

3 participants