Add JAX integration tests #1685

josh146 · 2021-09-24T09:14:42Z

Context: The JAX interface is missing many of the integration tests that the other interfaces have. This PR adds these integration tests in, and makes note of where the JAX interface may be lacking in feature parity.

Description of the Change:

Adds a JAX integration test to test_gradient_transform.py (including JIT tests).
Adds QNode integration tests tests/interfaces/test_batch_jax_qnode.py.
Adds a batch transform integration test test_batch_transform.py.
Modifies qml.math.get_trainable_indices() to correctly return results when JAX is performing a forward-only computation. This is done by making the following change:
- On forward-only passes, all DeviceArray objects are treated as trainable.
- On forwards+backwards passes, only jax.core.Tracer objects are treated as trainable (which matches the jax.grad(cost, argnum=...) argument).
This change allows the metric tensor/gradient transform functions --- which apply in forward-only mode but require knowledge of trainable parameters --- to apply to QNodes when using JAX.

This is required because, since JAX does not have a method of specifying trainable parameters on the forward pass, perviously gradient transforms would simply register no trainable parameters on forward passes, and return an empty list as a result! Paradoxically, differentiating the gradient transform would work fine, since the trainable parameters are now registered.

Benefits:

Better tests for JAX, and a better idea of what works, and what doesn't.
Gradient transforms, and the metric tensor, now works for the JAX interface.

Possible Drawbacks:

Several areas were noticed were JAX usage resulted in errors or issues, unlike other interfaces:

Ragged QNodes in backprop mode. E.g., return qml.expval(qml.PauliZ(0)), qml.probs(wires=[1]). This appears to be because line 230 in QubitDevice (results = self._asarray(results)) fails.
Hamiltonian expansion of expval(H) when using finite shots fails. When using finite shots, expval(H) results in the tape being expanded to multiple tapes, each one having a vector-valued output. Since the JAX parameter-shift interface does not support vector-valued tapes, an error occurs.
JIT mode with the adjoint method. Since the adjoint method is not using host_callback in the jax.py interface, it results in a failure attempting to JIT the QNode.
JIT mode with jax.grad(cost, argnums=...) where argnums is a subset of allowed arguments, e.g., argnums=[0, 2]. In this scenario, the JAX interface is only passing unwrapped trainable parameters to the host_callback - as a result, non-trainable parameters are remaining as jax.core.Tracer objects on the tape, which cannot be understood by the device.

Related GitHub Issues: n/a

github-actions · 2021-09-24T09:15:04Z

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.

codecov · 2021-09-28T12:59:41Z

Codecov Report

Merging #1685 (db973d5) into master (b352691) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #1685   +/-   ##
=======================================
  Coverage   98.90%   98.90%           
=======================================
  Files         206      206           
  Lines       15388    15396    +8     
=======================================
+ Hits        15219    15227    +8     
  Misses        169      169

Impacted Files	Coverage Δ
pennylane/math/multi_dispatch.py	`100.00% <100.00%> (ø)`
pennylane/math/utils.py	`100.00% <100.00%> (ø)`
pennylane/transforms/batch_transform.py	`100.00% <100.00%> (ø)`

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 b352691...db973d5. Read the comment docs.

josh146 · 2021-09-28T17:14:58Z

pennylane/math/utils.py

@@ -279,6 +279,6 @@ def requires_grad(tensor, interface=None):
    if interface == "jax":
        import jax

-        return isinstance(tensor, jax.interpreters.ad.JVPTracer)
+        return isinstance(tensor, jax.core.Tracer)


jax.core.Tracer is the original parent class, so this is a lot safer :) There are cases I discovered where JAX will use tracers that aren't JVPTracer.

josh146 · 2021-09-28T17:15:22Z

tests/collections/test_collections.py

