Make numerical representations of operators differentiable #1749

glassnotes · 2021-10-13T19:13:19Z

Context: Operator overhaul

Description of the Change: Converts the numerical representation of parametrized operators to use qml.math throughout, and return matrix representations in the same interface in which input parameters are specified.

Benefits: Operator representations are fully differentiable.

Possible Drawbacks: TODO

Related GitHub Issues: N/A

github-actions · 2021-10-13T19:13:40Z

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.

josh146 · 2021-10-14T06:14:21Z

@glassnotes one good way to test these changes --- if you go into devices/default_qubit_<interface>.py, you'll see that there is a class attribute self.parametric_ops.

This instructs the backprop devices to ignore Operator.matrix for parametric ops, and instead they use an equivalent function inside the file devices/<interface>_ops.py.

If you clear the self.parametric_ops list, these devices will instead query Operator.matrix directly in backprop mode 🙂

…tors-are

josh146 · 2021-10-14T16:08:57Z

pennylane/devices/default_qubit_jax.py

@@ -17,7 +17,6 @@
 import pennylane as qml
 from pennylane.operation import DiagonalOperation
 from pennylane.devices import DefaultQubit
-from pennylane.devices import jax_ops


josh146

mammoth task @glassnotes! 🦣

josh146 · 2021-10-14T16:09:51Z

pennylane/devices/default_qubit_jax.py

        if isinstance(unitary, DiagonalOperation):
            return unitary.eigvals



If the rest of the method is identical to default.qubit, we can probably delete it, and just rely on inheritance!

Good catch! I think this could be removed from a few of the devices.

josh146 · 2021-10-14T16:10:26Z

pennylane/math/single_dispatch.py

@@ -134,6 +134,7 @@ def _take_autograd(tensor, indices, axis=None):
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "arctan2"] = "tensorflow.math"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "diag"] = "tensorflow.linalg"
 ar.autoray._SUBMODULE_ALIASES["tensorflow", "kron"] = "tensorflow.experimental.numpy"
+ar.autoray._SUBMODULE_ALIASES["tensorflow", "moveaxis"] = "tensorflow.experimental.numpy"


this is the only one that needed to be added? I'm somewhat surprised (in a good way), I thought more would be needed

josh146 · 2021-10-14T16:11:18Z

pennylane/ops/qubit/parametric_ops.py

+        interface = qml.math.get_interface(theta)
+
+        c = qml.math.cos(theta / 2)
+        s = qml.math.sin(theta / 2)


very very minor suggestion, but perhaps the 1j* and -theta were included here previously as a speedup?

Tried with torch, 7 microseconds for the 2x2 cases and 15 for the 4x4 cases. I'll take it! ⏩ 👍

josh146 · 2021-10-14T16:11:54Z

pennylane/ops/qubit/parametric_ops.py

+        if len(qml.math.shape(theta)) > 0:
+            theta = theta[0]


wait, why does this line occur?

I might actually be able to remove these now; I'll check. In the original versions using qml.math.array(stuff, like=interface), I needed to include this for cases like

torch.tensor([0.3], requires_grad=True)

otherwise there would be an extra dimension in the returned matrix. But with stack this may no longer be a problem.

Looks like they are not covered by the test suite 🤔

pennylane/ops/qubit/parametric_ops.py

josh146 · 2021-10-14T16:21:04Z

pennylane/utils.py


-    return expanded_tensor.reshape((2 ** M, 2 ** M))
+    return qml.math.reshape(expanded_tensor, (2 ** M, 2 ** M))


woah, nice catch updating this. Which functions were calling this method?

it was in some of the functions in PauliRot and MultiRZ.

josh146

@glassnotes I just realized, but changes to default.qubit.autograd are missing :)

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

pennylane/ops/qubit/parametric_ops.py

josh146 · 2021-10-14T18:06:29Z

  File "/home/runner/work/pennylane/pennylane/pennylane/ops/qubit/qchem_ops.py", line 545, in _matrix
    e = cmath.exp(-1j * theta / 2)
TypeError: must be real number, not ArrayBox

looks like there are some additional QChem ops not converted yet? In particular, OrbitalRotation.

