Add base autograph transformer and if/else support #6406

lillian542 · 2024-10-16T17:58:27Z

First of 3 PRs to add support for python control flow with capture enabled by mapping the pythonic control-flow to JAX-compatible implementations in PennyLane, using autograph. The follow up PRs will add support for while loops and for loops.

This is a modification for PennyLane of the autograph module implemented in catalyst.

Context:
Without autograph, calling this circuit:

import pennylane as qml

qml.capture.enable()

dev = qml.device("default.qubit", wires=3)

@qml.qnode(dev)
def circuit(x):
    if x > 1.967:
        qml.Hadamard(2)
    else:
        qml.Y(1)
    return qml.state()

raises a TracerBoolConversionError error because the if/else logic isn't JAX-compatible.

Description of the Change:

The initial infrastructure for using autograph to convert pythonic control flow is added
The primitive for converting if/else to @qml.cond for JAX compatibility is added
The malt package is added as a dependency

Benefits:
For the circuit above, we can run

>>> circuit = run_autograph(circuit)
>>> circuit(2.3)

and now the circuit can execute.

[sc-71821]

…tives

…e into autograph_ctrl_flow

…tograph1

codecov · 2024-10-16T21:06:52Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.39%. Comparing base (661dd52) to head (d5eaa8d).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff            @@
##           master    #6406    +/-   ##
========================================
  Coverage   99.38%   99.39%            
========================================
  Files         452      455     +3     
  Lines       42789    42914   +125     
========================================
+ Hits        42527    42653   +126     
+ Misses        262      261     -1

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

pennylane/capture/autograph/ag_primitives.py

Co-authored-by: Christina Lee <[email protected]>

albi3ro · 2024-10-17T20:18:39Z

pennylane/capture/autograph/ag_primitives.py

+            def passthrough_wrapper(*args, **kwargs):
+                return converted_call(wrapped_fn, args, kwargs, caller_fn_scope, options)


Suggested change

def passthrough_wrapper(*args, **kwargs):

return converted_call(wrapped_fn, args, kwargs, caller_fn_scope, options)

def passthrough_wrapper(*inner_args, **inner_kwargs):

return converted_call(wrapped_fn, inner_args, inner_kwargs, caller_fn_scope, options)

Or some other sort of varation on args and kwargs, since those name are already taken in the scope.

albi3ro · 2024-10-17T20:19:48Z

pennylane/capture/autograph/ag_primitives.py

+            assert args and callable(args[0])
+            wrapped_fn = args[0]
+
+            def passthrough_wrapper(*args, **kwargs):


Suggested change

def passthrough_wrapper(*args, **kwargs):

@functools.wraps(wrapped_fn)

def passthrough_wrapper(*args, **kwargs):

Should we use a wrapper here too?

pennylane/capture/autograph/ag_primitives.py

albi3ro · 2024-10-17T21:15:55Z

pennylane/capture/autograph/transformer.py

+from malt.impl.api import PyToPy
+
+import pennylane as qml
+from pennylane.capture.autograph.ag_primitives import AutoGraphError


Relative import not work?

astralcai · 2024-10-18T14:07:44Z

pennylane/capture/autograph/transformer.py

+            pass
+        elif callable(obj):
+            # pylint: disable=unnecessary-lambda,unnecessary-lambda-assignment
+            fn = lambda *args, **kwargs: obj(*args, **kwargs)


What's the purpose of this assignment?

mudit2812 · 2024-10-18T14:37:35Z

pennylane/capture/autograph/ag_primitives.py

@@ -0,0 +1,178 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.


Suggested change

# Copyright 2023 Xanadu Quantum Technologies Inc.

# Copyright 2024 Xanadu Quantum Technologies Inc.

mudit2812 · 2024-10-18T14:37:41Z

pennylane/capture/autograph/transformer.py

@@ -0,0 +1,227 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.


Suggested change

# Copyright 2023 Xanadu Quantum Technologies Inc.

