_overrides.py

"""
Python implementation of __torch_function__

While most of the torch API and handling for __torch_function__ happens
at the C++ level, some of the torch API is written in Python so we need
python-level handling for __torch_function__ overrides as well. The main
developer-facing functionality in this file is the
torch_function_dispatch decorator. This function can be applied to
python functions in the torch.functional module to enable
__torch_function__ overrides for that function. See the examples in the
docstrings for torch_function_dispatch for details.

NOTE: heavily inspired by NumPy's ``__array_function__`` (see:
https://github.com/pytorch/pytorch/issues/24015 and
https://www.numpy.org/neps/nep-0018-array-function-protocol.html
)

"""

import functools
import textwrap
from . import _six
if _six.PY3:
    from inspect import getfullargspec
    import collections
    ArgSpec = collections.namedtuple('ArgSpec', 'args varargs keywords defaults')

    def getargspec(func):
        spec = getfullargspec(func)
        return ArgSpec(spec.args, spec.varargs, spec.varkw, spec.defaults)
else:
    from inspect import getargspec

from .tensor import Tensor


_TENSOR_ONLY = [Tensor]

def _get_overloaded_types_and_args(relevant_args):
    """Returns a list of arguments on which to call __torch_function__.

    Checks arguments in relevant_args for __torch_function__ implementations,
    storing references to the arguments and their types in overloaded_args and
    overloaded_types in order of calling precedence. Only distinct types are
    considered. If a type is a subclass of another type it will have higher
    precedence, otherwise the precedence order is the same as the order of
    arguments in relevant_args, that is, from left-to-right in the argument list.

    The precedence-determining algorithm implemented in this function is
    described in `NEP-0018`_.

    See torch::append_overloaded_arg for the equivalent function in the C++
    implementation.

    Parameters
    ----------
    relevant_args : iterable of array-like
        Iterable of array-like arguments to check for __torch_function__
        methods.

    Returns
    -------
    overloaded_types : collection of types
        Types of arguments from relevant_args with __torch_function__ methods.
    overloaded_args : list
        Arguments from relevant_args on which to call __torch_function__
        methods, in the order in which they should be called.

    .. _NEP-0018:
       https://numpy.org/neps/nep-0018-array-function-protocol.html

    """
    # Runtime is O(num_arguments * num_unique_types)
    overloaded_types = []
    overloaded_args = []
    for arg in relevant_args:
        arg_type = type(arg)
        # We only collect arguments if they have a unique type, which ensures
        # reasonable performance even with a long list of possibly overloaded
        # arguments.
        if (arg_type not in overloaded_types and hasattr(arg_type, '__torch_function__')):
            # Create lists explicitly for the first type (usually the only one
            # done) to avoid setting up the iterator for overloaded_args.
            if overloaded_types:
                overloaded_types.append(arg_type)
                # By default, insert argument at the end, but if it is
                # subclass of another argument, insert it before that argument.
                # This ensures "subclasses before superclasses".
                index = len(overloaded_args)
                for i, old_arg in enumerate(overloaded_args):
                    if issubclass(arg_type, type(old_arg)):
                        index = i
                        break
                overloaded_args.insert(index, arg)
            else:
                overloaded_types = [arg_type]
                overloaded_args = [arg]

    return overloaded_types, overloaded_args


def _implement_torch_function(
        implementation, public_api, relevant_args, args, kwargs):
    """Implement a function with checks for __torch_function__ overrides.

    See torch::autograd::handle_torch_function for the equivalent of this
    function in the C++ implementation.

    Arguments
    ---------
    implementation : function
        Function that implements the operation on ``torch.Tensor`` without
        overrides when called like ``implementation(*args, **kwargs)``.
    public_api : function
        Function exposed by the public torch API originally called like
        ``public_api(*args, **kwargs)`` on which arguments are now being
        checked.
    relevant_args : iterable
        Iterable of arguments to check for __torch_function__ methods.
    args : tuple
        Arbitrary positional arguments originally passed into ``public_api``.
    kwargs : tuple
        Arbitrary keyword arguments originally passed into ``public_api``.

    Returns
    -------
    Result from calling `implementation()` or an `__torch_function__`
    method, as appropriate.

    Raises
    ------
    TypeError : if no implementation is found.

    """
    # Check for __torch_function__ methods.
    types, overloaded_args = _get_overloaded_types_and_args(relevant_args)
    # Short-cut for common cases: no overload or only Tensor overload
    # (directly or with subclasses that do not override __torch_function__).
    if not overloaded_args or types == _TENSOR_ONLY:
        return implementation(*args, **kwargs)

    # Call overrides
    for overloaded_arg in overloaded_args:
        # Use `public_api` instead of `implementation` so __torch_function__
        # implementations can do equality/identity comparisons.
        result = overloaded_arg.__torch_function__(public_api, args, kwargs)

        if result is not NotImplemented:
            return result

    func_name = '{}.{}'.format(public_api.__module__, public_api.__name__)
    raise TypeError("no implementation found for '{}' on types that implement "
                    '__torch_function__: {}'
                    .format(func_name, list(map(type, overloaded_args))))


