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] First PR for enabling dynamic decompositions with PLxPR enabled #6859

Merged
merged 48 commits into from
Feb 10, 2025
Merged
Changes from 1 commit
Commits
Show all changes
48 commits
Select commit Hold shift + click to select a range
696acf3
E.C.
PietropaoloFrisoni Jan 20, 2025
9cdfae9
Creating an empty `DynamicDecomposeInterpreter` c;ass
PietropaoloFrisoni Jan 21, 2025
1e138fa
Sbattendo la testa contro il muro tante volte
PietropaoloFrisoni Jan 21, 2025
76c9250
Current prototype version
PietropaoloFrisoni Jan 22, 2025
c821ac9
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 22, 2025
cee6ec4
Fixing one more problem
PietropaoloFrisoni Jan 22, 2025
a208dc5
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 22, 2025
d16a9e0
Moving tests to separate file
PietropaoloFrisoni Jan 23, 2025
a8b9283
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 23, 2025
18bc43c
Pylint fixes (although premature)
PietropaoloFrisoni Jan 23, 2025
e2e8fd0
Removing reundandt tuple calls
PietropaoloFrisoni Jan 23, 2025
0abd620
Tests with dynamic wires
PietropaoloFrisoni Jan 23, 2025
1e3ffb6
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 23, 2025
1ae399b
Adding Autograph test
PietropaoloFrisoni Jan 23, 2025
9a54c3e
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 24, 2025
c5f2ae5
Removing unused parameters and adding a few tests
PietropaoloFrisoni Jan 24, 2025
497440c
Adding a few more tests
PietropaoloFrisoni Jan 24, 2025
c7da133
Removing import
PietropaoloFrisoni Jan 24, 2025
2f0417c
Pylint
PietropaoloFrisoni Jan 24, 2025
3c8bc37
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 24, 2025
e9ff110
Adding test with hyperparameters
PietropaoloFrisoni Jan 27, 2025
b9f5d03
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 27, 2025
f2437b4
Black
PietropaoloFrisoni Jan 27, 2025
8c615c5
A few more tests
PietropaoloFrisoni Jan 27, 2025
9fef95b
Changelog
PietropaoloFrisoni Jan 27, 2025
4a56150
Removing redundant operations
PietropaoloFrisoni Jan 28, 2025
e733762
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 29, 2025
9792df6
Pre-binding hyperparameters [ci skip]
PietropaoloFrisoni Jan 29, 2025
aa422b3
Removing redundant method
PietropaoloFrisoni Jan 30, 2025
b7a18cd
Pylint
PietropaoloFrisoni Jan 30, 2025
97cba03
Testing CI failures (JAX imports)
PietropaoloFrisoni Jan 30, 2025
e05963e
isort
PietropaoloFrisoni Jan 30, 2025
ec73757
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 30, 2025
bc594ca
Support for consts and hyperparameters
PietropaoloFrisoni Jan 31, 2025
4734388
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Jan 31, 2025
1ea4a14
Fixes neede after autograph PR merged on master
PietropaoloFrisoni Jan 31, 2025
f65ab4e
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Feb 5, 2025
f576486
[ci skip]
PietropaoloFrisoni Feb 5, 2025
fd817cc
Removing `DynamicDecomposeInterpreter`
PietropaoloFrisoni Feb 5, 2025
90fc1e5
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Feb 5, 2025
e7a2775
pylint
PietropaoloFrisoni Feb 5, 2025
f621156
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Feb 10, 2025
e9834af
Suggestions from code review (more tests)
PietropaoloFrisoni Feb 10, 2025
bd9473b
disabling wrong iimport order in test file (conflict between isort an…
PietropaoloFrisoni Feb 10, 2025
43a6838
Merge branch 'master' of https://github.com/PennyLaneAI/pennylane int…
PietropaoloFrisoni Feb 10, 2025
98fe17e
Suggestions from code review
PietropaoloFrisoni Feb 10, 2025
902e4f8
Refactoring changelog with program capture entries
PietropaoloFrisoni Feb 10, 2025
4f95b74
Merge branch 'master' into cond_dynamic_decomp
PietropaoloFrisoni Feb 10, 2025
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Adding a few more tests
PietropaoloFrisoni committed Jan 24, 2025
commit 497440c14129f0b17dcf83667c20b2522e4af0fc
46 changes: 23 additions & 23 deletions tests/capture/transforms/test_capture_dynamic_decompositions.py
Original file line number Diff line number Diff line change
@@ -37,7 +37,7 @@ class SimpleCustomOp(Operation):
"""Simple custom operation that contains a single gate in its decomposition"""

num_wires = 1
num_params = 1
num_params = 0

def _init__(self, phi, wires, id=None):
super().__init__(phi, wires=wires, id=id)
@@ -49,8 +49,8 @@ def _plxpr_decomposition(self) -> "jax.core.Jaxpr":
)