# Copyright 2024 Xanadu Quantum Technologies Inc.

mudit2812 · 2024-10-18T14:37:48Z

tests/capture/autograph/test_if_else.py

@@ -0,0 +1,308 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.


Suggested change

# Copyright 2023 Xanadu Quantum Technologies Inc.

# Copyright 2024 Xanadu Quantum Technologies Inc.

mudit2812 · 2024-10-18T14:37:56Z

tests/capture/autograph/test_autograph.py

@@ -0,0 +1,478 @@
+# Copyright 2023 Xanadu Quantum Technologies Inc.


Suggested change

# Copyright 2023 Xanadu Quantum Technologies Inc.

# Copyright 2024 Xanadu Quantum Technologies Inc.

mudit2812 · 2024-10-18T14:38:42Z

pennylane/capture/__init__.py

@@ -147,6 +147,7 @@ class MyCustomOp(qml.operation.Operator):
    def _(*args, **kwargs):
        return type.__call__(MyCustomOp, *args, **kwargs)
 """
+from pennylane.capture import autograph


Is this done so that we can import qml.capture.autograph? I feel like that shouldn't be necessary right?

mudit2812 · 2024-10-24T13:39:22Z

pennylane/capture/autograph/transformer.py

+    **Example**
+
+    .. code-block:: python
+
+        def decide(x):
+            if x < 5:
+                y = 15
+            else:
+                y = 1
+            return y
+
+        @qjit(autograph=True)
+        def func(x: int):
+            y = decide(x)
+            return y ** 2
+
+    >>> print(autograph_source(decide))
+    def decide_1(x):
+        with ag__.FunctionScope('decide', 'fscope', ag__.STD) as fscope:
+            def get_state():
+                return (y,)
+            def set_state(vars_):
+                nonlocal y
+                (y,) = vars_
+            def if_body():
+                nonlocal y
+                y = 15
+            def else_body():
+                nonlocal y
+                y = 1
+            y = ag__.Undefined('y')
+            ag__.if_stmt(x < 5, if_body, else_body, get_state, set_state, ('y',), 1)
+            return y
+    """


Might make sense to replace this with an example that doesn't use qjit.

albi3ro · 2024-10-24T19:48:40Z

pennylane/capture/autograph/transformer.py

+            self._cache.has(fn, TOPLEVEL_OPTIONS)
+            or self._cache.has(fn, NESTED_OPTIONS)
+            or self._cache.has(fn, STANDARD_OPTIONS)


What are TOP_LEVEL_OPTIONS, etc, and where do they come from? It doesn't look like they are imported up top 🤔

Line 210 below

andrijapau and others added 30 commits October 3, 2024 15:07

add user decompose function

5dfc38f

code factor

f1d9eee

update changelog

b24c97b

improve decomposition function

0d3bfb6

codefactor fix

0d66209

added modified generator

a25d6ba

copy code from catalyst

2ad96ac

remove catalyst from __init__ and transformer files

c471c0b

replace many catalyst dependencies/mentions

49c0fac

copy tests from catalyst

5dd1d42

update imports

c94da87

[ci skip]

ab8ae7a

Merge branch 'master' into add_autograph

ac7d53e

import autograph module in capture

4a6bdf3

remove final catalyst dependencies from ag_primitives

2c2f1dd

record queueing in FlatFn so we can access program_length in ag_primi…

9b18b97

…tives

most conditional tests work

cf92f2a

initial removal of queueing dependency

7a21764

the current state of the tests

72b3e42

remove index setting and logical operators

fee16ff

formatting

3ccd0e0

remove fallback

1ab03f3

remove disable_autograph

2bfaf72

remove autograph_include

42ce1e8

Merge branch 'master' into autograph_ctrl_flow

d3a807c

remove strict_conversion and ignore_fallbacks

fa0b415

Tidying up

54b3036

Tidying up more

26cbde4

fix more tests

211bd81

Merge branch 'autograph_ctrl_flow' of github.com:PennyLaneAI/pennylan…

