[OpRefactor] Add custom decomposition context manager to device #1900

glassnotes · 2021-11-16T15:49:03Z

Context: Operator overhaul

Description of the Change: Adds custom decompositions at the device level by creating a custom_expand_fn and context manager. Also adds a context manager to enable custom decompositions to be temporarily used by devices. Subsumes PR #1872.

Example: suppose we are running on a trapped-ion machine, and would like to implement a decomposition of the CNOT gate into IsingXX. We can do so by defining a decomposition (which has the same signature as that of the normal operation), and passing the device a dictionary detailing the custom decompositions.

def ion_trap_cnot(wires):
    return [
        qml.RY(np.pi/2, wires=wires[0]),
        qml.IsingXX(np.pi/2, wires=wires),
        qml.RX(-np.pi/2, wires=wires[0]),
        qml.RY(-np.pi/2, wires=wires[0]),
        qml.RY(-np.pi/2, wires=wires[1])
    ]

dev = qml.device('default.qubit', wires=2, custom_decomps={"CNOT" : ion_trap_cnot})

Now, we can see the decomposition get applied by drawing a QNode using the device expansion strategy.

>>> @qml.qnode(dev)
... def run_cnot():
...     qml.CNOT(wires=[0, 1])
...     return qml.expval(qml.PauliX(wires=1))
>>> print(qml.draw(run_cnot, expansion_strategy="device")())
 0: ──RY(1.57)──╭IsingXX(1.57)──RX(-1.57)──RY(-1.57)──┤     
 1: ────────────╰IsingXX(1.57)──RY(-1.57)─────────────┤ ⟨X⟩

Benefits: Provides an intuitive means of implementing an arbitrary number of custom decompositions at runtime. There is minimal modification to the code of the device, and no modification of any operator code, since everything is done via context managers.

Possible Drawbacks: Currently, this works by over-riding the decompose method of the specified operations within a given context. However this may not work for templates, which implement expand.

Related GitHub Issues: N/A

glassnotes · 2021-11-16T15:50:42Z

pennylane/_qubit_device.py

+            # within the context.
+            def custom_decomp_expand(self, circuit, max_expansion=10):
+                with self.custom_decomp_context():
+                    return self.custom_fn(circuit, max_expansion)


