From a1e2ea60ec65f0c45b54f99b0b6f1ea0959d9703 Mon Sep 17 00:00:00 2001 From: Qottmann Date: Wed, 12 Jan 2022 16:53:11 +0100 Subject: [PATCH 01/26] added multi_dispatch decorator --- pennylane/math/__init__.py | 2 + pennylane/math/multi_dispatch.py | 71 +++++++++++++++++++++++++++++++ tests/math/test_multi_disptach.py | 43 +++++++++++++++++++ 3 files changed, 116 insertions(+) create mode 100644 tests/math/test_multi_disptach.py diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py index 883f09ba259..edfe53508fc 100644 --- a/pennylane/math/__init__.py +++ b/pennylane/math/__init__.py @@ -35,6 +35,7 @@ from .multi_dispatch import ( _multi_dispatch, + multi_dispatch, block_diag, concatenate, diag, @@ -75,6 +76,7 @@ def __getattr__(name): __all__ = [ "_multi_dispatch", + "multi_dispatch", "allclose", "allequal", "block_diag", diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index e75a419ec03..2724aadb10b 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -14,6 +14,8 @@ """Multiple dispatch functions""" # pylint: disable=import-outside-toplevel,too-many-return-statements import warnings +from collections.abc import Sequence +import functools from autograd.numpy.numpy_boxes import ArrayBox from autoray import numpy as np @@ -80,6 +82,75 @@ def _multi_dispatch(values): return "numpy" +def multi_dispatch(argnum=None, tensor_list=None): + """Decorater to dispatch arguments handled by the interface. + + This helps simplify definitions of new functions inside pennylane. Instead of writing + + >>> def some_function(tensor1, tensor2, option): + ... interface = qml.math._multi_dispatch([tensor1, tensor2]) + ... ... + + We can decorate the function, indicating the arguments that are tensors handled + by the interface + + + >>> @qml.math.multi_dispatch(argnum=[0, 1]) + ... def some_function(tensor1, tensor2, option, like): + ... # the interface string is stored in `like`. + ... ... + + Args: + argnum (list[int]): a list of integers indicating indicating the indices + to dispatch (i.e. the arguments that are tensors handled by an interface) + If None, dispatch over all arguments + tensor_lists(list[int]): a list of integers indicating which indices + in argnum are lists of tensors. + If None, this option is ignored. + + Returns: + decorator: + + .. seealso:: :func:`pennylane.math.multi_dispatch._multi_dispatch` + + .. note:: + This decorator makes the interface argument "like" optional as it utilizes + the utility function `_multi_dispatch` to automatically detect the appropriate + interface based on the tensor types. + + ** Examples ** + We can redefine external functions to be suitable for pennylane. Here, we + redefine autoray's `stack` function. + >>> stack = multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack) + + """ + + def decorator(fn): + @functools.wraps(fn) + def wrapper(*args, **kwargs): + argnums = argnum or list(range(len(args))) + tensor_lists = tensor_list or [] + + if not isinstance(argnums, Sequence): + argnums = [argnums] + if not isinstance(tensor_lists, Sequence): + tensor_lists = [tensor_lists] + + dispatch_args = [] + + for a in argnums: + if a in tensor_lists: + dispatch_args.extend(args[a]) + else: + dispatch_args.append(args[a]) + + interface = kwargs.pop("like", None) + interface = interface or _multi_dispatch(dispatch_args) + kwargs["like"] = interface + + return fn(*args, **kwargs) + return wrapper + return decorator def block_diag(values): """Combine a sequence of 2D tensors to form a block diagonal tensor. diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_disptach.py new file mode 100644 index 00000000000..9e14466f7de --- /dev/null +++ b/tests/math/test_multi_disptach.py @@ -0,0 +1,43 @@ +# Copyright 2018-2020 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +""" Assertion test for multi_dispatch function/decorator +""" +import autoray +import numpy as onp +import pytest +from pennylane import numpy as np +from pennylane import math as fn + + +tf = pytest.importorskip("tensorflow", minversion="2.1") +torch = pytest.importorskip("torch") + +test_multi_dispatch_stack_data = [ + [[1.0, 0.0], [2.0, 3.0]], + ([1.0, 0.0], [2.0, 3.0]), + onp.array([[1.0, 0.0], [2.0, 3.0]]), + np.array([[1.0, 0.0], [2.0, 3.0]]), + # torch.tensor([[1.,0.],[2.,3.]]), + tf.Variable([[1.0, 0.0], [2.0, 3.0]]), + tf.constant([[1.0, 0.0], [2.0, 3.0]]), +] + + +@pytest.mark.parametrize("x", test_multi_dispatch_stack_data) +def test_multi_dispatch_stack(x): + """ Test that the decorated autoray function stack can handle all inputs """ + stack = fn.multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack) + res = stack(x) + print(res) + assert fn.allequal(res, [[1.0, 0.0], [2.0, 3.0]]) From 31661c1b59f0cd82ab151401da26ce8615ccd907 Mon Sep 17 00:00:00 2001 From: Qottmann Date: Wed, 12 Jan 2022 17:02:20 +0100 Subject: [PATCH 02/26] added release notes and ran black on only my files --- doc/releases/changelog-dev.md | 5 ++++- pennylane/math/multi_dispatch.py | 4 ++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 4a6d1393b7b..5f13183cc58 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -273,6 +273,9 @@ * The QAOA module now accepts both NetworkX and RetworkX graphs as function inputs. [(#1791)](https://github.com/PennyLaneAI/pennylane/pull/1791) +* Added `multi_dispatch` decorator that helps ease the definition of new functions. + [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2082) +

Breaking changes

* The behaviour of `RotosolveOptimizer` has been changed regarding @@ -341,4 +344,4 @@ This release contains contributions from (in alphabetical order): -Juan Miguel Arrazola, Ali Asadi, Esther Cruz, Olivia Di Matteo, Diego Guala, Ankit Khandelwal, Jay Soni, Antal Száva, David Wierichs, Shaoming Zhang +Juan Miguel Arrazola, Ali Asadi, Esther Cruz, Olivia Di Matteo, Diego Guala, Ankit Khandelwal, Korbinian Kottmann, Jay Soni, Antal Száva, David Wierichs, Shaoming Zhang diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 2724aadb10b..82b6211d0ca 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -82,6 +82,7 @@ def _multi_dispatch(values): return "numpy" + def multi_dispatch(argnum=None, tensor_list=None): """Decorater to dispatch arguments handled by the interface. @@ -149,9 +150,12 @@ def wrapper(*args, **kwargs): kwargs["like"] = interface return fn(*args, **kwargs) + return wrapper + return decorator + def block_diag(values): """Combine a sequence of 2D tensors to form a block diagonal tensor. From 96ff41843b12085688a382b3fe60fd230a170a77 Mon Sep 17 00:00:00 2001 From: Qottmann Date: Wed, 12 Jan 2022 17:03:27 +0100 Subject: [PATCH 03/26] typo in PR number, now 2084 --- doc/releases/changelog-dev.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 5f13183cc58..84e4af1e34d 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -274,7 +274,7 @@ [(#1791)](https://github.com/PennyLaneAI/pennylane/pull/1791) * Added `multi_dispatch` decorator that helps ease the definition of new functions. - [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2082) + [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084)

Breaking changes

From c606d34b06de8a72e67bee18b502ce295e1d3c36 Mon Sep 17 00:00:00 2001 From: Qottmann Date: Wed, 12 Jan 2022 19:28:33 +0100 Subject: [PATCH 04/26] updated black, reran on test_multi_dispatch --- tests/math/test_multi_disptach.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_disptach.py index 9e14466f7de..2a039da855a 100644 --- a/tests/math/test_multi_disptach.py +++ b/tests/math/test_multi_disptach.py @@ -36,7 +36,7 @@ @pytest.mark.parametrize("x", test_multi_dispatch_stack_data) def test_multi_dispatch_stack(x): - """ Test that the decorated autoray function stack can handle all inputs """ + """Test that the decorated autoray function stack can handle all inputs""" stack = fn.multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack) res = stack(x) print(res) From 790e52a1c6b5e3338c611773431d7fcf86df81fb Mon Sep 17 00:00:00 2001 From: Qottmann Date: Wed, 12 Jan 2022 19:37:21 +0100 Subject: [PATCH 05/26] added black version requirement, optional commit --- doc/requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/requirements.txt b/doc/requirements.txt index 810d15f25e8..a9c1da592e2 100644 --- a/doc/requirements.txt +++ b/doc/requirements.txt @@ -21,3 +21,4 @@ tensornetwork==0.3 toml torch==1.9.0+cpu torchvision==0.10.0+cpu +black>=21 From c864c4828dfefbec0489cf90bd7d47d46cdb2fdb Mon Sep 17 00:00:00 2001 From: Josh Izaac Date: Thu, 13 Jan 2022 15:39:49 +0800 Subject: [PATCH 06/26] Update doc/requirements.txt --- doc/requirements.txt | 1 - 1 file changed, 1 deletion(-) diff --git a/doc/requirements.txt b/doc/requirements.txt index a9c1da592e2..810d15f25e8 100644 --- a/doc/requirements.txt +++ b/doc/requirements.txt @@ -21,4 +21,3 @@ tensornetwork==0.3 toml torch==1.9.0+cpu torchvision==0.10.0+cpu -black>=21 From cc7fc37af502bb7bc5297eaef5d78268e2c4664c Mon Sep 17 00:00:00 2001 From: Qottmann Date: Thu, 13 Jan 2022 12:35:17 +0100 Subject: [PATCH 07/26] moved black>=21 requirement out of docs to /requirements.txt --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index e9cead71d73..dfeef6c5743 100644 --- a/requirements.txt +++ b/requirements.txt @@ -13,3 +13,4 @@ semantic_version==2.6 dask[delayed]==2021.10 autoray>=0.2.5 matplotlib==3.4 +black>=21 \ No newline at end of file From b4a6e8c75ef4d460f151af58d334daa8049422fb Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Thu, 13 Jan 2022 16:24:33 +0100 Subject: [PATCH 08/26] Update tests/math/test_multi_disptach.py Co-authored-by: Josh Izaac --- tests/math/test_multi_disptach.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_disptach.py index 2a039da855a..bbe01535202 100644 --- a/tests/math/test_multi_disptach.py +++ b/tests/math/test_multi_disptach.py @@ -39,5 +39,4 @@ def test_multi_dispatch_stack(x): """Test that the decorated autoray function stack can handle all inputs""" stack = fn.multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack) res = stack(x) - print(res) assert fn.allequal(res, [[1.0, 0.0], [2.0, 3.0]]) From a25e9b5e8d3e4b8e7d3b5144d84f0be5feb925e3 Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Thu, 13 Jan 2022 16:24:43 +0100 Subject: [PATCH 09/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 82b6211d0ca..0bbbf844d86 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -86,7 +86,7 @@ def _multi_dispatch(values): def multi_dispatch(argnum=None, tensor_list=None): """Decorater to dispatch arguments handled by the interface. - This helps simplify definitions of new functions inside pennylane. Instead of writing + This helps simplify definitions of new functions inside PennyLane. Instead of writing >>> def some_function(tensor1, tensor2, option): ... interface = qml.math._multi_dispatch([tensor1, tensor2]) From 7c97b7f673b0d355bb73c815085869a011c84fae Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Thu, 13 Jan 2022 16:25:04 +0100 Subject: [PATCH 10/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 0bbbf844d86..8d02c8cfe57 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -102,9 +102,9 @@ def multi_dispatch(argnum=None, tensor_list=None): ... ... Args: - argnum (list[int]): a list of integers indicating indicating the indices - to dispatch (i.e. the arguments that are tensors handled by an interface) - If None, dispatch over all arguments + argnum (list[int]): A list of integers indicating indicating the indices + to dispatch (i.e., the arguments that are tensors handled by an interface). + If ``None``, dispatch over all arguments. tensor_lists(list[int]): a list of integers indicating which indices in argnum are lists of tensors. If None, this option is ignored. From ba212d43b44a229eadab29c0e51b249434e6f531 Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Thu, 13 Jan 2022 16:25:13 +0100 Subject: [PATCH 11/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 8d02c8cfe57..70f22a3f828 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -105,9 +105,9 @@ def multi_dispatch(argnum=None, tensor_list=None): argnum (list[int]): A list of integers indicating indicating the indices to dispatch (i.e., the arguments that are tensors handled by an interface). If ``None``, dispatch over all arguments. - tensor_lists(list[int]): a list of integers indicating which indices - in argnum are lists of tensors. - If None, this option is ignored. + tensor_lists (list[int]): a list of integers indicating which indices + in ``argnum`` are expected to be lists of tensors. + If ``None``, this option is ignored. Returns: decorator: From bd3b15e2c8297190a247e94bea87739e8415842c Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Thu, 13 Jan 2022 16:25:33 +0100 Subject: [PATCH 12/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 70f22a3f828..da5e7933c85 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -110,7 +110,10 @@ def multi_dispatch(argnum=None, tensor_list=None): If ``None``, this option is ignored. Returns: - decorator: + func: A wrapped version of the function, which will automatically attempt + to dispatch to the correct autodifferentiation framework for the requested + arguments. Note that the ``like`` argument will be optional, but can be provided + if an explicit override is needed. .. seealso:: :func:`pennylane.math.multi_dispatch._multi_dispatch` From ac9a82e491795850c91653b46d13b647109be322 Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Thu, 13 Jan 2022 16:26:14 +0100 Subject: [PATCH 13/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index da5e7933c85..e117e458936 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -122,9 +122,11 @@ def multi_dispatch(argnum=None, tensor_list=None): the utility function `_multi_dispatch` to automatically detect the appropriate interface based on the tensor types. - ** Examples ** - We can redefine external functions to be suitable for pennylane. Here, we - redefine autoray's `stack` function. + **Examples** + + We can redefine external functions to be suitable for PennyLane. Here, we + redefine Autoray's ``stack`` function. + >>> stack = multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack) """ From 48b1834e30588ed777a2fc4bacf62304b31169f0 Mon Sep 17 00:00:00 2001 From: Qottmann Date: Thu, 13 Jan 2022 20:08:56 +0100 Subject: [PATCH 14/26] added example in multi_dispatch and expanded tests --- doc/releases/changelog-dev.md | 16 +++++++++++- pennylane/math/multi_dispatch.py | 31 +++++++++++++++++------ tests/math/test_multi_disptach.py | 42 +++++++++++++++++++++++++++++-- 3 files changed, 79 insertions(+), 10 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 84e4af1e34d..c4563a99885 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -273,7 +273,21 @@ * The QAOA module now accepts both NetworkX and RetworkX graphs as function inputs. [(#1791)](https://github.com/PennyLaneAI/pennylane/pull/1791) -* Added `multi_dispatch` decorator that helps ease the definition of new functions. +* Added `multi_dispatch` decorator that helps ease the definition of new functions + inside PennyLane. We can decorate the function, indicating the arguments that are + tensors handled by the interface + + >>> @qml.math.multi_dispatch(argnum=[0, 1]) + ... def some_function(tensor1, tensor2, option, like): + ... # the interface string is stored in ``like``. + ... ... + + Previously, this was done using the private utility function ``_multi_dispatch``. + + >>> def some_function(tensor1, tensor2, option): + ... interface = qml.math._multi_dispatch([tensor1, tensor2]) + ... ... + [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084)

Breaking changes

diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index e117e458936..1925f7e3c9e 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -86,13 +86,8 @@ def _multi_dispatch(values): def multi_dispatch(argnum=None, tensor_list=None): """Decorater to dispatch arguments handled by the interface. - This helps simplify definitions of new functions inside PennyLane. Instead of writing - - >>> def some_function(tensor1, tensor2, option): - ... interface = qml.math._multi_dispatch([tensor1, tensor2]) - ... ... - - We can decorate the function, indicating the arguments that are tensors handled + This helps simplify definitions of new functions inside PennyLane. We can + decorate the function, indicating the arguments that are tensors handled by the interface @@ -101,6 +96,7 @@ def multi_dispatch(argnum=None, tensor_list=None): ... # the interface string is stored in `like`. ... ... + Args: argnum (list[int]): A list of integers indicating indicating the indices to dispatch (i.e., the arguments that are tensors handled by an interface). @@ -129,6 +125,27 @@ def multi_dispatch(argnum=None, tensor_list=None): >>> stack = multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack) + We can also define more elaborate custom function. Here is an example of a ``custom_function`` that + computes :math:`c \sum_i (v_i)^T v_i`, where :math:`v_i` are vectors in ``values`` and + :math:`c` is a fixed ``coefficient``. Note how ``argnum=0`` only points to the first argument ``values``, + how ``tensor_list=0`` indicates that said first argument is a list of vectors, and that ``coefficient`` is not + dispatched. + + >>> @math.multi_dispatch(argnum=0,tensor_list=0) + >>> def custom_function(values, like, coefficient=10): + >>> # values is a list of vectors + >>> # like can force the interface (optional) + >>> return coefficient * np.sum([math.dot(v,v) for v in values]) + + We can then run + + >>> values = [np.array([1,2,3]) for _ in range(5)] + >>> print(custom_function(values)) + + and obtain + + >>> 700 + """ def decorator(fn): diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_disptach.py index bbe01535202..876cb1c6885 100644 --- a/tests/math/test_multi_disptach.py +++ b/tests/math/test_multi_disptach.py @@ -16,20 +16,22 @@ import autoray import numpy as onp import pytest +from autoray import numpy as anp from pennylane import numpy as np from pennylane import math as fn tf = pytest.importorskip("tensorflow", minversion="2.1") torch = pytest.importorskip("torch") +jax = pytest.importorskip("jax") test_multi_dispatch_stack_data = [ [[1.0, 0.0], [2.0, 3.0]], ([1.0, 0.0], [2.0, 3.0]), onp.array([[1.0, 0.0], [2.0, 3.0]]), + anp.array([[1.0, 0.0], [2.0, 3.0]]), np.array([[1.0, 0.0], [2.0, 3.0]]), - # torch.tensor([[1.,0.],[2.,3.]]), - tf.Variable([[1.0, 0.0], [2.0, 3.0]]), + jax.numpy.array([[1.0, 0.0], [2.0, 3.0]]), tf.constant([[1.0, 0.0], [2.0, 3.0]]), ] @@ -40,3 +42,39 @@ def test_multi_dispatch_stack(x): stack = fn.multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack) res = stack(x) assert fn.allequal(res, [[1.0, 0.0], [2.0, 3.0]]) + + +@pytest.mark.parametrize("x", test_multi_dispatch_stack_data) +def test_multi_dispatch_decorate(x): + """Test decorating a standard numpy function for PennyLane""" + + @fn.multi_dispatch(argnum=[0], tensor_list=[0]) + def tensordot(x, like, axes=None): + return np.tensordot(x[0], x[1], axes=axes) + + assert fn.allequal(tensordot(x, axes=(0, 0)).numpy(), 2) + + +test_data0 = [ + (1, 2, 3), + [1, 2, 3], + onp.array([1, 2, 3]), + anp.array([1, 2, 3]), + np.array([1, 2, 3]), + torch.tensor([1, 2, 3]), + jax.numpy.array([1, 2, 3]), + tf.constant([1, 2, 3]), +] + +test_data = [(x, x) for x in test_data0] + + +@pytest.mark.parametrize("t1,t2", test_data) +def test_multi_dispatch_decorate2(t1, t2): + """Test decorating a standard numpy function for PennyLane, this time using 2 separate inputs""" + + @fn.multi_dispatch(argnum=[0, 1], tensor_list=None) + def tensordot(tensor1, tensor2, like, axes=None): + return np.tensordot(tensor1, tensor2, axes=axes) + + assert fn.allequal(tensordot(t1, t2, axes=(0, 0)).numpy(), 14) From 590788088f7a83e90ea5a770d875bfcfdebb3ca3 Mon Sep 17 00:00:00 2001 From: Qottmann Date: Fri, 14 Jan 2022 11:46:18 +0100 Subject: [PATCH 15/26] changed argument handling to allow for argnum=0 and added custom function to complete tests --- pennylane/math/multi_dispatch.py | 6 +++--- tests/math/test_multi_disptach.py | 29 +++++++++++++++++++++++++++-- 2 files changed, 30 insertions(+), 5 deletions(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 1925f7e3c9e..e77f6fd2afd 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -126,7 +126,7 @@ def multi_dispatch(argnum=None, tensor_list=None): >>> stack = multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack) We can also define more elaborate custom function. Here is an example of a ``custom_function`` that - computes :math:`c \sum_i (v_i)^T v_i`, where :math:`v_i` are vectors in ``values`` and + computes :math:`c \\sum_i (v_i)^T v_i`, where :math:`v_i` are vectors in ``values`` and :math:`c` is a fixed ``coefficient``. Note how ``argnum=0`` only points to the first argument ``values``, how ``tensor_list=0`` indicates that said first argument is a list of vectors, and that ``coefficient`` is not dispatched. @@ -151,8 +151,8 @@ def multi_dispatch(argnum=None, tensor_list=None): def decorator(fn): @functools.wraps(fn) def wrapper(*args, **kwargs): - argnums = argnum or list(range(len(args))) - tensor_lists = tensor_list or [] + argnums = argnum if argnum is not None else list(range(len(args))) + tensor_lists = tensor_list if tensor_list is not None else [] if not isinstance(argnums, Sequence): argnums = [argnums] diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_disptach.py index 876cb1c6885..eee2ab0d8a0 100644 --- a/tests/math/test_multi_disptach.py +++ b/tests/math/test_multi_disptach.py @@ -71,10 +71,35 @@ def tensordot(x, like, axes=None): @pytest.mark.parametrize("t1,t2", test_data) def test_multi_dispatch_decorate2(t1, t2): - """Test decorating a standard numpy function for PennyLane, this time using 2 separate inputs""" + """Test decorating a standard numpy function for PennyLane, automatically dispatching all inputs by choosing argnum=None""" - @fn.multi_dispatch(argnum=[0, 1], tensor_list=None) + @fn.multi_dispatch(argnum=None, tensor_list=None) def tensordot(tensor1, tensor2, like, axes=None): return np.tensordot(tensor1, tensor2, axes=axes) assert fn.allequal(tensordot(t1, t2, axes=(0, 0)).numpy(), 14) + +test_data_values = [ + [[1,2,3] for _ in range(5)], + [(1,2,3) for _ in range(5)], + [np.array([1,2,3]) for _ in range(5)], + [onp.array([1,2,3]) for _ in range(5)], + [anp.array([1,2,3]) for _ in range(5)], + [torch.tensor([1,2,3]) for _ in range(5)], + [jax.numpy.array([1,2,3]) for _ in range(5)], + [tf.constant([1,2,3]) for _ in range(5)] +] + +@pytest.mark.parametrize("values", test_data_values) +def test_multi_dispatch_decorate3(values): + """Test decorating a custom function for PennyLane including a non-dispatchable parameter """ + @fn.multi_dispatch(argnum=0,tensor_list=0) + def custom_function(values, like, coefficient=10): + """ + A dummy custom function that computes coeff :math:`c \\sum_i (v_i)^T v_i` where :math:`v_i` are vectors in ``values`` + and :math:`c` is a fixed ``coefficient``. + values is a list of vectors + like can force the interface (optional) + """ + return coefficient * np.sum([fn.dot(v,v) for v in values]) + assert fn.allequal(custom_function(values),700) From e8cf9810021cc23c6e284d081a5b9b8d729c1ab9 Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Fri, 14 Jan 2022 15:32:28 +0100 Subject: [PATCH 16/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index e77f6fd2afd..40cfa10ef71 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -139,12 +139,9 @@ def multi_dispatch(argnum=None, tensor_list=None): We can then run - >>> values = [np.array([1,2,3]) for _ in range(5)] - >>> print(custom_function(values)) - - and obtain - - >>> 700 + >>> values = [np.array([1, 2, 3]) for _ in range(5)] + >>> custom_function(values) + 700 """ From a9455797e09c15442619bdbcaaffd27178c1b202 Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Fri, 14 Jan 2022 15:33:13 +0100 Subject: [PATCH 17/26] Update doc/releases/changelog-dev.md Co-authored-by: Josh Izaac --- doc/releases/changelog-dev.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 61321c15acc..6b49bc22347 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -279,7 +279,8 @@ * Added `multi_dispatch` decorator that helps ease the definition of new functions inside PennyLane. We can decorate the function, indicating the arguments that are - tensors handled by the interface + tensors handled by the interface: + [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084) >>> @qml.math.multi_dispatch(argnum=[0, 1]) ... def some_function(tensor1, tensor2, option, like): From a620e13b132ac4073646d5618b99d3511c47f039 Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Fri, 14 Jan 2022 15:33:46 +0100 Subject: [PATCH 18/26] Update doc/releases/changelog-dev.md Co-authored-by: Josh Izaac --- doc/releases/changelog-dev.md | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 6b49bc22347..e9e3b208fdc 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -282,16 +282,20 @@ tensors handled by the interface: [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084) + ```pycon >>> @qml.math.multi_dispatch(argnum=[0, 1]) ... def some_function(tensor1, tensor2, option, like): ... # the interface string is stored in ``like``. ... ... + ``` - Previously, this was done using the private utility function ``_multi_dispatch``. + Previously, this was done using the private utility function `_multi_dispatch`. + ```pycon >>> def some_function(tensor1, tensor2, option): ... interface = qml.math._multi_dispatch([tensor1, tensor2]) ... ... + ``` [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084) From 65110b1a3707320f6d307963262271ae21efee49 Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Fri, 14 Jan 2022 15:34:31 +0100 Subject: [PATCH 19/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 40cfa10ef71..0f29180dc82 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -88,7 +88,7 @@ def multi_dispatch(argnum=None, tensor_list=None): This helps simplify definitions of new functions inside PennyLane. We can decorate the function, indicating the arguments that are tensors handled - by the interface + by the interface: >>> @qml.math.multi_dispatch(argnum=[0, 1]) From 421eb7ca060ff334fedbbbca677a529123d5478c Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Fri, 14 Jan 2022 15:35:18 +0100 Subject: [PATCH 20/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 0f29180dc82..618b160482c 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -84,7 +84,7 @@ def _multi_dispatch(values): def multi_dispatch(argnum=None, tensor_list=None): - """Decorater to dispatch arguments handled by the interface. + r"""Decorater to dispatch arguments handled by the interface. This helps simplify definitions of new functions inside PennyLane. We can decorate the function, indicating the arguments that are tensors handled From 3bdf5a7108043e840283b548fa418aef645c1b1e Mon Sep 17 00:00:00 2001 From: Qottmann Date: Fri, 14 Jan 2022 15:45:40 +0100 Subject: [PATCH 21/26] small review --- doc/releases/changelog-dev.md | 2 -- pennylane/math/multi_dispatch.py | 2 ++ tests/math/test_multi_disptach.py | 28 ++++++++++++++++------------ 3 files changed, 18 insertions(+), 14 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index e9e3b208fdc..e790ad8d760 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -296,8 +296,6 @@ ... interface = qml.math._multi_dispatch([tensor1, tensor2]) ... ... ``` - - [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084)

Breaking changes

diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index 618b160482c..be1212543df 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -135,6 +135,8 @@ def multi_dispatch(argnum=None, tensor_list=None): >>> def custom_function(values, like, coefficient=10): >>> # values is a list of vectors >>> # like can force the interface (optional) + >>> if like == "tensorflow": + >>> # add interface-specific handling if necessary >>> return coefficient * np.sum([math.dot(v,v) for v in values]) We can then run diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_disptach.py index eee2ab0d8a0..d719efcf390 100644 --- a/tests/math/test_multi_disptach.py +++ b/tests/math/test_multi_disptach.py @@ -79,21 +79,24 @@ def tensordot(tensor1, tensor2, like, axes=None): assert fn.allequal(tensordot(t1, t2, axes=(0, 0)).numpy(), 14) + test_data_values = [ - [[1,2,3] for _ in range(5)], - [(1,2,3) for _ in range(5)], - [np.array([1,2,3]) for _ in range(5)], - [onp.array([1,2,3]) for _ in range(5)], - [anp.array([1,2,3]) for _ in range(5)], - [torch.tensor([1,2,3]) for _ in range(5)], - [jax.numpy.array([1,2,3]) for _ in range(5)], - [tf.constant([1,2,3]) for _ in range(5)] + [[1, 2, 3] for _ in range(5)], + [(1, 2, 3) for _ in range(5)], + [np.array([1, 2, 3]) for _ in range(5)], + [onp.array([1, 2, 3]) for _ in range(5)], + [anp.array([1, 2, 3]) for _ in range(5)], + [torch.tensor([1, 2, 3]) for _ in range(5)], + [jax.numpy.array([1, 2, 3]) for _ in range(5)], + [tf.constant([1, 2, 3]) for _ in range(5)], ] + @pytest.mark.parametrize("values", test_data_values) def test_multi_dispatch_decorate3(values): - """Test decorating a custom function for PennyLane including a non-dispatchable parameter """ - @fn.multi_dispatch(argnum=0,tensor_list=0) + """Test decorating a custom function for PennyLane including a non-dispatchable parameter""" + + @fn.multi_dispatch(argnum=0, tensor_list=0) def custom_function(values, like, coefficient=10): """ A dummy custom function that computes coeff :math:`c \\sum_i (v_i)^T v_i` where :math:`v_i` are vectors in ``values`` @@ -101,5 +104,6 @@ def custom_function(values, like, coefficient=10): values is a list of vectors like can force the interface (optional) """ - return coefficient * np.sum([fn.dot(v,v) for v in values]) - assert fn.allequal(custom_function(values),700) + return coefficient * np.sum([fn.dot(v, v) for v in values]) + + assert fn.allequal(custom_function(values), 700) From 3ff2742ce58b1092b5de468bb0e2d18bde51b15f Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Mon, 17 Jan 2022 11:59:47 +0100 Subject: [PATCH 22/26] Update doc/releases/changelog-dev.md Co-authored-by: Josh Izaac --- doc/releases/changelog-dev.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 3408b4c0db8..f8d799a09f7 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -289,7 +289,7 @@ [(#2062)](https://github.com/PennyLaneAI/pennylane/pull/2062) [(#2063)](https://github.com/PennyLaneAI/pennylane/pull/2063) -* Added `multi_dispatch` decorator that helps ease the definition of new functions +* Added a new `multi_dispatch` decorator that helps ease the definition of new functions inside PennyLane. We can decorate the function, indicating the arguments that are tensors handled by the interface: [(#2082)](https://github.com/PennyLaneAI/pennylane/pull/2084) From 7df030b227bdcd0f716aabfd970604ccdce977ab Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Mon, 17 Jan 2022 11:59:54 +0100 Subject: [PATCH 23/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index e56ad118a83..ab9eadff026 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -131,7 +131,7 @@ def multi_dispatch(argnum=None, tensor_list=None): how ``tensor_list=0`` indicates that said first argument is a list of vectors, and that ``coefficient`` is not dispatched. - >>> @math.multi_dispatch(argnum=0,tensor_list=0) + >>> @math.multi_dispatch(argnum=0, tensor_list=0) >>> def custom_function(values, like, coefficient=10): >>> # values is a list of vectors >>> # like can force the interface (optional) From cf81b3ec95ce60fbe6706d814719c4aed4124b24 Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Mon, 17 Jan 2022 12:00:05 +0100 Subject: [PATCH 24/26] Update pennylane/math/multi_dispatch.py Co-authored-by: Josh Izaac --- pennylane/math/multi_dispatch.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/pennylane/math/multi_dispatch.py b/pennylane/math/multi_dispatch.py index ab9eadff026..a61288dbd2a 100644 --- a/pennylane/math/multi_dispatch.py +++ b/pennylane/math/multi_dispatch.py @@ -125,7 +125,9 @@ def multi_dispatch(argnum=None, tensor_list=None): >>> stack = multi_dispatch(argnum=0, tensor_list=0)(autoray.numpy.stack) - We can also define more elaborate custom function. Here is an example of a ``custom_function`` that + We can also use the ``multi_dispatch`` decorator to dispatch + arguments of more more elaborate custom functions. Here is an example + of a ``custom_function`` that computes :math:`c \\sum_i (v_i)^T v_i`, where :math:`v_i` are vectors in ``values`` and :math:`c` is a fixed ``coefficient``. Note how ``argnum=0`` only points to the first argument ``values``, how ``tensor_list=0`` indicates that said first argument is a list of vectors, and that ``coefficient`` is not From fb79fbe44c37077b58ab8bf1e82aba8f26e0be0e Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Mon, 17 Jan 2022 12:00:19 +0100 Subject: [PATCH 25/26] Update tests/math/test_multi_disptach.py Co-authored-by: Josh Izaac --- tests/math/test_multi_disptach.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_disptach.py index d719efcf390..e0ba1aa688a 100644 --- a/tests/math/test_multi_disptach.py +++ b/tests/math/test_multi_disptach.py @@ -70,7 +70,7 @@ def tensordot(x, like, axes=None): @pytest.mark.parametrize("t1,t2", test_data) -def test_multi_dispatch_decorate2(t1, t2): +def test_multi_dispatch_decorate_argnum_none(t1, t2): """Test decorating a standard numpy function for PennyLane, automatically dispatching all inputs by choosing argnum=None""" @fn.multi_dispatch(argnum=None, tensor_list=None) From ba7cb0d27965bdb642d29648d70dd8246432eec0 Mon Sep 17 00:00:00 2001 From: Korbinian Kottmann <43949391+Qottmann@users.noreply.github.com> Date: Mon, 17 Jan 2022 12:00:28 +0100 Subject: [PATCH 26/26] Update tests/math/test_multi_disptach.py Co-authored-by: Josh Izaac --- tests/math/test_multi_disptach.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_disptach.py index e0ba1aa688a..406b68fb094 100644 --- a/tests/math/test_multi_disptach.py +++ b/tests/math/test_multi_disptach.py @@ -93,7 +93,7 @@ def tensordot(tensor1, tensor2, like, axes=None): @pytest.mark.parametrize("values", test_data_values) -def test_multi_dispatch_decorate3(values): +def test_multi_dispatch_decorate_non_dispatch(values): """Test decorating a custom function for PennyLane including a non-dispatchable parameter""" @fn.multi_dispatch(argnum=0, tensor_list=0)