Add support for batch execution to `qml.metric_tensor` #1638

josh146 · 2021-09-10T10:01:31Z

Context: With the low-level differentiable batch-execution pipeline now available in PennyLane, and the availability of @qml.batch_transform, we can now start porting our existing batch-tape transforms to take advantage of this framework.

Description of the Change:

When using qml.math.get_trainable_indices(), only NumPy arrays with requires_grad=True are taken as trainable when using Autograd. This is a requirement for batch transformations like the metric tensor to work correctly.
When using the QNode in backpropagation, trainable parameter indices are computed and stored. Previously, a backpropagation QNode would simply mark all parameters as trainable. This extra information is needed for the metric tensor.
The metric tensor has been converted into a batch transformation, that accepts both tapes and QNodes as input. This makes the old metric_tensor_tape function irrelevant, and it has been removed. The tests have been likewise updated.
Previously, qml.metric_tensor(qnode)(*args, **kwargs) would only return the metric tensor with respect to gate arguments, and ignore any classical processing inside the QNode, even very trivial classical processing such as parameter permutation. This would lead to many reported user bugs, such as QNGOptimizer returns TypeError when step method called #1154. In this new framework, the metric tensor now takes into account classical processing, and returns the metric tensor with respect to QNode arguments, not simply gate arguments.

To revert to the previous behaviour of returning the metric tensor with respect to gate arguments, qml.metric_tensor(qnode, hybrid=False) can be passed.

Benefits:

The qml.metric_tensor() function now makes use of batch execution for submission of circuits required for metric tensor computation.
The qml.metric_tensor() function now takes into account classical computation inside the QNode, for example, if gate arguments are repeated or permuted.
QNodes in backpropagation mode correctly track trainable parameters, fixing an issue that was long part of our test suite.

Possible Drawbacks: n/a

Related GitHub Issues: Closes #1154

github-actions · 2021-09-10T10:01:47Z

Hello. You may have forgotten to update the changelog!
Please edit .github/CHANGELOG.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-10T10:10:47Z

Codecov Report

Merging #1638 (44ec65b) into master (1f89139) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #1638   +/-   ##
=======================================
  Coverage   99.15%   99.15%           
=======================================
  Files         196      196           
  Lines       14294    14318   +24     
=======================================
+ Hits        14173    14197   +24     
  Misses        121      121

Impacted Files	Coverage Δ
pennylane/math/utils.py	`100.00% <100.00%> (ø)`
pennylane/qnode.py	`98.73% <100.00%> (+0.03%)`	⬆️
pennylane/transforms/__init__.py	`100.00% <100.00%> (ø)`
pennylane/transforms/classical_jacobian.py	`100.00% <100.00%> (ø)`
pennylane/transforms/metric_tensor.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 1f89139...44ec65b. Read the comment docs.

josh146 · 2021-09-13T11:51:46Z

[ch8957]

glassnotes

@josh146 thanks this is a great update 🎉 I've left a number of questions within

.github/CHANGELOG.md

glassnotes · 2021-09-14T14:40:16Z

.github/CHANGELOG.md