glassnotes · 2021-10-14T18:15:29Z

  File "/home/runner/work/pennylane/pennylane/pennylane/ops/qubit/qchem_ops.py", line 545, in _matrix
    e = cmath.exp(-1j * theta / 2)
TypeError: must be real number, not ArrayBox
looks like there are some additional QChem ops not converted yet? In particular, OrbitalRotation.

Oops, I missed those because they weren't included in the parametric_ops dictionary! They must be pretty new as well :) Thanks for the heads up!

josh146 · 2021-10-14T18:15:58Z

No worries! Sorry, I got a bit too invested in the CI, so I've already pushed a fix 😬 Waiting to see if the CI passes now

glassnotes · 2021-10-14T18:25:46Z

I've already pushed a fix

No complaints here 🎉 Thanks!

josh146 · 2021-10-14T18:29:46Z

All that remains are the Torch 'probabilities not summing to 1' issue 🤔

glassnotes · 2021-10-16T15:23:38Z

tests/devices/test_default_qubit_torch.py

@@ -1321,8 +1321,8 @@ def circuit(a, b):
            qml.CRX(b, wires=[0, 1])
            return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))

-        a = torch.tensor([-0.234], dtype=torch.float64, requires_grad=True)
-        b = torch.tensor([0.654], dtype=torch.float64, requires_grad=True)
+        a = torch.tensor(-0.234, dtype=torch.float64, requires_grad=True)


Removing the checks on the size of theta led to errors in test_backprop_gradient; I've thus adjusted this test, but I'm a bit confused because test_prob_differentiability creates the tensors in the same way and runs without issues.

josh146 · 2021-10-18T06:13:27Z

pennylane/ops/qubit/parametric_ops.py

+        # It might be that they are in different interfaces, e.g.,
+        # Rot(0.2, 0.3, tf.Variable(0.5), wires=0)
+        # So we need to make sure the matrix comes out having the right type
+        interfaces = [qml.math.get_interface(p) for p in [phi, theta, omega]]


@glassnotes you could replace this with

Suggested change

# It might be that they are in different interfaces, e.g.,

# Rot(0.2, 0.3, tf.Variable(0.5), wires=0)

# So we need to make sure the matrix comes out having the right type

interfaces = [qml.math.get_interface(p) for p in [phi, theta, omega]]

# It might be that they are in different interfaces, e.g.,

# Rot(0.2, 0.3, tf.Variable(0.5), wires=0)

# So we need to make sure the matrix comes out having the right type

interfaces = qml.math._multi_dispatch([phi, theta, omega])

_multi_dispatch was originally intended to be purely a private function, but since then I've realized it's kinda useful in general, so I might make it public.

What it does is that it iterates through a list of tensor-like objects, and decides the overall interface to choose. It also has some logic that 'prioritizes' certain interfaces over others.

E.g., if the list of tensors is [autograd, autograd, numpy, torch], it priorities torch over autograd.

josh146 · 2021-10-18T07:05:15Z

pennylane/ops/qubit/parametric_ops.py

    @classmethod
    def _matrix(cls, *params):


Suggested change

@classmethod

def _matrix(cls, *params):

@classmethod

@functools.lru_cache()

def _matrix(cls, *params):

I think this should lead to speedups, since the call to generate the matrix will be cached for the same input object (e.g., if theta has the same value and same type!)

You're right! First timing run is the original, second is with the caching. Beautiful 🤩

However this seems to cause trouble with the gradients in the Ising operators in the gradient tests for JAX/TF. Have you ever encountered this before?

def __hash__(self): if ops.Tensor._USE_EQUALITY and ops.executing_eagerly_outside_functions(): # pylint: disable=protected-access > raise TypeError("Variable is unhashable. " "Instead, use tensor.ref() as the key.") E TypeError: Variable is unhashable. Instead, use tensor.ref() as the key.

That's super weird 🤔 does qml.RX(tf.Variable(0.5), wires=0).matrix work okay?

