Allow specifying a subset of parameters to estimate the jacobian by passing the `argnum` kwarg #1371

antalszava · 2021-05-29T01:21:06Z

Context:
#958

Description of the Change:

Adds the argnum keyword argument to the JacobianTape.jacobian method inspired by the argnum kwarg of the qml.grad function. Using this keyword argument the trainable parameters to use for estimating the jacobian can be specified.

The logic is organized into a separate private method to allow easy unit testing and to prepare for further future extensions.

Note: as this is an addition to the JacobianTape, this option is only applicable for the differentiation methods that depend on such an object and its jacobian method: "parameter-shift", "finite-diff" and "reversible".

Benefits:

Jacobians can now be estimated by only using a subset of the trainable parameters.

Possible Drawbacks:
N/A

Related GitHub Issues:
Closes #958

github-actions · 2021-05-29T01:21:23Z

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-05-29T01:33:22Z

Codecov Report

Merging #1371 (c44bedc) into master (9a52aa1) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #1371   +/-   ##
=======================================
  Coverage   98.19%   98.19%           
=======================================
  Files         157      157           
  Lines       11713    11729   +16     
=======================================
+ Hits        11501    11517   +16     
  Misses        212      212

Impacted Files	Coverage Δ
pennylane/qnode.py	`97.85% <ø> (ø)`
pennylane/tape/jacobian_tape.py	`98.16% <100.00%> (+0.14%)`	⬆️

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 9a52aa1...c44bedc. Read the comment docs.

antalszava · 2021-05-29T01:57:05Z

tests/tape/test_jacobian_tape.py

@@ -454,6 +482,55 @@ def test_single_expectation_value(self, tol):
        expected = np.array([[-np.sin(y) * np.sin(x), np.cos(y) * np.cos(x)]])
        assert np.allclose(res, expected, atol=tol, rtol=0)

+    def test_single_expectation_value_with_num_params_all(self, tol):


This and the following test are using the test_single_expectation_value case to make sure that parameter selection works fine.

tests/tape/test_jacobian_tape.py

albi3ro · 2021-05-31T14:23:49Z

Maybe instead of providing num_params to tape.jacobian, it accepts a list of argnum like qml.grad. The choice of argnum for the random choice would occur upstream somewhere. This would make tape.jacobian much more versatile and useful as well, providing many more opportunities for speed-up.

albi3ro

Maybe we could instead put this functionality into _grad_method_validation? Non-selected parameters simply get their "grad_method" set to 0. This would keep the logic in tape.jacobian simpler for an already complicated function.

albi3ro · 2021-05-31T17:38:08Z

pennylane/tape/jacobian_tape.py

+        if num_params is None:
+            return enumerate(diff_methods)
+
+        if len(diff_methods) < num_params or num_params < 1:


I suggest breaking these two cases into separate warnings.

Then, when len(diff_methods) < num_params we can return enumerate(diff_methods) as demonstrated, but return an empty array when num_params < 1. I think that would be more intuitive behaviour.

Fair point! With the changed logic, the first doesn't make too much sense anymore, as there's a new error now. The second condition is treated separately and returns an empty list.

albi3ro · 2021-05-31T17:50:43Z

pennylane/tape/jacobian_tape.py