@staticmethod
def _compute_plxpr_decomposition(phi, wires):
qml.RX(phi, wires=wires)
def _compute_plxpr_decomposition(wires):
qml.Hadamard(wires=wires)


class CustomOpMultiWire(Operation):
@@ -70,9 +70,9 @@ def _plxpr_decomposition(self) -> "jax.core.Jaxpr":

@staticmethod
def _compute_plxpr_decomposition(phi, *wires):
qml.CNOT([wires[0], wires[1]])
qml.DoubleExcitation(phi, wires)
qml.CNOT([wires[0], wires[1]])
qml.CNOT(wires=[wires[0], wires[1]])
qml.DoubleExcitation(phi, wires=wires)
qml.CNOT(wires=[wires[0], wires[1]])
qml.RX(phi, wires=wires[0])
qml.RY(phi, wires=wires[1])
qml.RZ(phi, wires=wires[2])
@@ -206,51 +206,50 @@ def test_function_simple(self):
"""Test that a function with a custom operation is correctly decomposed."""

@DynamicDecomposeInterpreter()
def f(x):
qml.RY(x, wires=0)
SimpleCustomOp(x, wires=0)
def f():
qml.RY(0.1, wires=0)
SimpleCustomOp(wires=0)
return qml.expval(qml.Z(0))

jaxpr = jax.make_jaxpr(f)(0.5)
jaxpr = jax.make_jaxpr(f)()
assert len(jaxpr.eqns) == 4
assert jaxpr.eqns[0].primitive == qml.RY._primitive
assert jaxpr.eqns[1].primitive == qml.RX._primitive
assert jaxpr.eqns[1].primitive == qml.Hadamard._primitive
assert jaxpr.eqns[2].primitive == qml.PauliZ._primitive
assert jaxpr.eqns[3].primitive == qml.measurements.ExpectationMP._obs_primitive

############################
### QNode tests
############################

@pytest.mark.parametrize("x", [0.2, 0.8])
def test_qnode_simple(self, x):
def test_qnode_simple(self):
"""Test that a QNode with a custom operation is correctly decomposed."""

@DynamicDecomposeInterpreter()
@qml.qnode(device=qml.device("default.qubit", wires=2))
def circuit(x):
qml.RY(x, wires=0)
SimpleCustomOp(x, wires=0)
def circuit():
qml.RY(0.1, wires=0)
SimpleCustomOp(wires=0)
return qml.expval(qml.Z(0))

jaxpr = jax.make_jaxpr(circuit)(x)
jaxpr = jax.make_jaxpr(circuit)()

assert jaxpr.eqns[0].primitive == qnode_prim
qfunc_jaxpr = jaxpr.eqns[0].params["qfunc_jaxpr"]
assert qfunc_jaxpr.eqns[0].primitive == qml.RY._primitive
assert qfunc_jaxpr.eqns[1].primitive == qml.RX._primitive
assert qfunc_jaxpr.eqns[1].primitive == qml.Hadamard._primitive
assert qfunc_jaxpr.eqns[2].primitive == qml.PauliZ._primitive
assert qfunc_jaxpr.eqns[3].primitive == qml.measurements.ExpectationMP._obs_primitive

result = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts, x)
result = jax.core.eval_jaxpr(jaxpr.jaxpr, jaxpr.consts)

@qml.qnode(device=qml.device("default.qubit", wires=2))
def circuit_comparison(x):
qml.RY(x, wires=0)
qml.RX(x, wires=0)
def circuit_comparison():
qml.RY(0.1, wires=0)
qml.Hadamard(wires=0)
return qml.expval(qml.Z(0))

assert jax.numpy.allclose(*result, circuit_comparison(x))
assert jax.numpy.allclose(*result, circuit_comparison())

@pytest.mark.parametrize("wires", [[0, 1, 2, 3], [2, 3, 1, 0]])
def test_multi_wire(self, wires):
@@ -436,6 +435,7 @@ def circuit_comparison(x, wire):

return qml.expval(qml.Z(0))

# Autograph requires to capture the function first
jaxpr_comparison = qml.capture.make_plxpr(circuit_comparison)(x, wire)
result_comparison = jax.core.eval_jaxpr(
jaxpr_comparison.jaxpr, jaxpr_comparison.consts, x, wire