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

[Capture] Integrate plxpr transforms with execution pipeline #6914

Draft
wants to merge 12 commits into
base: master
Choose a base branch
from

Conversation

mudit2812
Copy link
Contributor

@mudit2812 mudit2812 commented Jan 31, 2025

Before submitting

Please complete the following checklist when submitting a PR:

  • All new features must include a unit test.
    If you've fixed a bug or added code that should be tested, add a test to the
    test directory!

  • All new functions and code must be clearly commented and documented.
    If you do make documentation changes, make sure that the docs build and
    render correctly by running make docs.

  • Ensure that the test suite passes, by running make test.

  • Add a new entry to the doc/releases/changelog-dev.md file, summarizing the
    change, and including a link back to the PR.

  • The PennyLane source code conforms to
    PEP8 standards.
    We check all of our code against Pylint.
    To lint modified files, simply pip install pylint, and then
    run pylint pennylane/path/to/file.py.

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.


Context:

Description of the Change:

Benefits:

Possible Drawbacks:

Related GitHub Issues:
[sc-83799] [sc-84017]

@mudit2812 mudit2812 changed the title Integrate plxpr transforms with execution pipeline [Capture] Integrate plxpr transforms with execution pipeline Jan 31, 2025
@mudit2812 mudit2812 force-pushed the capture-transforms-integration branch from 192abd6 to 7eeac8e Compare January 31, 2025 22:51
Comment on lines +572 to +574
if not self:
f = partial(qml.capture.PlxprInterpreter().eval, jaxpr, consts)
return jax.make_jaxpr(f)(*args)
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If there's no transforms, do we need to recapture it at all?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I only added this to adhere to the return type spec (jax.core.ClosedJaxpr). I don't want to change the return type spec because jax.core.ClosedJaxpr carries more info than jax.core.Jaxpr (both jaxpr and consts). I could update it to return tuple[jax.core.Jaxpr, Sequence[Number]], though that seems inconsistent with everything else we currently have.

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.

2 participants