add from_plxpr conversion function #837

albi3ro · 2024-06-20T18:05:14Z

Context:

This PR ports the pennylane PR #5883 to catalyst and renames it from_plxpr.

Adding this code to catalyst solves many of the dependency and testing issues that arise from placing this code in the pennylane resposity. In the future, if we move more of the catalyst frontend into pennylane, this function may move there as well.

Description of the Change:

Adds a function for converting pennylane variant jaxpr into catalyst variant jaxpr.

Benefits:

Opens up the ability to have improved program capture directly integrated with catalyst.

Possible Drawbacks:

Related GitHub Issues:

[sc-61537]

…into from-plxpr

dwierichs

Amazing work @albi3ro
This approach is definitely much nicer to read :) But that's also due to the clear structure you gave these functions 💯

I had a good bit of nitpicky language/naming/comment/error message comments, otherwise some optional suggestions and ideas.

doc/changelog.md

frontend/catalyst/from_plxpr.py

dwierichs · 2024-06-21T07:21:20Z

frontend/catalyst/from_plxpr.py

+    # Note that the value of rtd_kwargs is a string version of
+    # the info kwargs, not the info kwargs itself


Might be nice for future development to mention why we do this :)

Yeah... I'm not sure why. Think it has to do with mlir.

Okay! @dime10 , do you have context here? I think it might be nice to comment on it :)

Yes! This is just the easiest way to serialize the data since it is going through the IR. The alternative would be to provide typed attributes for each config option which is more work.

Thanks for the info. Adding that as a clarifying source code comment.

frontend/catalyst/from_plxpr.py

dwierichs · 2024-06-21T08:00:19Z

frontend/test/pytest/test_from_plxpr.py

+        jaxpr = jax.make_jaxpr(circuit)()
+        qml.capture.disable()
+
+        with pytest.raises(NotImplementedError):


Reusing a suggestion from the code itself.

Suggested change

with pytest.raises(NotImplementedError):

with pytest.raises(NotImplementedError, match="Only wire-based and observable-based measurements"):

frontend/test/pytest/test_from_plxpr.py

dwierichs · 2024-06-21T08:02:53Z

frontend/test/pytest/test_from_plxpr.py

+        call_jaxpr_pl = converted.eqns[0].params["call_jaxpr"]
+        call_jaxpr_c = catalxpr.eqns[0].params["call_jaxpr"]
+
+        compare_call_jaxprs(call_jaxpr_pl, call_jaxpr_c)


Could parametrize the expval, probs, var tests, they are very similar.

frontend/test/pytest/test_from_plxpr.py

**Context:** Catalyst PR #837 (PennyLaneAI/catalyst#837) needs a couple minor updates to the capture module. **Description of the Change:** 1) makes it possible to do `from pennylane.capture import AbstractOperator, AbstractMeasurement, qnode_prim` so we don't have to touch private functions 2) Adds `qnode` as a keyword argument that gets bound to the qnode primitive 3) Makes it so we can capture a sample measurement specified like `qml.sample(wires=1)` **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:** [sc-66703] --------- Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>

Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>

