Allow device to configure conversion to numpy and use of `pure_callback` #6788

albi3ro · 2025-01-08T19:00:02Z

Context:

While we have logic for sampling with jax, it does not really integrate very well into the workflow. While you can technically set diff_method=None right now and jit the execution end-to-end, trying to jit diff_method=None will cause incomprehensible error messages on non-DQ devices.

We want to forbid differentiation diff_method=None, but keep a way to jit a finite shot execution.

Description of the Change:

In order to allow jitting finite shot executions, we need a way for the device to be able to configure whether or not the data is converted to numpy. To do so, we simply add another property to the ExecutionConfig, convert_to_numpy. If False, then we will not use a pure_callback to convert the parameters to numpy. If True, we use a pure_callback and convert the parameters to numpy.

Benefits:

Speed ups due to being able to jit the entire execution.

Possible Drawbacks:

ExecutionConfig gets an addtional property, making it more complicated.

Related GitHub Issues:

Fixes #6054 Fixes #3259 Blocks #6770

github-actions · 2025-01-08T19:00:23Z

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

…I/pennylane into no-interface-boundary

codecov · 2025-01-09T15:58:52Z

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 99.60%. Comparing base (9350cb2) to head (9c7ea17).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #6788   +/-   ##
=======================================
  Coverage   99.60%   99.60%           
=======================================
  Files         476      476           
  Lines       45222    45248   +26     
=======================================
+ Hits        45045    45071   +26     
  Misses        177      177

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

tests/test_qnode.py

andrijapau · 2025-01-09T18:12:45Z

doc/releases/changelog-dev.md

