import torch


def wrap_operator(A):
    """ Wraps linear operator as function handle. """
    if isinstance(A, torch.Tensor):

        def _Afunc(x):
            if x.dtype == torch.double:
                return torch.matmul(x, A.T.double())
            else:
                return torch.matmul(x, A.T)

    elif callable(A):
        _Afunc = A

    else:
        raise ValueError(
            "Unable to wrap as a function handle. "
            "Expected a Torch tensor or callable but got "
            "{} instead.".format(type(A).__name__)
        )
    return _Afunc