The with statement does seem to alter the decomposition; if you do something like print(qml.CNOT.decomposition(wires=[0, 1]) with the new decomposition, it will show up correctly as the custom function, however when the custom expand function is later called, the decomposition is not overridden.

codecov · 2021-11-16T16:02:07Z

Codecov Report

Merging #1900 (e5769b8) into master (d4c98a1) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #1900   +/-   ##
=======================================
  Coverage   98.83%   98.83%           
=======================================
  Files         222      222           
  Lines       16989    17036   +47     
=======================================
+ Hits        16791    16838   +47     
  Misses        198      198

Impacted Files	Coverage Δ
pennylane/transforms/__init__.py	`100.00% <ø> (ø)`
pennylane/__init__.py	`100.00% <100.00%> (ø)`
pennylane/transforms/tape_expand.py	`100.00% <100.00%> (ø)`

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update d4c98a1...e5769b8. Read the comment docs.

glassnotes · 2021-11-16T18:25:26Z

pennylane/_qubit_device.py

+
+            original_decomp_method = obj.decompose
+
+            # This is the method that will override the current .decompose method of


The reason why we have to override decompose instead of decomposition here is because in the default operator expand function, this is what is called to perform the decomposition (see here).

nice! this is more general as well

The one disadvantage I can see is that, if a user wanted to define a custom decomposition for a template, this wouldn't work, because many of them don't have a decomposition, just expand. But this is maybe not so much the intended use case.

this will be changed shortly though, right? (@mariaschuld?)

josh146

the prototype seems to work really well 🎉

josh146 · 2021-11-17T12:20:59Z

doc/releases/changelog-dev.md

+  ```pycon
+  >>> print(qml.draw(decomp_qnode, expansion_strategy="device")(weights))
+   0: ──RX(0.4)──────────────────────╭C──RZ(3.14)──RY(1.57)──────────────────────────╭Z──RZ(3.14)──RY(1.57)──┤ ⟨Z⟩
+   1: ──RX(0.5)──RZ(3.14)──RY(1.57)──╰Z──RZ(3.14)──RY(1.57)──╭C──────────────────────│───────────────────────┤
+   2: ──RX(0.6)──RZ(3.14)──RY(1.57)──────────────────────────╰Z──RZ(3.14)──RY(1.57)──╰C──────────────────────┤
+  ```


I spent a long time looking at this, but it works out the recursive relationship between the custom CNOT and the custom Hadamard really well! 🎉

That's so wonderful @glassnotes, it feels right to throw this into the device!

josh146 · 2021-11-17T12:26:52Z

pennylane/_qubit_device.py

+            custom_decomp_condition =  qml.BooleanFn(
+                lambda obj: not isinstance(obj, qml.tape.QuantumTape) # Expand templates
+                and obj.name not in custom_op_names # Expand things that don't have custom decomp
+                and self.supports_operation(obj.name) # Expand things until supported on device
+            )


@glassnotes did

custom_decomp_condition = qml.BooleanFn(lambda obj: obj.name not in custom_op_names) custom_decomp_condition &= self.stopping_condition

not work?

this is the same as the following:

condition = qml.BooleanFn(lambda obj: obj.name not in custom_op_names) custom_fn = qml.transforms.create_expand_fn(depth=10, stop_at=condition, device=self)

Actually I had a weird amount of trouble with the stopping condition. It seemed like if I did not include that first condition there, it would not expand the templates 😕 Let me give this another go, though.

oh, the first one is included within the devices stopping condition though 🤔

Ahhh I found the problem I had originally. I was including both obj.name not in custom_op_names and self.supports_operation(obj.name) in the stopping condition I was passing to the create_expand_fn, and the second one was being cancelled out in the &= with the device criteria. Thanks!

josh146 · 2021-11-17T12:31:34Z

pennylane/_qubit_device.py

+
+            original_decomp_method = obj.decompose
+
+            # This is the method that will override the current .decompose method of


nice! this is more general as well

glassnotes · 2021-11-17T19:39:45Z

tests/transforms/test_tape_expand.py

+
+        assert len(decomp_ops) == 2
+
+        assert decomp_ops[0].name == "CRZ"


Honestly I was kind of surprised this worked out-of-the-box. It seems like the operation being controlled gets decomposed first, and then each gate in the decomposition gets the control applied.

@glassnotes I'm trying to think of any other tests cases, but it sounds like you've covered pretty much all! Only other thing I can think of is a decomposition where:

template A is decomposed into template B

gate C is decomposed into gate D

and then decomposing a circuit that has template A?

I also have to add gradient tests, which I just totally forgot about 😅

sure! Although my gut feeling is that this is one of the parts of the codebase where gradient tests aren't needed, since this PR doesn't introduce any classical processing.

mariaschuld · 2021-11-18T07:30:16Z

pennylane/_qubit_device.py

+        custom_decomps (Dict[Union(str, qml.operation.Operation), Callable]): Custom
+            decompositions to be applied by the device at runtime. When specified, the
+            device will create a custom expand function by combining the regular expansion
+            criteria with those specified by the custom decompositions.


Sorry to be a pain, but could we add an example to the docstring? Although the description is accurate, it is so much quicker to see it!

Actually, it might make more sense to add the example to the device loader, since that is the more user-facing docstring? E.g., here: https://pennylane.readthedocs.io/en/stable/code/api/pennylane.device.html

Oh, looks like we need to update that docstring anyway to match our latest conventions!

I meant to ask about this actually, I wasn't sure where it made sense to add this to the docs! Sure, I'll put it on qml.device page. @josh146 what do you mean by "update that docstring anyway to match our latest conventions"?

Oh, just following https://pennylane.readthedocs.io/en/latest/development/guide/documentation.html#functions-and-methods, e.g., having the example after the arguments etc.

mariaschuld · 2021-11-18T07:31:13Z

pennylane/_qubit_device.py

@@ -121,7 +125,7 @@ def _scatter(indices, array, new_dimensions):
        "Projector",
    }