" \\ z≈ç√ fd" \ `

**Context:** Catalyst PR #837 (PennyLaneAI/catalyst#837) needs a couple minor updates to the capture module. **Description of the Change:** 1) makes it possible to do `from pennylane.capture import AbstractOperator, AbstractMeasurement, qnode_prim` so we don't have to touch private functions 2) Adds `qnode` as a keyword argument that gets bound to the qnode primitive 3) Makes it so we can capture a sample measurement specified like `qml.sample(wires=1)` **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:** [sc-66703] --------- Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>

codecov · 2024-07-15T18:13:27Z

Codecov Report

Attention: Patch coverage is 96.29630% with 5 lines in your changes missing coverage. Please review.

Project coverage is 97.92%. Comparing base (5fa4b21) to head (f1f90aa).

Files	Patch %	Lines
frontend/catalyst/from_plxpr.py	96.29%	2 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #837      +/-   ##
==========================================
- Coverage   97.93%   97.92%   -0.02%     
==========================================
  Files          73       74       +1     
  Lines       10330    10465     +135     
  Branches     1170     1211      +41     
==========================================
+ Hits        10117    10248     +131     
- Misses        170      171       +1     
- Partials       43       46       +3

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

dime10

Nice work!

frontend/catalyst/from_plxpr.py

dime10 · 2024-07-15T18:06:21Z

frontend/catalyst/from_plxpr.py

+    for s in shots:
+        for m in measurements:
+            shape, dtype = m.abstract_eval(shots=s, num_device_wires=num_device_wires)
+            shapes.append(jax.core.ShapedArray(shape, dtype_map.get(dtype, dtype)))


This map is interesting, I wonder if it is actually needed? Abstract JAX arrays can already hold Python types like int, float, complex which are considered "weak types" and will adapt to the bitwidth of "strong types" like int64, or fall back to the same types defined in the map above.

Looks like tests still pass when getting rid of it, so I'm going to say we don't need it until proven otherwise.

frontend/catalyst/from_plxpr.py

dime10 · 2024-07-15T21:00:12Z

frontend/catalyst/from_plxpr.py

+    operator is consumed.
+    """
+
+    def read(self, var):


The code is really nicely structured! The only thing I'm wondering is whether we can't use the interpreter state for the whole conversion procedure, since it seems to duplicate environment and free standing read function above?

I'm thinking of defining a base class for jaxpr interpreters and reusing the structure for a variety of different algorithms using the template design pattern.

You can see a prototype here:

https://github.com/PennyLaneAI/pennylane/blob/plxpr-interpreter/pennylane/capture/interpreters.py

The design and implementation is going to take a little bit of work, but in the end, we will be able to use that design to make this code much nicer :) I do see this code getting restructured once that happens.

Right now we free standing functions with a mutable input, but we could also make them all class methods. I'm fine with promoting everything to class methods now too.

Rewriting te code to follow the structure I'm thinking of for a "Plxpr interpreter template".

dime10 · 2024-07-15T21:04:57Z

frontend/catalyst/from_plxpr.py

+
+    _deallocate(state)
+    # Read the final result of the Jaxpr from the environment
+    return [state.read(outvar) for outvar in measurements]


I'm guessing this is expected for now, but just noting that here we can only ever return MP results, not values from classical ops.

dime10 · 2024-07-15T21:08:58Z

frontend/catalyst/from_plxpr.py

+    wires = [state.get_wire(w) for w in wire_values]
+
+    invals = [state.read(invar) for invar in eqn.invars[:-n_wires]]
+    outvals = qinst_p.bind(


We have two custom primitives because they deviate from the form qinst_p has, whose conversion isn't included here:

qunitary_p

gphase_p

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

dime10 · 2024-07-30T15:24:10Z

Thanks for your comments @albi3ro, let me know once you are finished with them and I'll give it a final look :)

albi3ro added 4 commits June 20, 2024 14:01

add from_plxpr conversion function

606aaf3

Merge branch 'main' into from-plxpr

ae33010

sort imports

01714e3

Merge branch 'from-plxpr' of https://github.com/PennyLaneAI/catalyst …

7700595

…into from-plxpr

albi3ro requested review from dime10 and dwierichs June 20, 2024 18:08

This was referenced Jun 20, 2024

to_catalyst plxpr conversion function using a jaxpr interpreter PennyLaneAI/pennylane#5883

Closed

Minor program capture fixes PennyLaneAI/pennylane#5889

Merged

albi3ro added 2 commits June 20, 2024 15:42

black and pylint:

1e01682

Merge branch 'main' into from-plxpr

ab15a51

dwierichs requested changes Jun 21, 2024

View reviewed changes

Update frontend/catalyst/from_plxpr.py

e82d05e

albi3ro and others added 7 commits June 21, 2024 10:56

Apply suggestions from code review

71dae19

Co-authored-by: David Wierichs <david.wierichs@xanadu.ai>

Merge branch 'main' into from-plxpr

afa212b

bump pl version, minor test fixes

3dbd8c1

Merge branch 'main' into from-plxpr

29d58cd

black

399f622

isort

b5a1e16

linting

ac06787

" \\ z≈ç√ fd" \ `

albi3ro mentioned this pull request Jul 10, 2024

Add qml.capture.to_catalsyt conversion function PennyLaneAI/pennylane#5771

Closed

Merge branch 'main' into from-plxpr

f19890e

dime10 requested changes Jul 15, 2024

View reviewed changes

albi3ro and others added 3 commits July 24, 2024 12:41

Apply suggestions from code review

8bef814

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

Merge branch 'main' into from-plxpr

4802726

Merge branch 'main' into from-plxpr

f1f90aa

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

add from_plxpr conversion function #837

add from_plxpr conversion function #837

albi3ro commented Jun 20, 2024 •

edited

Loading

dwierichs left a comment

dwierichs Jun 21, 2024

albi3ro Jun 21, 2024 •

edited

Loading

dwierichs Jun 21, 2024

dime10 Jul 15, 2024

albi3ro Jul 29, 2024

dwierichs Jun 21, 2024

dwierichs Jun 21, 2024

codecov bot commented Jul 15, 2024 •

edited

Loading

dime10 left a comment

dime10 Jul 15, 2024

albi3ro Jul 29, 2024

dime10 Jul 15, 2024 •

edited

Loading

albi3ro Jul 29, 2024

albi3ro Jul 30, 2024

dime10 Jul 15, 2024

dime10 Jul 15, 2024

albi3ro Jul 29, 2024

dime10 commented Jul 30, 2024

		# Note that the value of rtd_kwargs is a string version of
		# the info kwargs, not the info kwargs itself

	with pytest.raises(NotImplementedError):
	with pytest.raises(NotImplementedError, match="Only wire-based and observable-based measurements"):

add from_plxpr conversion function #837

Are you sure you want to change the base?

add from_plxpr conversion function #837

Conversation

albi3ro commented Jun 20, 2024 • edited Loading

dwierichs left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

albi3ro Jun 21, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

codecov bot commented Jul 15, 2024 • edited Loading

Codecov Report

dime10 left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dime10 Jul 15, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

dime10 commented Jul 30, 2024

albi3ro commented Jun 20, 2024 •

edited

Loading

albi3ro Jun 21, 2024 •

edited

Loading

codecov bot commented Jul 15, 2024 •

edited

Loading

dime10 Jul 15, 2024 •

edited

Loading