Jax all devices #1065

chaserileyroberts · 2021-02-03T20:01:37Z

Context:
Previously, only the default.qubit.jax device supported the JAX interface. This PR adds JAX interface support to all devices.

Description of the Change:
The recent addition of jax.experimental.host_callback allows us to finally do jax -> numpy -> numpy -> jax within a jax.jit function! This PR adds the needed scaffolding to do it.

Benefits:
Everyone can start using JAX more and I won't take any more execuses.

Possible Drawbacks:
None.

Related GitHub Issues:
#943

TODOs

Gradient support
Vmap support

josh146

This is awesome, I'm excited to have this in!

Most of my suggestions in the PR itself are minor, but there are some other parts of the codebase that should be modified:

We likely need an integration test for this new functionality. Could you make a copy of tests/tape/interfaces/test_qnode_autograd.py as test_qnode_jax.py? Probably the entire file could stay as is, with minor modification to account for the autograd -> JAX change.
Small tweaks to the JAX page in the docs need to be made to account for this change

.github/CHANGELOG.md

josh146 · 2021-02-04T06:07:49Z

.github/CHANGELOG.md

+  dev = qml.device('cirq.simulator', wires=1)
+  @qml.qnode(dev, interface="jax")
+  def circuit(x):
+      qml.RX(x[1], wires=0)
+      qml.Rot(x[0], x[1], x[2], wires=0)
+      return qml.expval(qml.PauliZ(0))
+
+  weights = jnp.array([0.2, 0.5, 0.1])
+  print(circuit(weights)) # DeviceArray(...)


Woah. This is awesome.

pennylane/tape/interfaces/jax.py

josh146 · 2021-02-04T06:24:37Z

pennylane/tape/interfaces/jax.py

+        if len(self.observables) != 1:
+            raise ValueError("Only one return type is supported currently")
+        return_type = self.observables[0].return_type
+        if return_type is not Variance and return_type is not Expectation:
+            raise ValueError(
+                f"Only Variance and Expectation returns are support, given {return_type}"
+            )
+        exec_fn = partial(self.execute_device, device=device)


Could this be simplified using self.output_dim?

return host_callback.call( exec_fn, params, result_shape=jax.ShapeDtypeStruct(self.output_dim, JAXInterface.dtype) )

However, there would still need to be one validation; a validation to ensure that there is no mixture of probs and expval/var. E.g., anything that would result in a different output dim to what the tape estimates.

Can I leave that to a future PR where I fix the above todo?

josh146 · 2021-02-04T06:25:34Z

pennylane/tape/interfaces/jax.py

+    @classmethod
+    def apply(cls, tape):
+        """Apply the JAX interface to an existing tape in-place.
+
+        Args:
+            tape (.JacobianTape): a quantum tape to apply the JAX interface to
+
+        **Example**
+
+        >>> with JacobianTape() as tape:
+        ...     qml.RX(0.5, wires=0)
+        ...     expval(qml.PauliZ(0))
+        >>> JAXInterface.apply(tape)
+        >>> tape
+        <JAXQuantumTape: wires=<Wires = [0]>, params=1>
+        """
+        tape_class = getattr(tape, "__bare__", tape.__class__)
+        tape.__bare__ = tape_class
+        tape.__class__ = type("JAXQuantumTape", (cls, tape_class), {})
+        return tape


Tangent: In retrospect, this is a design decision I do not like. Even ignoring the dynamic class change, this feels more like a function than a class method?

I'm usually always against inplace transformations, and this one is pretty intense. tbh I don't really understand how this works.

iirc the in-place transform was solely to placate Autograd. Everything else worked at the time without it being in-place (tf, torch), but autograd would give a really bizarre and long traceback that vanished when it was made in-place instead.