6be1ff5

…e into autograph_ctrl_flow

lillian542 and others added 2 commits October 16, 2024 16:11

fix failing test

da0cd1d

Merge branch 'master' into autograph1

cea27f6

lillian542 marked this pull request as ready for review October 16, 2024 20:12

lillian542 added 4 commits October 16, 2024 16:22

package name is diastatic-malt

9da4afa

Merge branch 'autograph1' of github.com:PennyLaneAI/pennylane into au…

fbff208

…tograph1

fix decorator test

e443173

rename test file to avoid CI confusion

2ee18ca

lillian542 and others added 3 commits October 17, 2024 14:00

a few more tests and docstrings updates

fc59744

one more test for code coverage

a8946af

Merge branch 'master' into autograph1

908d6b1

lillian542 requested review from albi3ro and astralcai October 17, 2024 18:16

albi3ro reviewed Oct 17, 2024

View reviewed changes

pennylane/capture/autograph/ag_primitives.py Outdated Show resolved Hide resolved

Update pennylane/capture/autograph/ag_primitives.py

60cb38b

Co-authored-by: Christina Lee <[email protected]>

albi3ro reviewed Oct 17, 2024

View reviewed changes

lillian542 mentioned this pull request Oct 17, 2024

Add support for while loop with autograph #6413

Open

lillian542 requested a review from mudit2812 October 17, 2024 21:10

albi3ro reviewed Oct 17, 2024

View reviewed changes

pennylane/capture/autograph/ag_primitives.py Show resolved Hide resolved

albi3ro reviewed Oct 17, 2024

View reviewed changes

astralcai reviewed Oct 18, 2024

View reviewed changes

mudit2812 reviewed Oct 18, 2024

View reviewed changes

mudit2812 reviewed Oct 24, 2024

View reviewed changes

albi3ro reviewed Oct 24, 2024

View reviewed changes

Merge branch 'master' into autograph1

d5eaa8d

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add base autograph transformer and if/else support #6406

Add base autograph transformer and if/else support #6406

lillian542 commented Oct 16, 2024 •

edited

Loading

codecov bot commented Oct 16, 2024 •

edited

Loading

albi3ro Oct 17, 2024

albi3ro Oct 17, 2024

albi3ro Oct 17, 2024

astralcai Oct 18, 2024

mudit2812 Oct 18, 2024

mudit2812 Oct 18, 2024

mudit2812 Oct 18, 2024

mudit2812 Oct 18, 2024

mudit2812 Oct 18, 2024

mudit2812 Oct 24, 2024

albi3ro Oct 24, 2024

lillian542 Oct 28, 2024

		def passthrough_wrapper(args, *kwargs):
		return converted_call(wrapped_fn, args, kwargs, caller_fn_scope, options)

	def passthrough_wrapper(args, *kwargs):
	@functools.wraps(wrapped_fn)
	def passthrough_wrapper(args, *kwargs):

		@@ -0,0 +1,178 @@
		# Copyright 2023 Xanadu Quantum Technologies Inc.

	# Copyright 2023 Xanadu Quantum Technologies Inc.
	# Copyright 2024 Xanadu Quantum Technologies Inc.

		@@ -0,0 +1,227 @@
		# Copyright 2023 Xanadu Quantum Technologies Inc.

		@@ -0,0 +1,308 @@
		# Copyright 2023 Xanadu Quantum Technologies Inc.

		@@ -0,0 +1,478 @@
		# Copyright 2023 Xanadu Quantum Technologies Inc.

Add base autograph transformer and if/else support #6406

Are you sure you want to change the base?

Add base autograph transformer and if/else support #6406

Conversation

lillian542 commented Oct 16, 2024 • edited Loading

codecov bot commented Oct 16, 2024 • edited Loading

Codecov Report

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

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

lillian542 commented Oct 16, 2024 •

edited

Loading

codecov bot commented Oct 16, 2024 •

edited

Loading