[Capture] Compute the jvp of a jaxpr using default.qubit tools #6875

albi3ro · 2025-01-22T23:29:22Z

Context:

See PR #6905 for how this will fit in to the entire workflow.

Description of the Change:

Adds a execute_and_jvp function to devices/qubit/jaxpr_adjoint for computing the results and jvp of a jaxpr using adjoint jacobian.

Benefits:

More jaxpr derivatives.

Possible Drawbacks:

Not very featureful and robust right now.

Related GitHub Issues:

[sc-82169]

codecov · 2025-01-30T20:31:32Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.59%. Comparing base (ca9a547) to head (06dda62).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6875   +/-   ##
=======================================
  Coverage   99.59%   99.59%           
=======================================
  Files         478      479    +1     
  Lines       45294    45375   +81     
=======================================
+ Hits        45112    45193   +81     
  Misses        182      182

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

dwierichs

Thanks @albi3ro, this addition is really nice to see! 🎉
I like how the separation of forward and backward pass is so convenient to separate the two parts of the adjoint method algorithm 😍
Nice and clean tests as well!

The only real request I'd have is to add some dev comments to make it easier for future work on this code.

pennylane/devices/qubit/jaxpr_adjoint.py

dwierichs · 2025-02-04T20:17:55Z

tests/devices/qubit/test_jaxpr_adjoint_jvp.py

+
+        jaxpr = jax.make_jaxpr(f)(0.5).jaxpr
+
+        with pytest.raises(ValueError):


Shouldn't this a generator-undefined related error? 🤔

That's what I thought. I'm planning on fixing that at some point soon, but for now, ValueError is the behavior.

tests/devices/qubit/test_jaxpr_adjoint_jvp.py

dwierichs · 2025-02-04T20:20:14Z

tests/devices/qubit/test_jaxpr_adjoint_jvp.py

+        assert qml.math.allclose(dres, dexpected)
+
+    def test_jaxpr_consts(self):
+        """Test that we can execute jaxpr with consts."""


Are closed-over variables supported already as well? :) Could test that in the same test, possibly

consts and closed over variables are essentially the same thing from the jaxpr perspective. They both go into jaxpr.consts. Since I'm testing executing the jaxpr, and not capturing it, I think one way of producing consts should be fine.

dwierichs · 2025-02-04T20:21:53Z

pennylane/devices/qubit/jaxpr_adjoint.py

+    env[eqn.outvars[0]] = (op, ad.Zero(AbstractOperator()))
+
+    if any(not isinstance(t, ad.Zero) for t in tangents[1:]):
+        raise NotImplementedError("adjoint jvp only differentiable parameters in the 0 position.")


Does this stem from the restriction that the adjoint method only works for single-parameter/single-generator gates? Is the positional restriction an additional restriction? Maybe a dev comment could clarify :)

limitation for the moment to handle both the fact that generators are define with respect a single parameter, and it doesn' t make sense for wires to be differentiable.

Co-authored-by: David Wierichs <[email protected]>

PietropaoloFrisoni

Looks good to me!

PietropaoloFrisoni · 2025-02-05T22:26:06Z

doc/releases/changelog-dev.md

