In order to simplify the writting of universal functions it could be great to have a decorator function which hide the technical part of the code (convertion of input and output of the wrapped function/method).
For example, the code:

def my_universal_function(a, b, c):
    # Convert all inputs to EagerPy tensors
    a, b, c = ep.astensors(a, b, c)

    # performs some computations
    result = (a + b * c).square()

    # and return a native tensor
    return result.raw

would become:

@eager_function
def my_universal_function(a, b, c):
    return (a + b * c).square()

In addition, we could add the feature that if the input tensors are already eagerpy tensors, then no convertion to raw format should done on the output tensors.

I wrote a prototype of such a decorator function. It should not work on any type of arguments and so its usage would require that the wrapped function has a rather "simple" signature (with args and kwargs constituted of tensors or nested containers with tensors on leaves: dict, list, tuple or namedtuple like containers).

Would you consider to have this feature in eagerpy?