+           [0.        , 0.28750832]])
+    ```
+
+    To revert to the previous behaviour of returning the metric tensor with respect to gate


This is great 💯

Totally optional suggestion: Is there a more descriptive keyword than hybrid we could use here? It's not immediately clear just from looking at the arguments what it means. Maybe... use_gate_args, or something like this?

I agree that this example is looking great and there could be a more descriptive keyword here

I admit, I used hybrid only because the qml.gradients module uses it 🙈

I have a slight preference for keeping it consistent (for now), while opening up an issue to replace the keyword name in both places?

Options I can think of are:

quantum_only

include_classical

circuit_only

?

I'm a fan of include_classical :)

glassnotes · 2021-09-14T14:44:06Z

pennylane/transforms/metric_tensor.py

+from .batch_transform import batch_transform
+
+
+SUPPORTED_OPS = ["RX", "RY", "RZ", "PhaseShift"]


Is it necessary to support PhaseShift, since its decomposition is in terms of RZ?

For the metric tensor, we can support any single-parameter operation that has a generator, so this limits us to these four (I believe?)

As to why PhaseShift; on the off-chance someone does have a circuit with PhaseShift, I think this will lead to a slight reduction in overhead, since the expansion can be avoided.

glassnotes · 2021-09-14T14:44:22Z

pennylane/transforms/metric_tensor.py

+    new_tape = tape.expand(depth=2, stop_at=_stopping_critera)
+    params = new_tape.get_parameters(trainable_only=False)


Why depth 2?

No reason except the existing metric tensor code uses depth 2!

pennylane/transforms/metric_tensor.py

glassnotes · 2021-09-14T14:49:49Z

pennylane/transforms/metric_tensor.py

@@ -152,7 +187,7 @@ def metric_tensor_tape(tape, diag_approx=False, wrt=None):
        # to measure in the basis of the parametrized layer generators.
        with tape.__class__() as layer_tape:
            for op in queue:
-                op.queue()
+                qml.apply(op)


Couldn't help myself!

glassnotes · 2021-09-14T14:51:41Z

pennylane/transforms/metric_tensor.py

+        if is_square and qml.math.allclose(cjac, qml.numpy.eye(cjac.shape[0])):
+            # Classical Jacobian is the identity. No classical processing
+            # is present inside the QNode.
+            return mt


So this covers the case where even if hybrid=True, no actual classical processing happened so we just return the mt w.r.t. the gate arguments?

glassnotes · 2021-09-14T14:54:52Z

tests/transforms/test_metric_tensor.py

@@ -58,9 +57,6 @@ def circuit(weights):
        assert tapes[2].operations[0].data == [1]
        assert tapes[2].operations[1].data == [2]

-        result = qml.metric_tensor(circuit)(params)


Why isn't the result still being tested?

I think it's superfluous --- this test is testing that the rotation gate is correctly decomposed, and the execution of the compiled circuit feels a bit out of scope

glassnotes · 2021-09-14T15:00:23Z

pennylane/math/utils.py

-        # Currently, in the Autograd interface, we assume
-        # that all objects are differentiable by default.
-        return getattr(tensor, "requires_grad", True)
+        return getattr(tensor, "requires_grad", False)


Is this setting the default value of requires_grad everywhere to False? If so that seems like a major change, it should maybe be added separately to the CHANGELOG and docs

Luckily not, it turns out that this function is not used anywhere important yet 😆 I was super nervous about making this change, expecting tonnes of tests to fail.

But good idea to mention this in the changelog

Have added this to the changelog!

glassnotes · 2021-09-14T15:01:45Z

tests/transforms/test_metric_tensor.py

+            qml.RX(a[1], wires=0)
+            qml.RY(a[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PhaseShift(b, wires=1)


There is no classical processing of the arguments in this QNode

There is, but it's subtle. The classical processing function is

f: ([a0, a1], b) -> (a1, a0, b)

So the classical Jacobians will be a permutation matrix and an identity matrix:

classical_jacobian(circuit)(a, b) == ([[0, 1], [1, 0]], [[1]])

Have added a comment here :)

Wow I totally missed that! Thanks, the comment is helpful 💯

Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

anthayes92 · 2021-09-14T17:44:37Z

.github/CHANGELOG.md

+           [0.        , 0.28750832]])
+    ```
+
+    To revert to the previous behaviour of returning the metric tensor with respect to gate


I agree that this example is looking great and there could be a more descriptive keyword here

anthayes92 · 2021-09-14T18:04:35Z

pennylane/transforms/metric_tensor.py