-    def __init__(self, wires=1, shots=None, cache=0, analytic=None):
+    def __init__(self, wires=1, shots=None, cache=0, analytic=None, custom_decomps=None):


@mlxd just thought I'd tag you because we are adding stuff to the QubitDevice class that may have to be ported to Device one day

Great, thanks for the heads-up @mariaschuld !

mariaschuld · 2021-11-18T07:32:51Z

pennylane/devices/default_qubit.py

+    def __init__(self, wires, *, shots=None, cache=0, analytic=None, custom_decomps=None):
+        super().__init__(
+            wires, shots, cache=cache, analytic=analytic, custom_decomps=custom_decomps
+        )


If custom_decomps is used with external devices (that do not pass it to super), will nothing happen? Is this a problem?

Ah, this is a very good point. I changed the changelog example to use lightning.qubit, and got

TypeError: __init__() got an unexpected keyword argument 'custom_decomps'

However, this works fine:

dev = qml.device("lightning.qubit", wires=3) custom_decomps = {qml.CNOT : custom_cnot, "Hadamard" : custom_hadamard} custom_decomps = qml.transforms.create_custom_decomp_expand_fn(custom_decomps, dev) dev.custom_expand(custom_decomps) @qml.qnode(dev) def circuit(weights): qml.BasicEntanglerLayers(weights, wires=[0, 1, 2]) return qml.expval(qml.PauliX(0)) print(qml.draw(circuit, expansion_strategy="device")(weights))

@glassnotes you could create a utility function that combines the calls to create_custom_decomp_expand_fn and dev.custom_expand:

@contextlib.contextmanager def decomposition_context(dev, decomps: dict = None): expand_fn = create_custom_decomp_expand_fn(decomps, dev) try: dev.custom_expand(expand_fn) yield finally: dev.custom_expand_fn = None

Since dev.custom_expand will globally change the state of the device, I am always concerned about having that be user-facing without wrapping it in a context manager (which provides a context where the change of state is explicit).

This then simplifies the above down to

dev = qml.device("lightning.qubit", wires=3) @qml.qnode(dev) def circuit(weights): qml.BasicEntanglerLayers(weights, wires=[0, 1, 2]) return qml.expval(qml.PauliX(0)) custom_decomps = {qml.CNOT : custom_cnot, "Hadamard" : custom_hadamard} with decomposition_context(dev, custom_decomps): print(qml.draw(circuit, expansion_strategy="device")(weights))

which will work on every device without modification. What do you think?

I must admit I am very much in favour of the latter approach, since it feels more maintainable. Changes to the decomposition API do not require changes to the device API, and vice versa.

It also provides a place in the documentation for documentation and examples.

Ooh this gives me an idea :) In the device loader, you could do

def device(name, *args, **kwargs): custom_decomp = kwargs.pop("custom_decomp", None) dev = ... # load the device if custom_decomp is not None: custom_decomp_expand_fn = qml.transforms.create_custom_decomp_expand_fn( custom_decomps, dev ) dev.custom_expand(custom_decomp_expand_fn) ...

and the subclasses never have to worry :)

with decomposition_context(dev, custom_decomps):
print(qml.draw(circuit, expansion_strategy="device")(weights))

This is like what was implemented in #1872 .

Since dev.custom_expand will globally change the state of the device, I am always concerned about having that be user-facing without wrapping it in a context manager (which provides a context where the change of state is explicit).

In this case though, if the argument is passed upon creation of the device, I would expect the behaviour to be global, and that any time that device is used, the custom decomposition is used.

I do really like the idea of putting this in the device loader though; way more flexible, plus we were going to put the docs there anyways.

