Repository: PennyLaneAI/pennylane

Issue:
Simplify adding multiple dispatch functions to the `qml.math` module
### Feature details

Currently, the file `pennylane/math/multi_dispatch.py` includes various 'multiple dispatch' functions, that accept multiple tensors, and determine the correct autodifferentiation framework to call. They do this by querying the utility function [`qml.math._multi_dispatch`](https://pennylane.readthedocs.io/en/latest/code/api/pennylane.math._multi_dispatch.html):

```python
def some_function(tensor1, tensor2, option):
    interface = qml.math._multi_dispatch([tensor1, tensor2])
    ...
```


### Implementation


We can simplify this process, while better documenting the API, by factoring out this boilerplate into a convenient decorator:

```python
@qml.math.multi_dispatch(argnum=[0, 1])
def some_function(tensor1, tensor2, option, like):
    # the interface string is stored in `like`.
    ...
```

Here, the `@qml.math.multi_dispatch` decorator would automate the following:

- it would intercept and determine the `interface` argument from the specified arguments of the wrapped function.
- it would make the `interface` argument optional, for enforcing a particular interface if required.

E.g.,

```python
def multi_dispatch(argnum=None, tensor_lists=None):
    # argnum list[int]: The argument indices to dispatch.
    #    If None, dispatch over all arguments.
    # tensor_list list[int]: Argument indices which correspond
    #    to lists of tensors.

    def decorator(fn)
        @functools.wraps(fn)
        def wrapper(*args, **kwargs):
            argnum = argnum or list(range(len(args))
            tensor_lists = tensor_list or []

            if not isinstance(argnum, Sequence):
                argnum = [argnum]

            if not isinstance(tensor_lists, Sequence):
                tensor_lists = [tensor_lists]

            dispatch_args = []

            for a in argnum:
                if a in tensor_list:
                    dispatch_args .extend(args[a])
                else:
                    dispatch_args .append(args[a])

            interface = kwargs.pop("like", None)
            interface = interface or qml.math._multi_dispatch(dispatch_args)
            kwargs["like"] = interface

            return fn(*args, **kwargs)
    return decorator
```

This would significantly simplify adding some functions, e.g.,

```python
stack = multi_dispatch(argnum=0, tensor_lists=0)(autoray.numpy.stack)
```

### How important would you say this feature is?

1: Not important. Would be nice to have.

### Additional information

_No response_


Discussion:

Resolving patch:
diff --git a/pennylane/math/__init__.py b/pennylane/math/__init__.py
--- 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,
@@ -76,6 +77,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
--- 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
@@ -81,6 +83,102 @@ def _multi_dispatch(values):
     return "numpy"
 
 
+def multi_dispatch(argnum=None, tensor_list=None):
+    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
+    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 expected to be lists of tensors.
+            If ``None``, this option is ignored.
+
+    Returns:
+        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`
+
+    .. 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)
+
+    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
+    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)
+    >>>     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
+
+    >>> values = [np.array([1, 2, 3]) for _ in range(5)]
+    >>> custom_function(values)
+    700
+
+    """
+
+    def decorator(fn):
+        @functools.wraps(fn)
+        def wrapper(*args, **kwargs):
+            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]
+            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.
 


Test patch:
diff --git a/tests/math/test_multi_disptach.py b/tests/math/test_multi_disptach.py
new file mode 100644
--- /dev/null
+++ b/tests/math/test_multi_disptach.py
@@ -0,0 +1,109 @@
+# 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 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]]),
+    jax.numpy.array([[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)
+    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_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)
+    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_decorate_non_dispatch(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)