import torch


def matrix_from_mvp(mvp, dims, device=None):
    """Compute matrix representation from matrix-vector product.

    Parameters:
    -----------
    mvp : function
        Matrix-vector multiplication routine.
    dims : tuple(int)
        Dimensions of the matrix

    Returns:
    --------
    torch.Tensor
        Matrix representation on ``device``
    """
    if device is None:
        device = torch.device("cpu")

    assert len(dims) == 2
    matrix = torch.zeros(*dims).to(device)
    for i in range(dims[1]):
        v = torch.zeros(dims[0]).to(device)
        v[i] = 1.0
        matrix[:, i] = mvp(v)
    return matrix