mariaschuld · 2021-11-18T07:33:19Z

pennylane/transforms/tape_expand.py

@@ -25,6 +27,11 @@
    not_tape,
 )

+# Needed for custom decomposition context manager
+from pennylane.transforms.qfunc_transforms import NonQueuingTape


I only learnt about this recently :)

The NonQueuingTape is something we should definitely move out of the transforms folder! we should put it in tape/ lol

pennylane/transforms/tape_expand.py

josh146

this feels so natural and intuitive @glassnotes!

I still have to think about this, because my hope here was to eliminate the need for the user to work with the contexts at all, and have everything handled by the device (well, now, the device loader).

Yeah, I guess I just wanted some secret advanced context manager somewhere, just in case I need to help someone with code and an existing device that has already been instantiated/used 😆

the following always works:

expand_fn = qml.transforms.create_decomp_expand_fn(custom_decomps, dev)
dev.custom_expand(expand_fn)

but I was thinking that

with set_decompositions(custom_decomp, dev):
    # existing models and executions

could be really nice to whip out in advanced situations/debugging, and a lot safer than calling dev.custom_expand(), which will cause global state issues

josh146 · 2021-11-18T16:29:23Z

pennylane/__init__.py

+        dev = qml.device(
+            'default.qubit', wires=2, custom_decomps={"CNOT" : ion_trap_cnot}
+        )


this feels very natural and intuitive 🙌

pennylane/__init__.py

pennylane/transforms/tape_expand.py

josh146 · 2021-11-18T16:34:37Z

tests/transforms/test_tape_expand.py

+    def test_custom_decomp_template_to_template(self):
+        """Test that decomposing a template into another template and some
+        gates yields the correct results."""


Co-authored-by: Josh Izaac <josh146@gmail.com>

glassnotes · 2021-11-18T17:46:33Z

Yeah, I guess I just wanted some secret advanced context manager somewhere, just in case I need to help someone with code and an existing device that has already been instantiated/used laughing

Ohhhh I see. I misunderstood, I interpreted this as being the way a user would implement this, and not just an alternative method. This is a very easy utility function to add :)

pennylane/transforms/tape_expand.py

mariaschuld · 2021-11-19T06:36:41Z

pennylane/transforms/tape_expand.py

+                if self.num_params == 0:
+                    return fn(self.wires)
+                return fn(*self.parameters, self.wires)
+            return tape


I am momentarily confused, will this line ever be used? Since there is a return function above without an if fork? Maybe this is why the coverage complains?

Whoa!! No, it won't be, removing it has absolutely no effect. I swear, when I checked the coverage yesterday everything was passing 😕

mariaschuld

Sorry, just one tiny last thing - the one line that is not covered by tests, would you be able to add a test? I'm particularly asking because I am unsure if the line is ever used?

…b.com/PennyLaneAI/pennylane into custom_decomposition_device_expansion

glassnotes · 2021-11-19T14:31:44Z

[sc-7225]

Try setting custom expand function with context manager.

dfdd2f1

glassnotes commented Nov 16, 2021

View reviewed changes

Olivia Di Matteo added 7 commits November 16, 2021 11:15

Untangle some of the Booleans.

79f4acb

Working version with decompose.

38abe2d

Get it to work with templates.

c5f6701

Update CHANGELOG.

1293802

Fix typo.

5644463

Fix spacing.

577878d

Fix some codefactor issues.

99092a1

glassnotes commented Nov 16, 2021

View reviewed changes

josh146 reviewed Nov 17, 2021

View reviewed changes

Olivia Di Matteo and others added 7 commits November 17, 2021 08:50

Merge branch 'master' into custom_decomposition_device_expansion

507e610

Fix stopping condition and run black.

2ddd724

Merge branch 'master' into custom_decomposition_device_expansion

5ae8c62

Move context manager into file with expansion functions.

2e2239e

Remove beta from example in CHANGELOG.

e3b10e6

Run black and appease pylint.

0b8985a