+* Devices can now configure whether or not the data is converted to numpy and `jax.pure_callback`
+  is used by the new `ExecutionConfig.convert_to_numpy` property. Finite shot executions
+  on `default.qubit` can now be jitted end-to-end, even with parameter shift.
+  [(#6788)](https://github.com/PennyLaneAI/pennylane/pull/6788)
+


Suggested change

* Devices can now configure whether or not the data is converted to numpy and `jax.pure_callback`

is used by the new `ExecutionConfig.convert_to_numpy` property. Finite shot executions

on `default.qubit` can now be jitted end-to-end, even with parameter shift.

[(#6788)](https://github.com/PennyLaneAI/pennylane/pull/6788)

* Devices can now configure whether or not the data is converted to numpy enabling `jax.pure_callback` to be used by the new `ExecutionConfig.convert_to_numpy` property. Finite shot executions on `default.qubit` can now be jitted end-to-end leading to performance improvements, even with parameter shift.

[(#6788)](https://github.com/PennyLaneAI/pennylane/pull/6788)

I think the new option sounds a little confusing to me. What was the problem with the original note? I'll see if I can find a better way to phrase it. Maybe just leave out the reference to jax.pure_callback?

For this improvement, I recommend swapping the sentences around. We should always assume a reader is less likely to keep reading -- what is the first thing we want them to takeaway?

In this case, the first sentence should likely speak to what the user-facing improvement is (and how a user can do it, if relevant), followed by the actual implementation change.

andrijapau

Just some comments on my initial pass through.

pennylane/devices/default_qubit.py

astralcai

LGTM

mudit2812 · 2025-01-10T16:22:26Z

tests/devices/default_qubit/test_default_qubit_preprocessing.py

+    def test_convert_to_numpy_with_adjoint(self):
+        """Test that we will convert to numpy with adjoint."""
+        config = qml.devices.ExecutionConfig(gradient_method="adjoint", interface="jax-jit")
+        dev = qml.device("default.qubit")
+        processed = dev.setup_execution_config(config)
+        assert processed.convert_to_numpy
+
+    @pytest.mark.parametrize("interface", ("autograd", "torch", "tf"))
+    def test_convert_to_numpy_non_jax(self, interface):
+        """Test that other interfaces are still converted to numpy."""
+        config = qml.devices.ExecutionConfig(gradient_method="adjoint", interface=interface)
+        dev = qml.device("default.qubit")
+        processed = dev.setup_execution_config(config)
+        assert processed.convert_to_numpy


Can you include "jax" as an interface in the testing?

Also, I'm curious if converting to numpy with adjoint has negative effects on performance.

I do think we could make adjoint jittable and get some nice speed boosts, but i think that would need to be a follow on task.

Also, the jax interface is tested in test_not_convert_to_numpy_with_jax.

mudit2812 · 2025-01-10T16:25:40Z

tests/templates/test_state_preparations/test_mottonen_state_prep.py

@@ -417,7 +417,7 @@ def circuit(state):


 @pytest.mark.jax
-@pytest.mark.parametrize("shots, atol", [(None, 0.005), (1000000, 0.05)])
+@pytest.mark.parametrize("shots, atol", [(None, 0.005), (1000000, 0.1)])


This is a lot of shots...

If we're using a seed, maybe it's worth reducing the number of shots.

Given the tolerance is still almost to high... the problem is that this test is now doing finite shot finite differences with float32. Im not sure we can get any accuracy out of that.

Do we even need a finite shot test here?

Definitely seems like the wrong file for this type of test...

I'm just going to remove the finite shot version.

tests/workflow/interfaces/execute/test_jax.py

allow device to configure conversion to numpy

cbe1f20

albi3ro and others added 7 commits January 8, 2025 15:17

testing and changelog

cef0e31

updating xfailing pulse gradient tests

f6949e1

fixing more tests

1b069b8

Merge branch 'master' into no-interface-boundary

a733cca

fixing more tests

fbce8cc

Merge branch 'no-interface-boundary' of https://github.com/PennyLaneA…

df4bd13

…I/pennylane into no-interface-boundary

final test

2110c9b

albi3ro requested review from astralcai, mudit2812 and andrijapau January 9, 2025 16:45

andrijapau reviewed Jan 9, 2025

View reviewed changes

tests/test_qnode.py Outdated Show resolved Hide resolved

andrijapau reviewed Jan 9, 2025

View reviewed changes

pennylane/devices/default_qubit.py Outdated Show resolved Hide resolved

astralcai reviewed Jan 9, 2025

View reviewed changes

pennylane/devices/default_qubit.py Show resolved Hide resolved

astralcai reviewed Jan 9, 2025

View reviewed changes

pennylane/devices/default_qubit.py Show resolved Hide resolved

astralcai approved these changes Jan 9, 2025

View reviewed changes

mudit2812 reviewed Jan 10, 2025

View reviewed changes

albi3ro and others added 3 commits January 13, 2025 10:48

responses to review

3f412a1

oops

1ce1b70

Merge branch 'master' into no-interface-boundary

a089c37

albi3ro requested a review from andrijapau January 13, 2025 17:11

changelog and plugins page

1c627e1

albi3ro requested a review from mudit2812 January 13, 2025 17:34

albi3ro mentioned this pull request Jan 13, 2025

[Exploration] Natively sampling in torch in with default.qubit #6814

Open

remove bad mottonen state prep test

9c7ea17

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Allow device to configure conversion to numpy and use of `pure_callback` #6788

Allow device to configure conversion to numpy and use of `pure_callback` #6788

albi3ro commented Jan 8, 2025 •

edited

Loading

github-actions bot commented Jan 8, 2025

codecov bot commented Jan 9, 2025 •

edited

Loading

andrijapau Jan 9, 2025

albi3ro Jan 13, 2025

josh146 Jan 13, 2025

andrijapau left a comment

astralcai left a comment

mudit2812 Jan 10, 2025

albi3ro Jan 10, 2025

albi3ro Jan 13, 2025

mudit2812 Jan 10, 2025

mudit2812 Jan 10, 2025

albi3ro Jan 10, 2025

albi3ro Jan 13, 2025

mudit2812 Jan 13, 2025

albi3ro Jan 13, 2025

Allow device to configure conversion to numpy and use of pure_callback #6788

Are you sure you want to change the base?

Allow device to configure conversion to numpy and use of pure_callback #6788

Conversation

albi3ro commented Jan 8, 2025 • edited Loading

github-actions bot commented Jan 8, 2025

codecov bot commented Jan 9, 2025 • 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

andrijapau left a comment

Choose a reason for hiding this comment

astralcai 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Allow device to configure conversion to numpy and use of `pure_callback` #6788

Allow device to configure conversion to numpy and use of `pure_callback` #6788

albi3ro commented Jan 8, 2025 •

edited

Loading

codecov bot commented Jan 9, 2025 •

edited

Loading