added multi_dispatch decorator

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/2084)
+
 <h3>Breaking changes</h3>
 
 * 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

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",

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
@@ -81,6 +83,79 @@ 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.
 

requirements.txt

@@ -13,3 +13,4 @@ semantic_version==2.6
 dask[delayed]==2021.10
 autoray>=0.2.5
 matplotlib==3.4
+black>=21

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]])