Merge branch 'master' into custom_decomposition_device_expansion

bd52d19

glassnotes marked this pull request as ready for review November 17, 2021 18:26

glassnotes mentioned this pull request Nov 17, 2021

[WIP] Context manager for custom decompositions #1872

Closed

5 tasks

Add tests; fix bug in decomp of parametrized gates.

19d1366

glassnotes commented Nov 17, 2021

View reviewed changes

glassnotes requested review from josh146 and mariaschuld November 17, 2021 19:40

mariaschuld reviewed Nov 18, 2021

View reviewed changes

pennylane/transforms/tape_expand.py Show resolved Hide resolved

josh146 approved these changes Nov 18, 2021

View reviewed changes

glassnotes and others added 3 commits November 18, 2021 12:22

Merge branch 'master' into custom_decomposition_device_expansion

003a266

Update pennylane/__init__.py

4154b04

Co-authored-by: Josh Izaac <josh146@gmail.com>

Add example to create_decomp_expand_fn

8005f03

Olivia Di Matteo added 3 commits November 18, 2021 12:55

Add new context manager for temporary decomposition setting

a4e7aec

Add test of context manager.

1a1d7cc

Run black.

fa13a2f

glassnotes commented Nov 18, 2021

View reviewed changes

pennylane/transforms/tape_expand.py Outdated Show resolved Hide resolved

glassnotes commented Nov 18, 2021

View reviewed changes

pennylane/transforms/tape_expand.py Outdated Show resolved Hide resolved

glassnotes and others added 3 commits November 18, 2021 13:02

Apply suggestions from code review

bff8e48

Run black.

8a4cd64

Remove return statement.

831e7ec

glassnotes requested a review from mariaschuld November 18, 2021 18:25

Update CHANGELOG to add description of new context manager.

201218f

mariaschuld reviewed Nov 19, 2021

View reviewed changes

mariaschuld requested changes Nov 19, 2021

View reviewed changes

mariaschuld and others added 4 commits November 19, 2021 11:45

Merge branch 'master' into custom_decomposition_device_expansion

dc01d2a

Add new argument to specify depth.

71016bc

Merge branch 'custom_decomposition_device_expansion' of https://githu…

ed24c2c

…b.com/PennyLaneAI/pennylane into custom_decomposition_device_expansion

Fix codefactor issues.

f61b82a

glassnotes requested a review from mariaschuld November 19, 2021 14:08

Merge branch 'master' into custom_decomposition_device_expansion

e9cdab5

Merge branch 'master' into custom_decomposition_device_expansion

a399eb9

mariaschuld approved these changes Nov 22, 2021

View reviewed changes

mariaschuld and others added 2 commits November 22, 2021 06:59

Merge branch 'master' into custom_decomposition_device_expansion

d072957

Merge branch 'master' into custom_decomposition_device_expansion

e5769b8

glassnotes merged commit 854d4b4 into master Nov 22, 2021

glassnotes deleted the custom_decomposition_device_expansion branch November 22, 2021 13:22


		original_decomp_method = obj.decompose

		# This is the method that will override the current .decompose method of


		assert len(decomp_ops) == 2

		assert decomp_ops[0].name == "CRZ"

[OpRefactor] Add custom decomposition context manager to device #1900

[OpRefactor] Add custom decomposition context manager to device #1900

Conversation

glassnotes commented Nov 16, 2021 • edited Loading

Choose a reason for hiding this comment

codecov bot commented Nov 16, 2021 • 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

josh146 left a 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

glassnotes Nov 17, 2021 • 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

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

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

josh146 left a comment • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

glassnotes commented Nov 18, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

mariaschuld left a comment

Choose a reason for hiding this comment

glassnotes commented Nov 19, 2021

glassnotes commented Nov 16, 2021 •

edited

Loading

codecov bot commented Nov 16, 2021 •

edited

Loading

glassnotes Nov 17, 2021 •

edited

Loading

josh146 left a comment •

edited

Loading