@@ -400,7 +400,7 @@ def test_dot_product_qnodes_tensor(self, qnodes, interface, tf_support, torch_su
            coeffs = coeffs.numpy()

        expected = np.dot(qcval, coeffs)
-        assert np.all(res == expected)
+        assert np.allclose(res, expected)


For some reason, this test was failing for me on CI (but not locally) 🤔

antalszava

Overall looks good to me! 🙂 Left some questions, but no major blocker.

If I recall correctly, we plan to discontinue support for QNodes with ragged outputs (right? 🤔). Would that affect the failing case of Ragged QNodes in backprop mode?

doc/releases/changelog-dev.md

antalszava · 2021-10-19T14:55:15Z

pennylane/math/multi_dispatch.py

+    if interface == "jax":
+        import jax
+
+        if not any(isinstance(v, jax.core.Tracer) for v in values):


This looks like a great change to have 🥇

How come it's placed here, instead of into the JAX branch of requires_grad?

I think it needs to be here, since the not any check can only be done here, it cannot be done inside the requires_grad check (which only checks a single tensor at a time) 🤔

I could be wrong though, let me know if you see a way around this!

No, I think you're right. 🤔 At least nothing else comes to mind that we could use here.

pennylane/transforms/batch_transform.py

tests/gradients/test_finite_difference.py

antalszava · 2021-10-19T15:39:40Z

tests/interfaces/test_batch_jax_qnode.py

+            qml.CNOT(wires=[0, 1])
+            return qml.expval(qml.PauliY(1))
+
+        res = jax.grad(cost_fn, argnums=[0, 1])(a, b, shots=30000)


How does the test case make sure that we had shots=30000 instead of shots=100? Would the deviation from the expected analytic result be even bigger than 0.1 with shots=100?

Yes, it is not the most ideal way of testing this, but I spent a while and couldn't come up with anything better 🤔

The other interfaces use shots=[(1, 1000)], which is nicer since the output shape changes. However, it can't be used for JAX, since JAX only supports scalar outputs :(

antalszava · 2021-10-19T15:41:03Z

tests/interfaces/test_batch_jax_qnode.py

+        assert spy.call_args[1]["gradient_fn"] is qml.gradients.param_shift
+
+        # if we set the shots to None, backprop can now be used
+        cost_fn(a, b, shots=None)


Should we then pass in diff_method="param-shift" if we really wanted to use parameter shift with shots=None?

yep! The reason the internal diff_method is changing here is because, by default, diff_method="best". If you instead set diff_method="param-shift", then it will not change dynamically.

antalszava · 2021-10-19T15:44:54Z

tests/interfaces/test_batch_jax_qnode.py

+
+        if diff_method not in {"backprop"}:
+            pytest.skip("Test only supports backprop")


Would this have any effect at this point?

Suggested change

if diff_method not in {"backprop"}:

pytest.skip("Test only supports backprop")

I think I would prefer to leave it in, since we may have support for vector valued QNodes in parameter-shift mode at some point in the future!

Right! Would we only test assert res.dtype is np.dtype("complex128") for non-backprop diff methods? That seems to be the only statement before skipping the test.

tests/interfaces/test_batch_jax_qnode.py

tests/transforms/test_metric_tensor.py

Co-authored-by: antalszava <antalszava@gmail.com>

antalszava

Looks good to me! 💯 Thank you for adding these tests to check the JAX interface. 😍

Just double-checking this one:

If I recall correctly, we plan to discontinue support for QNodes with ragged outputs (right? thinking). Would that affect the failing case of Ragged QNodes in backprop mode?

This is just to understand better what the priority of the drawbacks would be that were identified in the PR description. 🙂

josh146 · 2021-10-20T05:11:07Z

Would that affect the failing case of Ragged QNodes in backprop mode?

Yes! Hopefully, this would allow these QNodes to work with JAX, once we make that change 🙂

Add JAX integration tests

43b1fe0

add test

b6e78c6

josh146 mentioned this pull request Sep 25, 2021

Add GateFabric Template #1687

Merged

antalszava and others added 6 commits September 27, 2021 19:43

Merge branch 'master' into batch-qnode-jax

5e7d680

Merge branch 'master' into batch-qnode-jax

7f7a65e

test

653a546

update

92cb81d

update

eedeabb

Merge branch 'master' into batch-qnode-jax

c180bf2

josh146 added 3 commits September 29, 2021 01:07

Add jit tests

04dfcf5

Merge branch 'master' into batch-qnode-jax

61a60ef

more

31eae7c

josh146 requested a review from antalszava September 28, 2021 17:12

josh146 added the review-ready 👌 PRs which are ready for review by someone from the core team. label Sep 28, 2021

josh146 commented Sep 28, 2021

View reviewed changes

josh146 added 3 commits September 29, 2021 14:12

Merge branch 'master' into batch-qnode-jax

c9eab59

Merge branch 'master' into batch-qnode-jax

a8db871

fix

b8e511d

antalszava reviewed Oct 19, 2021

View reviewed changes

josh146 and others added 2 commits October 20, 2021 00:08

Apply suggestions from code review

4d50817

Co-authored-by: antalszava <antalszava@gmail.com>

Merge branch 'master' into batch-qnode-jax

eb75c4c

josh146 requested a review from antalszava October 19, 2021 16:10

josh146 and others added 2 commits October 20, 2021 00:23

Merge branch 'master' into batch-qnode-jax

d09a382

Merge branch 'master' into batch-qnode-jax

db973d5

antalszava approved these changes Oct 19, 2021

View reviewed changes

josh146 merged commit f20c19d into master Oct 20, 2021

josh146 deleted the batch-qnode-jax branch October 20, 2021 05:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add JAX integration tests #1685

Add JAX integration tests #1685

josh146 commented Sep 24, 2021 •

edited

Loading

github-actions bot commented Sep 24, 2021

codecov bot commented Sep 28, 2021 •

edited

Loading

josh146 Sep 28, 2021

josh146 Sep 28, 2021

antalszava left a comment

antalszava Oct 19, 2021

josh146 Oct 19, 2021

antalszava Oct 19, 2021

antalszava Oct 19, 2021

josh146 Oct 19, 2021

antalszava Oct 19, 2021

josh146 Oct 19, 2021

antalszava Oct 19, 2021

josh146 Oct 19, 2021

antalszava Oct 19, 2021

antalszava left a comment

josh146 commented Oct 20, 2021


		if diff_method not in {"backprop"}:
		pytest.skip("Test only supports backprop")

Add JAX integration tests #1685

Add JAX integration tests #1685

Conversation

josh146 commented Sep 24, 2021 • edited Loading

github-actions bot commented Sep 24, 2021

codecov bot commented Sep 28, 2021 • edited Loading

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

antalszava 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

Choose a reason for hiding this comment

antalszava left a comment

Choose a reason for hiding this comment

josh146 commented Oct 20, 2021

josh146 commented Sep 24, 2021 •

edited

Loading

codecov bot commented Sep 28, 2021 •

edited

Loading