(it could have been torch that complained rather than autograd, I can't remember fully).

josh146 · 2021-02-04T06:26:14Z

tests/devices/test_default_qubit_jax.py

-    def test_non_backprop_error(self):
-        """Test that an error is raised in tape mode if the diff method is not backprop"""


josh146 · 2021-02-04T06:34:44Z

tests/tape/interfaces/test_tape_jax.py

+        # Easiest way to test object is a device array instead of np.array
+        assert "DeviceArray" in res.__repr__()


Co-authored-by: Josh Izaac <josh146@gmail.com>

codecov · 2021-02-04T19:06:18Z

Codecov Report

Merging #1065 (17364b0) into master (85dd193) will decrease coverage by 0.03%.
The diff coverage is 88.09%.

@@            Coverage Diff             @@
##           master    #1065      +/-   ##
==========================================
- Coverage   97.74%   97.71%   -0.04%     
==========================================
  Files         153      154       +1     
  Lines       11579    11551      -28     
==========================================
- Hits        11318    11287      -31     
- Misses        261      264       +3

Impacted Files	Coverage Δ
pennylane/devices/autograd_ops.py	`100.00% <ø> (+2.38%)`	⬆️
pennylane/devices/default_mixed.py	`100.00% <ø> (ø)`
pennylane/devices/default_qubit.py	`100.00% <ø> (ø)`
pennylane/devices/default_qubit_autograd.py	`100.00% <ø> (ø)`
pennylane/devices/default_qubit_jax.py	`93.75% <ø> (ø)`
pennylane/devices/default_qubit_tf.py	`90.14% <ø> (ø)`
pennylane/devices/jax_ops.py	`97.43% <ø> (+2.31%)`	⬆️
pennylane/devices/tf_ops.py	`100.00% <ø> (ø)`
pennylane/templates/embeddings/amplitude.py	`100.00% <ø> (ø)`
pennylane/vqe/vqe.py	`93.78% <ø> (ø)`
... and 4 more

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 85dd193...594edff. Read the comment docs.

chaserileyroberts · 2021-02-04T19:39:03Z

We likely need an integration test for this new functionality. Could you make a copy of tests/tape/interfaces/test_qnode_autograd.py as test_qnode_jax.py? Probably the entire file could stay as is, with minor modification to account for the autograd -> JAX change.

I don't have gradient support yet, so I'll only be able to support some of the tests.

chaserileyroberts · 2021-02-06T01:42:44Z

@josh146 Looks like this is still broken despite check on the build passing. Can you tag your github support on this PR?

Co-authored-by: Josh Izaac <josh146@gmail.com>

josh146 · 2021-02-08T16:49:16Z

.github/CHANGELOG.md

+      return qml.expval(qml.PauliZ(0))
+
+  weights = jnp.array([0.2, 0.5, 0.1])
+  print(circuit(weights)) # DeviceArray(...)


I can't suggest it, but this block is missing an end triple-back-quote

https://github.com/PennyLaneAI/pennylane/runs/1855194494#step:4:2022

Aw thank you. I was wondering why the doc build kept failing here but worked locally.

Chase Roberts added 9 commits January 28, 2021 00:54

Added JAX interface for non-jax devices

1d94f4c

Minor fixes

9b4634f

Fix test

7218f63

Added qnode test

b7e0db5

black

18e7ff0

Added to Changelog

a2aeea9

Fix hanging test

21ac7df

typo

d276c07

typo

bd06af1

chaserileyroberts requested a review from josh146 February 3, 2021 20:01

chaserileyroberts added the review-ready 👌 PRs which are ready for review by someone from the core team. label Feb 3, 2021

josh146 requested changes Feb 4, 2021

View reviewed changes

Chase Roberts and others added 3 commits February 4, 2021 10:58

Update .github/CHANGELOG.md

dd7e1a2

Co-authored-by: Josh Izaac <josh146@gmail.com>

Update .github/CHANGELOG.md

12597a2

Co-authored-by: Josh Izaac <josh146@gmail.com>

Update pennylane/tape/interfaces/jax.py

bccb33c

Co-authored-by: Josh Izaac <josh146@gmail.com>

Chase Roberts and others added 2 commits February 5, 2021 07:31

Merge branch 'master' into jax_all_devices

fe7ba43

Merge branch 'master' into jax_all_devices

5b625b1

josh146 and others added 10 commits February 8, 2021 19:53

Merge branch 'master' into jax_all_devices

4533546

Update pennylane/tape/interfaces/jax.py

f74a99f

Co-authored-by: Josh Izaac <josh146@gmail.com>

Update pennylane/tape/interfaces/jax.py

50993c4

Co-authored-by: Josh Izaac <josh146@gmail.com>

Update pennylane/tape/interfaces/jax.py

765059a

Co-authored-by: Josh Izaac <josh146@gmail.com>

Fixed docs autograd -> jax

7c0953f

Docs fix autograd -> jax

04bb91d

Fixed merge conflicts

7b7a520

Added simple qnode integration test

61a7a2f

Doc fix

88699ff

pytest import skip

357451f

Test fix

876178f

josh146 reviewed Feb 8, 2021

View reviewed changes

Doc fixes

594edff

chaserileyroberts mentioned this pull request Feb 8, 2021

Jax interface for all devices. #1076

Merged

chaserileyroberts closed this Feb 8, 2021

josh146 deleted the jax_all_devices branch April 29, 2021 09:46

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Jax all devices #1065

Jax all devices #1065

chaserileyroberts commented Feb 3, 2021 •

edited

Loading

josh146 left a comment

josh146 Feb 4, 2021

josh146 Feb 4, 2021

chaserileyroberts Feb 4, 2021

josh146 Feb 5, 2021

josh146 Feb 4, 2021 •

edited

Loading

chaserileyroberts Feb 4, 2021

josh146 Feb 5, 2021

josh146 Feb 5, 2021

josh146 Feb 4, 2021

josh146 Feb 4, 2021

codecov bot commented Feb 4, 2021 •

edited

Loading

chaserileyroberts commented Feb 4, 2021

chaserileyroberts commented Feb 6, 2021

josh146 Feb 8, 2021

josh146 Feb 8, 2021

chaserileyroberts Feb 8, 2021

		def test_non_backprop_error(self):
		"""Test that an error is raised in tape mode if the diff method is not backprop"""

		# Easiest way to test object is a device array instead of np.array
		assert "DeviceArray" in res.__repr__()

Jax all devices #1065

Jax all devices #1065

Conversation

chaserileyroberts commented Feb 3, 2021 • edited Loading

josh146 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

josh146 Feb 4, 2021 • edited Loading

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

codecov bot commented Feb 4, 2021 • edited Loading

Codecov Report

chaserileyroberts commented Feb 4, 2021

chaserileyroberts commented Feb 6, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

chaserileyroberts commented Feb 3, 2021 •

edited

Loading

josh146 Feb 4, 2021 •

edited

Loading

codecov bot commented Feb 4, 2021 •

edited

Loading