No it does not :(

yeah, I just ran the test suite with lru_cache() turned on for only RX, this is what happened:

Huge shame :(

I wonder if we can turn it on programmatically for just autograd and torch?

josh146

Got the tests passing :) Interestingly, it looks like there are no tests for:

tf and PhaseShift
tf and ControlledPhaseShift
tf and PauliRot
tf and CRZ
tf and U1
tf and U2
tf and U3
tf and qml.utils.expand

which codecov is now picking up, since there are tf branches in those operations

glassnotes · 2021-10-18T15:32:43Z

@josh146 CI is passing and all changes are in, this is ready for re-review. Codecov is still complaining about an overall decrease in coverage, but is not indicating any specific lines have been missed, as far as I can see.

josh146 · 2021-10-18T16:29:00Z

Will go for a review now 🙂

It looks like the reduction in coverage is coming from the devices, weird 🤔 Maybe codecov is just confused since the number of lines in the device class has dropped...

josh146

Thanks @glassnotes, been really looking forward to this change!

josh146 · 2021-10-18T16:30:13Z

doc/releases/changelog-dev.md

+* All qubit operations have been re-written to use the `qml.math` framework
+  for internal classical processing and the generation of their matrix representations.
+  As a result these representations are now fully differentiable, and the
+  framework-specific device classes no longer need to maintain framework-specific
+  versions of these matrices.
+  [(#1749)](https://github.com/PennyLaneAI/pennylane/pull/1749)


💯 for avoiding 'interface' hahaha

josh146 · 2021-10-18T16:30:54Z

doc/releases/changelog-dev.md

+- The updates made to the operator matrix representations in order to
+  enable full differentiability in all interfaces now require that single
+  parameters be passed as rank-0 tensors in order for the matrix dimensions
+  to be correct (i.e., use `qml.RX(tf.Variable(0.3), wires=0)` instead of
+  `qml.RX(tf.Variable([0.3]), wires=0)`.
+  [(#1749)](https://github.com/PennyLaneAI/pennylane/pull/1749)  


I almost tempted not to include this here, since technically this isn't a breaking change (we never guaranteed this behaviour in earlier versions)

I actually wasn't sure about it; let's just remove it then if the behaviour was never guaranteed anyways.

pennylane/ops/qubit/matrix_ops.py

pennylane/ops/qubit/parametric_ops.py

josh146 · 2021-10-18T16:46:34Z

pennylane/ops/qubit/parametric_ops.py

+
+        if qml.math.get_interface(theta) == "tensorflow":
+            theta = qml.math.cast_like(theta, 1j)
+
+        exp_part = qml.math.exp(-0.5j * theta)
+
+        return qml.math.diag([1, 1, exp_part, qml.math.conj(exp_part)])


This is very minor, no action needed, but I just noticed that this could simply be written as

def _matrix(cls, *params): return qml.math.diag(cls._eigvals(*params))

since it is diagonal. Although, I don't know if there are any unintentional issues associated with this.

josh146 · 2021-10-18T16:48:34Z

pennylane/ops/qubit/qchem_ops.py

+            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0],
+            [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0],
+        ]


No action needed here, but I wonder if this could be done via qml.math.scatter_element_add.

U = qml.math.eye(16, like=interface) indices = [(3, 3, 12, 12), (3, 12, 3, 12)] values = qml.math.stack([c - 1.0, -s, s, c - 1.0]) U = qml.math.scatter_element_add(U, indices, values)

josh146 · 2021-10-18T16:57:01Z

tests/devices/test_default_qubit_tf.py

+        obtained_mat = qml.utils.expand(tf_mat, wires, list(range(4)))
+
+        assert qml.math.get_interface(obtained_mat) == "tensorflow"
+        assert qml.math.allclose(qml.math.unwrap(obtained_mat), expected_mat)


Thanks for adding these @glassnotes!

…tors-are

doc/releases/changelog-dev.md

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

…tors-are

Convert most parametric ops to use qml.math.

1960665

Run black.

8a1b084

Olivia Di Matteo and others added 7 commits October 14, 2021 11:25

Use stack for everything. Update helper functions.

4c26d23

Remove operations that have been converted.

a93d794

Update single excitations and devices.

56b4d9e

Black utils.

9f3231e

Update all devices to use the default operations.

0b8ece6

Remove empty dict.

b460a68

Merge branch 'master' into sc-7223-numerical-representations-of-opera…

0c8779b

…tors-are

glassnotes requested a review from josh146 October 14, 2021 15:55

Merge branch 'master' into sc-7223-numerical-representations-of-opera…

64a9c25

…tors-are

josh146 reviewed Oct 14, 2021

View reviewed changes

josh146 requested changes Oct 14, 2021

View reviewed changes

josh146 reviewed Oct 14, 2021

View reviewed changes

Olivia Di Matteo and others added 2 commits October 14, 2021 12:47

Remove get unitary matrix from devices where not needed.

c1a8de5

Apply suggestions from code review

91c924f

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

josh146 reviewed Oct 14, 2021

View reviewed changes

pennylane/ops/qubit/parametric_ops.py Outdated Show resolved Hide resolved

Update pennylane/ops/qubit/parametric_ops.py

22e14df

josh146 added 3 commits October 15, 2021 02:12

fix

ca6bad7

fix

5295a0a

linting

7c17da9

josh146 added 2 commits October 15, 2021 02:34

fix torch errors

9872936

fix torch errors

91be029

Fix incorrect parameters in tensorflow section.

9ab79a7

glassnotes commented Oct 16, 2021

View reviewed changes

josh146 reviewed Oct 18, 2021

View reviewed changes

josh146 added 2 commits October 18, 2021 14:16

merge master

9acb2f4

fix tests

d9d58ca

josh146 reviewed Oct 18, 2021

View reviewed changes

fix tests

e504cfd

josh146 reviewed Oct 18, 2021

View reviewed changes

josh146 and others added 6 commits October 18, 2021 15:34

fix

438bc17

Fix matrix op and add speedup with 1j mult.

4a78be2

Add missing tests.

82b1731

Fix tensor dimensions.

716f811

Update changelog. Add missing eigval test.

6583f1d

Run black.

ed16fa0

glassnotes requested a review from josh146 October 18, 2021 15:30

josh146 approved these changes Oct 18, 2021

View reviewed changes

Merge branch 'master' into sc-7223-numerical-representations-of-opera…

29a6da2

…tors-are

glassnotes commented Oct 18, 2021

View reviewed changes

doc/releases/changelog-dev.md Outdated Show resolved Hide resolved

glassnotes and others added 2 commits October 18, 2021 13:21

Apply suggestions from code review

83b9b0b

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

Merge branch 'master' into sc-7223-numerical-representations-of-opera…

310c2ec

…tors-are

josh146 merged commit 1bca151 into master Oct 18, 2021

josh146 deleted the sc-7223-numerical-representations-of-operators-are branch October 18, 2021 18:15

glassnotes mentioned this pull request Oct 18, 2021

Add differentiation support to the qml.transforms.get_unitary_matrix transform #1682

Closed

This was referenced Oct 25, 2021

Remove complex warning from operator matrices #1802

Merged

[BUG] Operator matrices should be defined using holomorphic functions #1819

Open

This was referenced Nov 22, 2021

[WIP] Fix parametric ops GPU compatibility #1921

Closed

Fix parametric ops GPU compatibility (v0.19.1) #1927

Merged

antalszava mentioned this pull request Aug 17, 2022

Add PSWAP operation [unitaryhack] #2667

Merged

		if isinstance(unitary, DiagonalOperation):
		return unitary.eigvals


		return expanded_tensor.reshape((2 M, 2 M))
		return qml.math.reshape(expanded_tensor, (2 M, 2 M))

Make numerical representations of operators differentiable #1749

Make numerical representations of operators differentiable #1749

Conversation

glassnotes commented Oct 13, 2021 • edited Loading

github-actions bot commented Oct 13, 2021

josh146 commented Oct 14, 2021

Choose a reason for hiding this comment

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

glassnotes Oct 18, 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

josh146 left a comment

Choose a reason for hiding this comment

josh146 commented Oct 14, 2021

glassnotes commented Oct 14, 2021

josh146 commented Oct 14, 2021

glassnotes commented Oct 14, 2021

josh146 commented Oct 14, 2021

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

josh146 left a comment

Choose a reason for hiding this comment

glassnotes commented Oct 18, 2021

josh146 commented Oct 18, 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

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

glassnotes commented Oct 13, 2021 •

edited

Loading

glassnotes Oct 18, 2021 •

edited

Loading

josh146 commented Oct 18, 2021 •

edited

Loading