def _verify_matching_signatures(implementation, dispatcher):
    """Verify that a dispatcher function has the right signature."""
    implementation_spec = getargspec(implementation)
    dispatcher_spec = getargspec(dispatcher)

    if (implementation_spec.args != dispatcher_spec.args or
        implementation_spec.varargs != dispatcher_spec.varargs or
        implementation_spec.keywords != dispatcher_spec.keywords or
        (bool(implementation_spec.defaults) !=
         bool(dispatcher_spec.defaults)) or
        (implementation_spec.defaults is not None and
         len(implementation_spec.defaults) !=
         len(dispatcher_spec.defaults))):
        raise RuntimeError('implementation and dispatcher for %s have '
                           'different function signatures' % implementation)


_wrapped_func_source = textwrap.dedent("""
    @functools.wraps(implementation)
    def {name}(*args, **kwargs):
        relevant_args = dispatcher(*args, **kwargs)
        return implement_torch_function(
            implementation, {name}, relevant_args, args, kwargs)
    """)

def torch_function_dispatch(dispatcher, module=None, verify=True):
    """Decorator for adding dispatch with the __torch_function__ protocol.

    If you define a function in Python and would like to permit user-defined
    tensor-like types to override it using __torch_function__, please apply this
    decorator on this function together with a custom dispatcher that indicates
    which arguments should be checked for the presence of __torch_function__.

    Suppose we'd like to apply this function to torch.frob, which has the
    following definition:

        def frob(input, bias, option=None):
            return input + bias

    We'd need to define a dispatcher for frob that has the same signature and
    returns the elements of the signature that should be checked for
    `__torch_function__`. If any of the arguments has a `__torch_function__`
    attribute, that function will be called to handle custom dispatch. Assuming
    that `bias` can be a tensor-like, our dispatcher would look like:

        def _frob_dispatcher(input, bias, option=None):
            return (input, bias)

    The dispatcher must return an iterable, so return a single-element tuple if
    only one argument should be checked. We would then modify the original
    definition for torch.frob to look like:

        @torch_function_dispatch(_frob_dispatcher)
        def frob(input, bias, option=None):
             return input + bias

    See ``torch/functional.py`` for more usage examples.

    Parameters
    ----------
    dispatcher : callable
        Function that when called like ``dispatcher(*args, **kwargs)`` with
        arguments from the NumPy function call returns an iterable of
        array-like arguments to check for ``__torch_function__``.
    module : str, optional
        ``__module__`` attribute to set on new function, e.g.,
        ``module='torch'``.  By default, module is copied from the decorated
        function.
    verify : bool, optional
        If True, verify the that the signature of the dispatcher and decorated
        function signatures match exactly: all required and optional arguments
        should appear in order with the same names, but the default values for
        all optional arguments should be ``None``. Only disable verification
        if the dispatcher's signature needs to deviate for some particular
        reason, e.g., because the function has a signature like
        ``func(*args, **kwargs)``.

    Returns
    -------
    dispatcher : callable
        Function suitable for decorating the implementation of a NumPy
        function.

    Notes
    -----
    The dispatcher should normally return a tuple containing all input
    arguments that may have a ``__torch_function__`` attribute.

    In some cases where that's not easily possible, e.g. ``torch.cat``, it is
    also valid (if a little slower) to make the dispatcher function a generator
    (i.e. use ``yield`` to return arguments one by one).

    """
    def decorator(implementation):
        if verify:
            _verify_matching_signatures(implementation, dispatcher)

        # Equivalently, we could define this function directly instead of using
        # exec. This version has the advantage of giving the helper function a
        # more interpretable name. Otherwise, the original function does not
        # show up at all in many cases, e.g., if it's written in C++ or if the
        # dispatcher gets an invalid keyword argument.
        source = _wrapped_func_source.format(name=implementation.__name__)

        source_object = compile(
            source, filename='<__torch_function__ internals>', mode='exec')
        scope = {
            'implementation': implementation,
            'dispatcher': dispatcher,
            'functools': functools,
            'implement_torch_function': _implement_torch_function,
        }
        _six.exec_(source_object, scope)

        public_api = scope[implementation.__name__]

        if module is not None:
            public_api.__module__ = module

        public_api._implementation = implementation

        return public_api

    return decorator