@@ -393,6 +395,39 @@ def hessian_pd(self, i, j, params, **options):
        """
        raise NotImplementedError

+    @staticmethod
+    def _choose_params_with_methods(diff_methods, num_params):


Should there be a way to specify the seed for reproducibility?

Could we do that via random.seed(my_seed) in the test? Note that here the random library was used for sampling (though open for change if another way or using numpy could be preferred).

albi3ro · 2021-05-31T17:53:29Z

pennylane/tape/jacobian_tape.py

+                computing the jacobian. Specifies the number of parameters to sample.
+
+        Returns:
+            object or list: map of the trainable parameter indices and


Suggested change

object or list: map of the trainable parameter indices and

generator or object or list: map of the trainable parameter indices and

Was this a suggestion for the enumerate case? Changing object to enumerate here, though not sure if we also have a generator in other cases.

pennylane/tape/jacobian_tape.py

josh146

Thanks @antalszava! Great tests, and I'm happy with this PR in it's current form (once @albi3ro's comments are addressed).

@albi3ro makes a good point regarding choosing the random number of parameters higher up in the stack. Currently, the Jacobian method computes the derivative based on the tape.trainable_params, which is set by the interface. This is worth considering, but should not be a blocker to merging this PR; we can continue to improve this functionality in later PRs.

.github/CHANGELOG.md

josh146 · 2021-05-31T16:11:21Z

.github/CHANGELOG.md

@@ -2,6 +2,9 @@

 <h3>New features since last release</h3>

+* The number of parameters used to estimate the jacobian can now be specified.
+  [(#1371)](https://github.com/PennyLaneAI/pennylane/pull/1371)


I suggest adding an example here, as this is quite a user-facing feature!

Sure! Added

Co-authored-by: Christina Lee <christina@xanadu.ai>

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

antalszava · 2021-06-08T18:04:37Z

Thank you for the review comments! 🙂 Not sure, if _grad_method_validation would be a good place exactly as that would serve the purpose of validation only, but would be happy to move the logic around in a later addition. For now, might leave it at its current place if that's okay.

albi3ro

Reading this right now, I'm confused about "number of parameters" versus "which parameters". argnum to me connotes the usage in qml.grad . Before I continue to review, I'd like that distinction to be clearer.

antalszava · 2021-06-09T01:04:09Z

Hi @albi3ro, with the latest behaviour of the code, I would think that we have the same behaviour as with qml.grad: argnum is the argument to specify the indices of the parameters wrt. which we compute/estimate the jacobian. Do we have the same understanding on this one?

Let me know if the confusion comes from the fact that the PR name, changelog, etc. still use the "number of parameters" phrasing (can update those places).

albi3ro

Mostly, the PR just has some hold over's from the previous functionality.

Also, I still think this functionality merges well with _grad_method_validation

albi3ro · 2021-06-09T02:43:48Z

pennylane/tape/jacobian_tape.py

+        diff_methods_to_sample = map(diff_methods.__getitem__, argnum)
+        mapped_diff_methods = zip(argnum, diff_methods_to_sample)
+
+        return random.sample(list(mapped_diff_methods), k=num_params)


So we are still selecting a random sample? But num_params is the length of argnum, and thus mapped_diff_methods, so randomly sampling all of them? Is this just permuting their order?

Oh, very good point! Removed the sampling, indeed, was a left-over part from the previous behaviour.

albi3ro · 2021-06-09T02:47:05Z

pennylane/tape/jacobian_tape.py

+        if argnum is None:
+            return enumerate(diff_methods)
+
+        if not isinstance(argnum, list):


Maybe more specifically say isinstance(argnum, int)? Since the options for argnum are int, list, or None. We already returned if None, so if its not a list, then its an int. I think...

Sure that's updated?

albi3ro · 2021-06-09T02:50:13Z

pennylane/tape/jacobian_tape.py

+        if not isinstance(argnum, list):
+            argnum = [argnum]
+
+        if not all(ind in self.trainable_params for ind in argnum):


For argnum in qml.grad, the specification of argnum overrides any requires_grad information. Should we be able to override and set a non-trainable parameter trainable via argnum?

Good point, after our chat it seems that we'll not consider this (also, seems like qml.grad doesn't cover all cases here).

pennylane/tape/jacobian_tape.py

.github/CHANGELOG.md

albi3ro · 2021-06-09T03:02:02Z

pennylane/tape/jacobian_tape.py

+                with respect to. When there are fewer parameters specified than the
+                total number of trainable parameters, the jacobian is being estimated by
+                sampling ``argnum`` many parameters from the set of trainable parameters.


Is this sampling thing still the case? Does this need to get updated?

Good point, removed that part.

albi3ro · 2021-06-09T03:44:36Z

tests/tape/test_jacobian_tape.py

+        res = res.flatten()
+        expected = expected.flatten()
+
+        assert any(np.allclose(r, e, atol=tol, rtol=0) for r, e in zip(res, expected))


This is deterministic now, so instead have

expected = np.array([[0, np.cos(y)*np.cos(x)]])

?

Great point! Updated

Co-authored-by: Christina Lee <christina@xanadu.ai>

…ylane into allow_jacobian_estimate

antalszava · 2021-06-09T16:23:28Z

@albi3ro thank you for the comments and suggestions here! Great catches. Should have them updated now.

albi3ro

Thanks for fixing all the things.

I made a comment about the if not isinstance(argnum, list): line, but that's not blocking.

pennylane/tape/jacobian_tape.py

antalszava added 4 commits May 28, 2021 20:57

unit tests

a8220c6

adjust, add another warn case, add an integration test

4f35362

further integration tests

a2b42f0

formatting

0a8739b

antalszava added 5 commits May 28, 2021 21:21

isort

4281e42

change test case

3a99d6e

Merge branch 'master' into allow_jacobian_estimate

7b57065

format

b67b86e

changelog

48a88a8

antalszava added 2 commits May 28, 2021 21:41

docstrings

488f55b

docstrings

1c6dda7

antalszava commented May 29, 2021

View reviewed changes

tests/tape/test_jacobian_tape.py Outdated Show resolved Hide resolved

Update tests/tape/test_jacobian_tape.py

b9e0156

antalszava requested a review from josh146 May 29, 2021 13:35

josh146 requested a review from albi3ro May 31, 2021 16:10

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

albi3ro reviewed May 31, 2021

View reviewed changes

josh146 approved these changes Jun 1, 2021

View reviewed changes

antalszava and others added 7 commits June 7, 2021 16:54

Update pennylane/tape/jacobian_tape.py

0b7f60c

Co-authored-by: Christina Lee <christina@xanadu.ai>

Update .github/CHANGELOG.md

28c8bc2

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

update based on suggestions

9dbe273

warn and error

a6f95ac

changelog

c73aa92

format

5db9053

changelog example

556a2a1

antalszava requested a review from albi3ro June 8, 2021 18:00

albi3ro reviewed Jun 9, 2021

View reviewed changes

albi3ro requested changes Jun 9, 2021

View reviewed changes

antalszava and others added 2 commits June 9, 2021 10:45

Update pennylane/tape/jacobian_tape.py

93c12b6

Co-authored-by: Christina Lee <christina@xanadu.ai>

Update .github/CHANGELOG.md

8776648

Co-authored-by: Christina Lee <christina@xanadu.ai>

antalszava changed the title ~~Allow specifying a number of parameters to estimate the jacobian~~ Allow specifying a subset of parameters to estimate the jacobian by passing the argnum kwarg Jun 9, 2021

antalszava added 8 commits June 9, 2021 10:47

Merge branch 'master' into allow_jacobian_estimate

cc351f7

update comment

a4371df

Merge branch 'allow_jacobian_estimate' of github.com:PennyLaneAI/penn…

fbc5591

…ylane into allow_jacobian_estimate

update test with exact expected expression

41369bc

formatting

234e634

Changelog

dc1f4ba

Jacobian

f1cf0ed

no sampling

ade0343

antalszava requested a review from albi3ro June 9, 2021 16:23

albi3ro approved these changes Jun 9, 2021

View reviewed changes

Merge branch 'master' into allow_jacobian_estimate

5c4ec67

antalszava commented Jun 9, 2021

View reviewed changes

pennylane/tape/jacobian_tape.py Outdated Show resolved Hide resolved

antalszava added 5 commits June 9, 2021 12:52

Update pennylane/tape/jacobian_tape.py

3b864d0

Merge branch 'master' into allow_jacobian_estimate

60debe0

changelog

4ccfc71

test for qnode with decorator

e480621

format

c44bedc

antalszava merged commit 92c0ae7 into master Jun 9, 2021

antalszava deleted the allow_jacobian_estimate branch June 9, 2021 21:11

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Allow specifying a subset of parameters to estimate the jacobian by passing the `argnum` kwarg #1371

Allow specifying a subset of parameters to estimate the jacobian by passing the `argnum` kwarg #1371

antalszava commented May 29, 2021 •

edited

Loading

github-actions bot commented May 29, 2021

codecov bot commented May 29, 2021 •

edited

Loading

antalszava May 29, 2021

albi3ro commented May 31, 2021

albi3ro left a comment

albi3ro May 31, 2021

antalszava Jun 8, 2021

albi3ro May 31, 2021

antalszava Jun 8, 2021

albi3ro May 31, 2021

antalszava Jun 8, 2021

josh146 left a comment

josh146 May 31, 2021

antalszava Jun 8, 2021

antalszava commented Jun 8, 2021

albi3ro left a comment

antalszava commented Jun 9, 2021

albi3ro left a comment

albi3ro Jun 9, 2021

antalszava Jun 9, 2021

albi3ro Jun 9, 2021

antalszava Jun 9, 2021

albi3ro Jun 9, 2021

albi3ro Jun 9, 2021

antalszava Jun 9, 2021

albi3ro Jun 9, 2021

antalszava Jun 9, 2021

albi3ro Jun 9, 2021

antalszava Jun 9, 2021

antalszava commented Jun 9, 2021

albi3ro left a comment

	object or list: map of the trainable parameter indices and
	generator or object or list: map of the trainable parameter indices and

Allow specifying a subset of parameters to estimate the jacobian by passing the argnum kwarg #1371

Allow specifying a subset of parameters to estimate the jacobian by passing the argnum kwarg #1371

Conversation

antalszava commented May 29, 2021 • edited Loading

github-actions bot commented May 29, 2021

codecov bot commented May 29, 2021 • edited Loading

Codecov Report

Choose a reason for hiding this comment

albi3ro commented May 31, 2021

albi3ro 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

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

antalszava commented Jun 8, 2021

albi3ro left a comment

Choose a reason for hiding this comment

antalszava commented Jun 9, 2021

albi3ro 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

antalszava commented Jun 9, 2021

albi3ro left a comment

Choose a reason for hiding this comment

Allow specifying a subset of parameters to estimate the jacobian by passing the `argnum` kwarg #1371

Allow specifying a subset of parameters to estimate the jacobian by passing the `argnum` kwarg #1371

antalszava commented May 29, 2021 •

edited

Loading

codecov bot commented May 29, 2021 •

edited

Loading