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

Autograph with make plxpr #6645

Merged
merged 105 commits into from
Dec 5, 2024
Merged
Changes from 102 commits
Commits
Show all changes
105 commits
Select commit Hold shift + click to select a range
5dfc38f
add user decompose function
andrijapau Oct 3, 2024
f1d9eee
code factor
andrijapau Oct 3, 2024
b24c97b
update changelog
andrijapau Oct 3, 2024
0d3bfb6
improve decomposition function
andrijapau Oct 4, 2024
0d66209
codefactor fix
andrijapau Oct 4, 2024
a25d6ba
added modified generator
andrijapau Oct 4, 2024
2ad96ac
copy code from catalyst
lillian542 Oct 4, 2024
c471c0b
remove catalyst from __init__ and transformer files
lillian542 Oct 4, 2024
49c0fac
replace many catalyst dependencies/mentions
lillian542 Oct 7, 2024
5dd1d42
copy tests from catalyst
lillian542 Oct 7, 2024
c94da87
update imports
lillian542 Oct 7, 2024
ab8ae7a
[ci skip]
lillian542 Oct 8, 2024
ac7d53e
Merge branch 'master' into add_autograph
lillian542 Oct 8, 2024
4a6bdf3
import autograph module in capture
lillian542 Oct 8, 2024
2c2f1dd
remove final catalyst dependencies from ag_primitives
lillian542 Oct 8, 2024
9b18b97
record queueing in FlatFn so we can access program_length in ag_primi…
lillian542 Oct 8, 2024
cf92f2a
most conditional tests work
lillian542 Oct 9, 2024
7a21764
initial removal of queueing dependency
lillian542 Oct 10, 2024
72b3e42
the current state of the tests
lillian542 Oct 10, 2024
fee16ff
remove index setting and logical operators
lillian542 Oct 11, 2024
3ccd0e0
formatting
lillian542 Oct 11, 2024
1ab03f3
remove fallback
lillian542 Oct 12, 2024
2bfaf72
remove disable_autograph
lillian542 Oct 12, 2024
42ce1e8
remove autograph_include
lillian542 Oct 12, 2024
d3a807c
Merge branch 'master' into autograph_ctrl_flow
lillian542 Oct 12, 2024
fa0b415
remove strict_conversion and ignore_fallbacks
lillian542 Oct 12, 2024
54b3036
Tidying up
lillian542 Oct 12, 2024
26cbde4
Tidying up more
lillian542 Oct 12, 2024
211bd81
fix more tests
lillian542 Oct 16, 2024
6be1ff5
Merge branch 'autograph_ctrl_flow' of github.com:PennyLaneAI/pennylan…
lillian542 Oct 16, 2024
75b0c3c
only cond in ag_primitives
lillian542 Oct 16, 2024
5bfecb4
clean up transformer and utils
lillian542 Oct 16, 2024
ac51fad
re-organize tests
lillian542 Oct 16, 2024
d4d7220
remove utils and clean up docstrings
lillian542 Oct 16, 2024
207ecc6
update changelog
lillian542 Oct 16, 2024
d2f6e34
add malt as dependency of PL
lillian542 Oct 16, 2024
da0cd1d
fix failing test
lillian542 Oct 16, 2024
cea27f6
Merge branch 'master' into autograph1
lillian542 Oct 16, 2024
9da4afa
package name is diastatic-malt
lillian542 Oct 16, 2024
fbff208
Merge branch 'autograph1' of github.com:PennyLaneAI/pennylane into au…
lillian542 Oct 16, 2024
e443173
fix decorator test
lillian542 Oct 16, 2024
2ee18ca
rename test file to avoid CI confusion
lillian542 Oct 16, 2024
3ab7c9c
add while_loop implementation
lillian542 Oct 16, 2024
900cfe1
add test file
lillian542 Oct 17, 2024
fc59744
a few more tests and docstrings updates
lillian542 Oct 17, 2024
a8946af
one more test for code coverage
lillian542 Oct 17, 2024
908d6b1
Merge branch 'master' into autograph1
lillian542 Oct 17, 2024
60cb38b
Update pennylane/capture/autograph/ag_primitives.py
lillian542 Oct 17, 2024
9c8c74d
add initial tests
lillian542 Oct 17, 2024
bfc39b2
Merge branch 'autograph1' into autograph_while_loop
lillian542 Oct 17, 2024
3972b89
some small test changes
lillian542 Oct 17, 2024
ef2d503
update changelog
lillian542 Oct 21, 2024
19130e9
xfail test that includes for loop
lillian542 Oct 21, 2024
4b73551
add for loop support and tests
lillian542 Oct 21, 2024
d5eaa8d
Merge branch 'master' into autograph1
lillian542 Oct 28, 2024
7da5662
Merge branch 'autograph1' into autograph_while_loop
lillian542 Oct 28, 2024
a03acf8
Merge branch 'master' into autograph1
lillian542 Nov 21, 2024
1a01dea
Apply suggestions from code review
lillian542 Nov 21, 2024
42c91f9
use inner_args to avoid taken arguments
lillian542 Nov 21, 2024
78307b5
update copyright year
lillian542 Nov 21, 2024
e36fa65
use functools.wraps
lillian542 Nov 21, 2024
998ce1a
replace qjit example with pl example
lillian542 Nov 21, 2024
83a934e
add import path for run_autograph, autograph_source
lillian542 Nov 21, 2024
36267d8
change import structure + update example
lillian542 Nov 21, 2024
c06e768
fix a couple docstring mistakes
lillian542 Nov 21, 2024
d6fae5d
Merge branch 'autograph1' into autograph_while_loop
lillian542 Nov 21, 2024
a63009d
Apply suggestions from code review
lillian542 Nov 21, 2024
793e7f4
remove unneeded check
lillian542 Nov 26, 2024
1f1f097
small test fixes
lillian542 Nov 26, 2024
705eb35
Merge branch 'autograph_while_loop' of github.com:PennyLaneAI/pennyla…
lillian542 Nov 26, 2024
c387988
Merge branch 'master' into autograph1
lillian542 Nov 26, 2024
84ca81b
Merge branch 'autograph1' into autograph_while_loop
lillian542 Nov 26, 2024
e8f23ab
Merge branch 'autograph_while_loop' into autograph_for_loop
lillian542 Nov 26, 2024
ed9ea11
remove source_info function
lillian542 Nov 26, 2024
0faba0e
reoraganize and update tests
lillian542 Nov 26, 2024
4c5d969
clean up error msgs and docstrings
lillian542 Nov 26, 2024
b981ade
update tests
lillian542 Nov 26, 2024
6432807
update changelog
lillian542 Nov 26, 2024
3d6d40e
pylint complaint
lillian542 Nov 26, 2024
9403d0d
add autograph to make_plxpr
lillian542 Nov 27, 2024
4173a31
update tests
lillian542 Nov 27, 2024
947a8ab
add autograph argument to make_plxpr
lillian542 Nov 27, 2024
558901b
remove old code from previous implementation
lillian542 Dec 3, 2024
c439ed0
Merge branch 'master' into autograph_with_make_plxpr
lillian542 Dec 3, 2024
f4d8b07
update docstring
lillian542 Dec 3, 2024
5ab68c7
add malt to doc requirements
lillian542 Dec 3, 2024
be51555
make usage details collapsable
lillian542 Dec 4, 2024
ee22c0c
update changelog
lillian542 Dec 4, 2024
a73c584
does this make a collapsable tab?
lillian542 Dec 4, 2024
2682f18
Merge branch 'master' into autograph_with_make_plxpr
lillian542 Dec 4, 2024
0f57ce4
this is my least favourite game to play with sphinx
lillian542 Dec 4, 2024
c3b3c5e
Update doc/releases/changelog-dev.md
lillian542 Dec 4, 2024
4e2a206
try to add basic details section
lillian542 Dec 4, 2024
aa78fba
Merge branch 'autograph_with_make_plxpr' of github.com:PennyLaneAI/pe…
lillian542 Dec 4, 2024
7c8520d
add some of the content back in
lillian542 Dec 4, 2024
8f9189a
changelog update
lillian542 Dec 4, 2024
b6f2695
fix merge conflicts
lillian542 Dec 4, 2024
3bdd4a6
fingers crossed it all renders now
lillian542 Dec 4, 2024
f4c7a9c
add space between example cells
lillian542 Dec 4, 2024
e30dbdb
Apply suggestions from code review
lillian542 Dec 4, 2024
42a2656
re-organize changelog
lillian542 Dec 4, 2024
3234898
Merge branch 'master' into autograph_with_make_plxpr
lillian542 Dec 4, 2024
f845987
don't reorganize changelog
lillian542 Dec 5, 2024
3649d94
don't reorganize changelog pt2
lillian542 Dec 5, 2024
945a5bd
Merge branch 'master' into autograph_with_make_plxpr
lillian542 Dec 5, 2024
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
74 changes: 37 additions & 37 deletions doc/releases/changelog-dev.md
Original file line number Diff line number Diff line change
@@ -49,6 +49,42 @@
a `capabilities` property.
[(#6632)](https://github.com/PennyLaneAI/pennylane/pull/6632)


* Added `unary_mapping()` function to map `BoseWord` and `BoseSentence` to qubit operators, using unary mapping
[(#6576)](https://github.com/PennyLaneAI/pennylane/pull/6576);
added `binary_mapping()` function to map `BoseWord` and `BoseSentence` to qubit operators, using standard-binary mapping.
[(#6564)](https://github.com/PennyLaneAI/pennylane/pull/6564)

* Support is added for `if`/`else` statements and `for` and `while` loops in circuits executed with `qml.capture.enabled`, via Autograph.
Autograph conversion is now used by default in `make_plxpr`, but can be skipped with the keyword arg `autograph=False`.
[(#6406)](https://github.com/PennyLaneAI/pennylane/pull/6406)
[(#6413)](https://github.com/PennyLaneAI/pennylane/pull/6413)
[(#6426)](https://github.com/PennyLaneAI/pennylane/pull/6426)
[(#6645)](https://github.com/PennyLaneAI/pennylane/pull/6645)

* Added `christiansen_mapping()` function to map `BoseWord` and `BoseSentence` to qubit operators, using christiansen mapping.
[(#6623)](https://github.com/PennyLaneAI/pennylane/pull/6623)

* The `qml.qchem.factorize` function now supports new methods for double factorization:
Cholesky decomposition (`cholesky=True`) and compressed double factorization (`compressed=True`).
[(#6573)](https://github.com/PennyLaneAI/pennylane/pull/6573)
[(#6611)](https://github.com/PennyLaneAI/pennylane/pull/6611)

* Added `qml.qchem.symmetry_shift` function to perform the
[block-invariant symmetry shift](https://arxiv.org/pdf/2304.13772) on the electronic integrals.
[(#6574)](https://github.com/PennyLaneAI/pennylane/pull/6574)

* Added submodule for calculating vibrational Hamiltonians
* Implemented helper functions for geometry optimization, harmonic analysis,
and normal-mode localization.
[(#6453)](https://github.com/PennyLaneAI/pennylane/pull/6453)
* Implemented helper functions for calculating one-mode PES, two-mode PES, and
three-mode PES.
[(#6616)](https://github.com/PennyLaneAI/pennylane/pull/6616)
* Implemented wrapper function for vibrational Hamiltonian calculation and dataclass
for storing the data.
[(#6652)](https://github.com/PennyLaneAI/pennylane/pull/6652)

<h4>New `labs` module `dla` for handling dynamical Lie algebras (DLAs)</h4>

* Added a dense implementation of computing the Lie closure in a new function
@@ -66,11 +102,6 @@
* Added utility functions for handling dense matrices in the Lie theory context.
[(#6563)](https://github.com/PennyLaneAI/pennylane/pull/6563)

* Added `unary_mapping()` function to map `BoseWord` and `BoseSentence` to qubit operators, using unary mapping
[(#6576)](https://github.com/PennyLaneAI/pennylane/pull/6576);
added `binary_mapping()` function to map `BoseWord` and `BoseSentence` to qubit operators, using standard-binary mapping.
[(#6564)](https://github.com/PennyLaneAI/pennylane/pull/6564)


<h4>New API for Qubit Mixed</h4>

@@ -95,38 +126,6 @@ added `binary_mapping()` function to map `BoseWord` and `BoseSentence` to qubit
* Added submodule `devices.qubit_mixed.measure` as a necessary step for the new API, featuring a `measure` function for measuring qubits in mixed-state devices.
[(#6637)](https://github.com/PennyLaneAI/pennylane/pull/6637)

* Added submodule `devices.qubit_mixed.sampling` as a necessary step for the new API, featuring functions `sample_state`, `measure_with_samples` and `sample_probs` for sampling qubits in mixed-state devices.
[(#6639)](https://github.com/PennyLaneAI/pennylane/pull/6639)

* Support is added for `if`/`else` statements and `for` and `while` loops in circuits executed with `qml.capture.enabled`, via `autograph`
[(#6406)](https://github.com/PennyLaneAI/pennylane/pull/6406)
[(#6413)](https://github.com/PennyLaneAI/pennylane/pull/6413)
[(#6426)](https://github.com/PennyLaneAI/pennylane/pull/6426)

* Added `christiansen_mapping()` function to map `BoseWord` and `BoseSentence` to qubit operators, using christiansen mapping.
[(#6623)](https://github.com/PennyLaneAI/pennylane/pull/6623)

* The `qml.qchem.factorize` function now supports new methods for double factorization:
Cholesky decomposition (`cholesky=True`) and compressed double factorization (`compressed=True`).
[(#6573)](https://github.com/PennyLaneAI/pennylane/pull/6573)
[(#6611)](https://github.com/PennyLaneAI/pennylane/pull/6611)

* Added `qml.qchem.symmetry_shift` function to perform the
[block-invariant symmetry shift](https://arxiv.org/pdf/2304.13772) on the electronic integrals.
[(#6574)](https://github.com/PennyLaneAI/pennylane/pull/6574)

* Added submodule for calculating vibrational Hamiltonians
* Implemented helper functions for geometry optimization, harmonic analysis,
and normal-mode localization.
[(#6453)](https://github.com/PennyLaneAI/pennylane/pull/6453)
[(#6666)](https://github.com/PennyLaneAI/pennylane/pull/6666)
* Implemented helper functions for calculating one-mode PES, two-mode PES, and
three-mode PES.
[(#6616)](https://github.com/PennyLaneAI/pennylane/pull/6616)
* Implemented wrapper function for vibrational Hamiltonian calculation and dataclass
for storing the data.
[(#6652)](https://github.com/PennyLaneAI/pennylane/pull/6652)

<h3>Improvements 🛠</h3>

* Raises a comprehensive error when using `qml.fourier.qnode_spectrum` with standard numpy
@@ -370,6 +369,7 @@ Utkarsh Azad,
Astral Cai,
Yushao Chen,
Diksha Dhawan,
Lillian Frederiksen,
Pietropaolo Frisoni,
Austin Huang,
Korbinian Kottmann,
1 change: 1 addition & 0 deletions doc/requirements.txt
Original file line number Diff line number Diff line change
@@ -36,3 +36,4 @@ pennylane-sphinx-theme
tomli~=2.0.0 # Drop once minimum Python version is 3.11
sphinxext-opengraph==0.6.3 # Note: latest version 0.9.0 requires Sphinx >=4.0
matplotlib==3.8.0
diastatic-malt
56 changes: 55 additions & 1 deletion pennylane/capture/make_plxpr.py
Original file line number Diff line number Diff line change
@@ -18,14 +18,18 @@

import pennylane as qml

from .autograph import run_autograph

has_jax = True
try:
import jax
except ImportError: # pragma: no cover
has_jax = False


def make_plxpr(func: Callable, static_argnums: Union[int, Sequence[int]] = (), **kwargs):
def make_plxpr(
func: Callable, static_argnums: Union[int, Sequence[int]] = (), autograph=True, **kwargs
):
r"""Takes a function and returns a ``Callable`` that, when called, produces a PLxPR representing
the function with the given args.
@@ -39,6 +43,8 @@ def make_plxpr(func: Callable, static_argnums: Union[int, Sequence[int]] = (), *
Kwargs:
static_argnums (Union(int, Sequence[int])): optional, an ``int`` or collection of ``int``\ s
that specify which positional arguments to treat as static (trace- and compile-time constant).
autograph (bool): whether to use AutoGraph to convert Python control flow to native PennyLane
control flow. Defaults to True.
Returns:
Callable: function that, when called, returns the PLxPR representation of ``func`` for the specified inputs.
@@ -78,6 +84,51 @@ def circ(x):
] a
in (b,) }
.. details ::
:title: Usage Details
The ``autograph`` argument is ``True`` by default, converting Pythonic control flow to PennyLane
supported control flow. This requires the diastatic-malt package, a standalone fork of the AutoGraph
module in TensorFlow (`official documentation <https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/index.md>`_
).
On its own, capture of standard Python control flow is not supported:
.. code-block:: python
def fn(x):
if x > 5:
return x+1
return x+2
For this function, capture doesn't work without autograph:
>>> plxpr_fn = qml.capture.make_plxpr(fn, autograph=False)
>>> plxpr = plxpr_fn(3)
TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
With AutoGraph, the control flow is automatically converted to the native PennyLane control
flow implementation, and succeeds:
>>> plxpr_fn = qml.capture.make_plxpr(fn)
>>> plxpr = plxpr_fn(3)
>>> plxpr
{ lambda ; a:i64[]. let
b:bool[] = gt a 5
_:bool[] c:i64[] = cond[
args_slice=slice(4, None, None)
consts_slices=[slice(2, 3, None), slice(3, 4, None)]
jaxpr_branches=[{ lambda a:i64[]; . let in (True, a) }, { lambda a:i64[]; . let b:i64[] = add a 2 in (True, b) }]
] b True a a
in (c,) }
We can evaulate this to get the results:
>>> jax.core.eval_jaxpr(plxpr.jaxpr, plxpr.consts, 2)
[Array(4, dtype=int64, weak_type=True)]
>>> jax.core.eval_jaxpr(plxpr.jaxpr, plxpr.consts, 7)
[Array(8, dtype=int64, weak_type=True)]
"""
if not has_jax: # pragma: no cover
raise ImportError(
@@ -91,4 +142,7 @@ def circ(x):
"You can enable capture with ``qml.capture.enable()``"
)

if autograph:
func = run_autograph(func)

return jax.make_jaxpr(func, static_argnums=static_argnums, **kwargs)
115 changes: 108 additions & 7 deletions tests/capture/test_make_plxpr.py
Original file line number Diff line number Diff line change
@@ -27,6 +27,8 @@
jax = pytest.importorskip("jax")

# must be below jax importorskip
from jax import numpy as jnp # pylint: disable=wrong-import-position, wrong-import-order

from pennylane.capture import make_plxpr # pylint: disable=wrong-import-position


@@ -67,8 +69,9 @@ def circ(x):
assert hasattr(plxpr, "jaxpr")
isinstance(plxpr, jax._src.core.ClosedJaxpr) # pylint: disable=protected-access

@pytest.mark.parametrize("autograph", [True, False])
@pytest.mark.parametrize("static_argnums", [[0], [1], [0, 1], []])
def test_static_argnums(self, static_argnums, mocker):
def test_static_argnums(self, static_argnums, autograph, mocker):
"""Test that passing static_argnums works as expected"""

dev = qml.device("default.qubit", wires=1)
@@ -85,16 +88,19 @@ def circ(x, y):
params = [1.2, 2.3]
non_static_params = [params[i] for i in (0, 1) if i not in static_argnums]

plxpr = make_plxpr(circ, static_argnums=static_argnums)(*params)
plxpr = make_plxpr(circ, autograph=autograph, static_argnums=static_argnums)(*params)

# most recent call is to make a jaxpr of something else, so we can't use assert_called_with
spy.assert_has_calls([call(circ, static_argnums=static_argnums)])
if not autograph:
# when using autograph, we don't have the function make_jaxpr was called with
# (it was produced by run_autograph), so we can't check for the function call.
spy.assert_has_calls([call(circ, static_argnums=static_argnums)])

# plxpr behaves as expected wrt static argnums
res = jax.core.eval_jaxpr(plxpr.jaxpr, plxpr.consts, *non_static_params)
assert np.allclose(res, circ(*params))

def test_kwargs(self, mocker):
@pytest.mark.parametrize("autograph", [True, False])
def test_kwargs(self, mocker, autograph):
"""Test additional kwargs are passed through to make_jaxpr"""

dev = qml.device("default.qubit", wires=1)
@@ -106,6 +112,101 @@ def circ():
qml.Hadamard(0)
return qml.expval(qml.X(0))

output = make_plxpr(circ, autograph=autograph, return_shape=True)()

# assert new value for return_shape is passed to make_jaxpr
_ = make_plxpr(circ, return_shape=True)()
spy.assert_has_calls([call(circ, static_argnums=(), return_shape=True)])
if not autograph:
# when using autograph, we don't have the function make_jaxpr was called with
# (it was produced by run_autograph), so we can't check for the function call.
spy.assert_has_calls([call(circ, static_argnums=(), return_shape=True)], any_order=True)

# output is as expected for return_shape=True
assert len(output) == 2
isinstance(output[0], jax._src.core.ClosedJaxpr) # pylint: disable=protected-access
isinstance(output[0], jax._src.api.ShapeDtypeStruct) # pylint: disable=protected-access


@pytest.mark.usefixtures("enable_disable_plxpr")
class TestAutoGraphIntegration:
"""Test autograph integration for converting Python control flow into native PennyLane
`cond`, `for_loop` and `while_loop`. Note that autograph defaults to True in make_plxpr."""

def test_if_stmt(self):
"""Test that an if statement is converted to a jaxpr with a ``cond`` function, and
that in the case of a QNode, the resulting plxpr can be evaluated as expected"""

def func(x):
if x > 1.967:
qml.Hadamard(0)
else:
qml.Y(0)
return qml.state()

dev = qml.device("default.qubit", wires=1)
qnode = qml.QNode(func, dev)

plxpr1 = qml.capture.make_plxpr(func)(2)
plxpr2 = qml.capture.make_plxpr(qnode)(2)

# the plxpr includes a representation of a `cond` function
assert "cond[" in str(plxpr1)
assert "cond[" in str(plxpr2)

def eval(x):
return jax.core.eval_jaxpr(plxpr2.jaxpr, plxpr2.consts, x)

assert np.allclose(eval(2), [0.70710678, 0.70710678])
assert np.allclose(eval(1), [0, 1j])

def test_while_loop(self):
"""Test that a while loop is converted to a jaxpr with a ``while_loop`` function, and
that in the case of a QNode, the resulting plxpr can be evaluated as expected"""

def func(counter):
while counter < 10:
qml.RX(np.pi * 0.1, wires=0)
counter += 1

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

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

plxpr1 = qml.capture.make_plxpr(func)(0)
plxpr2 = qml.capture.make_plxpr(qnode)(0)

# the plxpr includes a representation of a `while_loop` function
assert "while_loop[" in str(plxpr1)
assert "while_loop[" in str(plxpr2)

def eval(x):
return jax.core.eval_jaxpr(plxpr2.jaxpr, plxpr2.consts, x)

assert np.allclose(eval(0), [-1])
assert np.allclose(eval(5), [0])

def test_for_loop(self):
"""Test that a for loop is converted to a jaxpr with a ``for_loop`` function, and
that in the case of a QNode, the resulting plxpr can be evaluated as expected"""

def func(angles):
for i, x in enumerate(angles):
qml.RX(x, wires=i)

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

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

plxpr1 = qml.capture.make_plxpr(func)(jnp.array([0.0, 0.0]))
plxpr2 = qml.capture.make_plxpr(qnode)(jnp.array([0.0, 0.0]))

# the plxpr includes a representation of a `for_loop` function
assert "for_loop[" in str(plxpr1)
assert "for_loop[" in str(plxpr2)

def eval(x):
x = jnp.array(x)
return jax.core.eval_jaxpr(plxpr2.jaxpr, plxpr2.consts, x)

assert np.allclose(eval([np.pi, np.pi / 2]), [-1, 0])