import torch


def Ax(A, x):
    """Returns the Matrix-vector product Ax.
    
    Arguments:
        A (function or torch.Tensor): Matrix-vector product function or 
            explicit matrix.
        x (torch.Tensor): A vector.
    """
    if callable(A):
        Ax = A(x)
    else:
        Ax = torch.mm(A, x)
    return Ax