Function decorator for specifying input and/or output array dtypes.

def dtype_contract(input_dtype=None, output_dtype=None):
    """Function decorator for specifying input and/or output array dtypes."""
    def wrap(function):
        @wraps(function)
        def wrapped_function(*args, **kwargs):
            if input_dtype is not None:
                check_dtype(args[0], input_dtype)
            array = function(*args, **kwargs)
            if output_dtype is not None:
                check_dtype(array, output_dtype)
            return array
        return wrapped_function
    return wrap