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 (5efeffb) to head (2110c9b).
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       45232    45242   +10     
=======================================
+ Hits        45055    45065   +10     
  Misses        177      177

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

andrijapau · 2025-01-09T17:03:01Z

tests/test_qnode.py

+            out = jitted_qnode1(0.123)
+            print(out)


Suggested change

out = jitted_qnode1(0.123)

print(out)

jitted_qnode1(0.123)

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)

andrijapau

Just some comments on my initial pass through.

andrijapau · 2025-01-09T18:15:40Z

pennylane/devices/default_qubit.py

@@ -581,6 +591,12 @@ def _setup_execution_config(self, execution_config: ExecutionConfig) -> Executio
        """
        updated_values = {}

+        updated_values["convert_to_numpy"] = (
+            execution_config.interface.value not in {"jax", "jax-jit"}


Do we want to continue using the Interface enum instead?

Suggested change

execution_config.interface.value not in {"jax", "jax-jit"}

execution_config.interface not in {qml.math.Interface.JAX, qml.math.Interface.JAX_JIT}

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.

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.

mudit2812 · 2025-01-10T16:27:47Z

tests/workflow/interfaces/execute/test_jax.py


            result = execute([tape1, tape2], dev, diff_method=param_shift, max_diff=1)
            return result[0] + result[1][0]

        res = cost_fn(params)
        x, y = params
        expected = 0.5 * (3 + jnp.cos(x) ** 2 * jnp.cos(2 * y))
-        assert np.allclose(res, expected, atol=tol, rtol=0)
+        assert np.allclose(res, expected, atol=2e-2, rtol=0)


If the tol isn't being used, it can be removed as an argument to this test.

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

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

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

andrijapau Jan 9, 2025

andrijapau Jan 9, 2025

andrijapau left a comment

andrijapau Jan 9, 2025

astralcai left a comment

mudit2812 Jan 10, 2025

albi3ro Jan 10, 2025

mudit2812 Jan 10, 2025

mudit2812 Jan 10, 2025

albi3ro Jan 10, 2025

mudit2812 Jan 10, 2025

	execution_config.interface.value not in {"jax", "jax-jit"}
	execution_config.interface not in {qml.math.Interface.JAX, qml.math.Interface.JAX_JIT}

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

Codecov Report

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

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

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