+* The adjoint jvp of a jaxpr can be computed using default.qubit tooling.
+  [(#6875)](https://github.com/PennyLaneAI/pennylane/pull/6875)
+
+* The source code has been updated use black 25.1.0


Suggested change

* The source code has been updated use black 25.1.0

This entry already exists (I am sure this is just the result of a merge conflict in the changelog)

PietropaoloFrisoni · 2025-02-05T22:29:44Z

pennylane/devices/qubit/jaxpr_adjoint.py

+            include the consts followed by the inputs.
+        num_wires (int): the number of wires to use.
+
+    Note that the consts for the jaxpr should be included in the beginning of both the ``args``


Suggested change

Note that the consts for the jaxpr should be included in the beginning of both the ``args``

Note that the consts for the jaxpr should be included at the beginning of both the ``args``

PietropaoloFrisoni · 2025-02-05T22:40:40Z

pennylane/devices/qubit/jaxpr_adjoint.py

+    for eqn in jaxpr.eqns:
+        if getattr(eqn.primitive, "prim_type", "") == "operator":
+            ket = _operator_forward_pass(eqn, env, ket)
+
+        elif getattr(eqn.primitive, "prim_type", "") == "measurement":
+            bra = _measurement_forward_pass(eqn, env, ket)
+            bras.append(bra)
+        else:
+            _other_prim_forward_pass(eqn, env)


Suggested change

for eqn in jaxpr.eqns:

if getattr(eqn.primitive, "prim_type", "") == "operator":

ket = _operator_forward_pass(eqn, env, ket)

elif getattr(eqn.primitive, "prim_type", "") == "measurement":

bra = _measurement_forward_pass(eqn, env, ket)

bras.append(bra)

else:

_other_prim_forward_pass(eqn, env)

for eqn in jaxpr.eqns:

prim_type = getattr(eqn.primitive, "prim_type", "")

if prim_type == "operator":

ket = _operator_forward_pass(eqn, env, ket)

elif prim_type == "measurement":

bras.append(_measurement_forward_pass(eqn, env, ket))

else:

_other_prim_forward_pass(eqn, env)

What do you think about this version with only one call to getattr?

PietropaoloFrisoni · 2025-02-05T22:43:41Z

pennylane/devices/qubit/jaxpr_adjoint.py

+                t = _read(env, eqn.invars[0])[1]
+                if not isinstance(t, ad.Zero):
+                    disable()
+                    try:
+                        ket_temp = apply_operation(generator(op, format="observable"), ket)
+                    finally:
+                        enable()
+                    modified = True
+                    for i, bra in enumerate(bras):
+                        out_jvps[i] += -2 * t * jnp.imag(jnp.sum(jnp.conj(bra) * ket_temp))


Suggested change

t = _read(env, eqn.invars[0])[1]

if not isinstance(t, ad.Zero):

disable()

try:

ket_temp = apply_operation(generator(op, format="observable"), ket)

finally:

enable()

modified = True

for i, bra in enumerate(bras):

out_jvps[i] += -2 * t * jnp.imag(jnp.sum(jnp.conj(bra) * ket_temp))

tang = _read(env, eqn.invars[0])[1]

if not isinstance(tang , ad.Zero):

disable()

try:

ket_temp = apply_operation(generator(op, format="observable"), ket)

finally:

enable()

modified = True

for i, bra in enumerate(bras):

out_jvps[i] += -2 * tang * jnp.imag(jnp.sum(jnp.conj(bra) * ket_temp))

Can we add a more self-explanatory name than t? I prosed tang, but feel free to change it : )

PietropaoloFrisoni · 2025-02-05T22:52:48Z

pennylane/devices/qubit/jaxpr_adjoint.py

+
+def _measurement_forward_pass(eqn, env, ket):
+    """Perform a measurement during the forward pass of the adjoint jvp."""
+    invals, tangents = tuple(zip(*(_read(env, var) for var in eqn.invars)))


Suggested change

invals, tangents = tuple(zip(*(_read(env, var) for var in eqn.invars)))

invals, tangents = zip(*(_read(env, var) for var in eqn.invars))

I suggested this, although I suspect you wanted to explicitly convert this to a tuple for efficiency. Is that so?

PietropaoloFrisoni · 2025-02-05T22:55:17Z

pennylane/devices/qubit/jaxpr_adjoint.py

+
+    mp = eqn.primitive.impl(*invals, **eqn.params)
+    bra = apply_operation(mp.obs, ket)
+    result = jnp.real(jnp.sum(jnp.conj(bra) * ket))


Suggested change

result = jnp.real(jnp.sum(jnp.conj(bra) * ket))

result = jnp.real(jnp.vdot(bra, ket))

Should be faster

albi3ro and others added 8 commits January 20, 2025 13:47

adjoint jvp on capture dq

98bd43f

adjoint jvp of jaxpr

5cdd1cb

Merge branch 'master' into capture-dq-adjoint-jvp

353a3af

consts are just args

b2ef102

cleaning up

8bd6338

starting on tests

26a8da0

testing

c013510

Merge branch 'master' into capture-dq-adjoint-jvp

06dda62

albi3ro marked this pull request as ready for review January 30, 2025 20:13

albi3ro requested review from PietropaoloFrisoni and dwierichs January 31, 2025 16:57

dwierichs requested changes Feb 4, 2025

View reviewed changes

albi3ro and others added 2 commits February 4, 2025 18:02

Apply suggestions from code review

bb3bf09

Co-authored-by: David Wierichs <[email protected]>

Merge branch 'master' into capture-dq-adjoint-jvp

92d549a

PietropaoloFrisoni reviewed Feb 5, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Capture] Compute the jvp of a jaxpr using default.qubit tools #6875

[Capture] Compute the jvp of a jaxpr using default.qubit tools #6875

albi3ro commented Jan 22, 2025 •

edited

Loading

codecov bot commented Jan 30, 2025

dwierichs left a comment

dwierichs Feb 4, 2025

albi3ro Feb 5, 2025

dwierichs Feb 4, 2025

albi3ro Feb 4, 2025

dwierichs Feb 4, 2025

albi3ro Feb 4, 2025

PietropaoloFrisoni left a comment

PietropaoloFrisoni Feb 5, 2025

PietropaoloFrisoni Feb 5, 2025

PietropaoloFrisoni Feb 5, 2025

PietropaoloFrisoni Feb 5, 2025

PietropaoloFrisoni Feb 5, 2025 •

edited

Loading

PietropaoloFrisoni Feb 5, 2025


		jaxpr = jax.make_jaxpr(f)(0.5).jaxpr

		with pytest.raises(ValueError):

	Note that the consts for the jaxpr should be included in the beginning of both the ``args``
	Note that the consts for the jaxpr should be included at the beginning of both the ``args``

	invals, tangents = tuple(zip(*(_read(env, var) for var in eqn.invars)))
	invals, tangents = zip(*(_read(env, var) for var in eqn.invars))

	result = jnp.real(jnp.sum(jnp.conj(bra) * ket))
	result = jnp.real(jnp.vdot(bra, ket))

[Capture] Compute the jvp of a jaxpr using default.qubit tools #6875

Are you sure you want to change the base?

[Capture] Compute the jvp of a jaxpr using default.qubit tools #6875

Conversation

albi3ro commented Jan 22, 2025 • edited Loading

codecov bot commented Jan 30, 2025

Codecov Report

dwierichs 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

PietropaoloFrisoni 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

PietropaoloFrisoni Feb 5, 2025 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

albi3ro commented Jan 22, 2025 •

edited

Loading

PietropaoloFrisoni Feb 5, 2025 •

edited

Loading