Examples
    --------
    >>> @flatten_args
    ... def f(x):
    ...     return 2*x

    >>> x, y, z = f(np.array([1., 2.]), 3., np.array([[1., 2.],[.5, .9]]))

    >>> x
    array([ 2.,  4.])

    >>> y
    6.0

    >>> z
    array([[ 2. ,  4. ],
           [ 1. ,  1.8]])

def flatten_args(fn):
    """
    Examples
    --------
    >>> @flatten_args
    ... def f(x):
    ...     return 2*x

    >>> x, y, z = f(np.array([1., 2.]), 3., np.array([[1., 2.],[.5, .9]]))

    >>> x
    array([ 2.,  4.])

    >>> y
    6.0

    >>> z
    array([[ 2. ,  4. ],
           [ 1. ,  1.8]])
    """
    @wraps(fn)
    def new_fn(*args):
        args_flat, shapes = flatten(args)
        result = fn(args_flat)
        return unflatten(result, shapes=shapes)

    return new_fn