+    >>> grad_fn = qml.grad(lambda x: met_fn(x)[3, 2])
+    >>> grad_fn(weights)
+    array([[ 0.04867729, -0.00049502,  0.        ],


Might be worth mentioning that we can find the Jacobian but agreed that the gradient example is more digestible here!

pennylane/transforms/metric_tensor.py

anthayes92 · 2021-09-14T18:24:43Z

pennylane/transforms/metric_tensor.py


-        qnode.construct(args, kwargs)
+                for c in cjac:
+                    if c is not None:


how can None values appear in the classical jacobian?

depending on the autodiff framework, None appears if the corresponding QNode argument is non-differentiable:

@qml.qnode(dev) def circuit(x, y): qml.RX(2*x, wires=0) qml.RY(y**3, wires=0) return qml.expval(qml.PauliZ(0) >>> x = np.array(0.4, requires_grad=True) >>> y = np.array(0.4, requires_grad=False)

anthayes92 · 2021-09-14T18:36:23Z

Another fine addition @josh146! Just a couple of minor suggestions and questions

Co-authored-by: anthayes92 <34694788+anthayes92@users.noreply.github.com>

pennylane/transforms/metric_tensor.py

josh146 · 2021-09-15T12:13:46Z

Thanks @glassnotes and @anthayes92! Your feedback was extremely helpful - I've incorporated it in 9fdf474..9f01b05 (you can view the combined diff here).

Note that the behaviour of qml.transforms.classical_jacobian changed in master since the code review, so there is a bit of a change to the classical_jacobian logic and tests to fix bugs that was causing failing tests.

glassnotes

Thanks for the clarifications @josh146 , looks good to go ⭐

glassnotes · 2021-09-15T12:32:14Z

.github/CHANGELOG.md

+           [0.        , 0.28750832]])
+    ```
+
+    To revert to the previous behaviour of returning the metric tensor with respect to gate


I'm a fan of include_classical :)

glassnotes · 2021-09-15T12:33:35Z

.github/CHANGELOG.md

+  Please use the `qml.metric_tensor` transform instead.
+  [(#1638)](https://github.com/PennyLaneAI/pennylane/pull/1638)
+
+- The utility function `qml.math.requires_grad` now returns `True` when using Autograd


This is not going in 0.18 right? I'm just thinking now whether we will have to update some of the demos / other docs to incorporate this.

Nope, this is being merged into master only/v0.19-dev

glassnotes · 2021-09-15T12:39:09Z

tests/transforms/test_metric_tensor.py

+            qml.RX(a[1], wires=0)
+            qml.RY(a[0], wires=0)
+            qml.CNOT(wires=[0, 1])
+            qml.PhaseShift(b, wires=1)


Wow I totally missed that! Thanks, the comment is helpful 💯

josh146 added 8 commits September 10, 2021 15:41

Bug fixes for batch execution

b2c4baf

more tests

317514b

more tests

a937d4d

more tests

c6dc629

more tests

0a5756a

more tests

490f106

changelog

fa10491

Add metric tensor

d70f0c8

fixes

e0b45f5

josh146 changed the title ~~[WIP] Batch metric tensor~~ Add support for batch execution to qml.metric_tensor Sep 10, 2021

josh146 added 7 commits September 10, 2021 22:42

add test

523b765

Merge branch 'master' into batch-bug-fixes

f8af2d2

update

377d8f3

remove xfail

7c400cc

Merge branch 'master' into batch-bug-fixes

57f83a0

Merge branch 'master' into batch-bug-fixes

0852693

Merge branch 'master' into batch-bug-fixes

852f2fe

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

Merge branch 'batch-bug-fixes' into batch-metric-tensor

63f2990

Base automatically changed from batch-bug-fixes to v0.18.0-rc0 September 13, 2021 14:25

josh146 changed the base branch from v0.18.0-rc0 to master September 13, 2021 14:28

josh146 added 3 commits September 13, 2021 22:28

Merge branch 'master' into batch-metric-tensor

337295f

linting:

626450e

coverage

258e288

josh146 requested review from glassnotes and anthayes92 September 13, 2021 18:09

Merge branch 'master' into batch-metric-tensor

3387d59

add changelog entry

9fdf474

glassnotes requested changes Sep 14, 2021

View reviewed changes

josh146 and others added 2 commits September 15, 2021 00:04

Apply suggestions from code review

22716aa

Co-authored-by: Olivia Di Matteo <2068515+glassnotes@users.noreply.github.com>

Merge branch 'master' into batch-metric-tensor

976a58c

anthayes92 approved these changes Sep 14, 2021

View reviewed changes

Update pennylane/transforms/metric_tensor.py

e787e82

Co-authored-by: anthayes92 <34694788+anthayes92@users.noreply.github.com>

josh146 commented Sep 15, 2021

View reviewed changes

pennylane/transforms/metric_tensor.py Outdated Show resolved Hide resolved

josh146 added 3 commits September 15, 2021 14:39

Update pennylane/transforms/metric_tensor.py

8f1552f

more tests

4b8249c

more tests

9f01b05

josh146 requested a review from glassnotes September 15, 2021 12:13

glassnotes approved these changes Sep 15, 2021

View reviewed changes

Merge branch 'master' into batch-metric-tensor

44ec65b

josh146 merged commit 17b1551 into master Sep 15, 2021

josh146 deleted the batch-metric-tensor branch September 15, 2021 13:30

		from .batch_transform import batch_transform


		SUPPORTED_OPS = ["RX", "RY", "RZ", "PhaseShift"]

		new_tape = tape.expand(depth=2, stop_at=_stopping_critera)
		params = new_tape.get_parameters(trainable_only=False)

Add support for batch execution to qml.metric_tensor #1638

Add support for batch execution to qml.metric_tensor #1638

Conversation

josh146 commented Sep 10, 2021 • edited Loading

github-actions bot commented Sep 10, 2021

codecov bot commented Sep 10, 2021 • edited Loading

Codecov Report

josh146 commented Sep 13, 2021

glassnotes 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

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

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

anthayes92 commented Sep 14, 2021

josh146 commented Sep 15, 2021

glassnotes 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

Add support for batch execution to `qml.metric_tensor` #1638

Add support for batch execution to `qml.metric_tensor` #1638

josh146 commented Sep 10, 2021 •

edited

Loading

codecov bot commented Sep 10, 2021 •

edited

Loading