import torch


def MaskedDenseMM(X, Y, rows, cols):
    """Helper routine for Backward computation for spmm
    Z = X[rows,:] x Y[cols,:]
    Args:
        X: dense matrix
        Y: dense matrix
        rows: row indices of X
        cols: row indices of Y

    Returns:
        Z as defined above
    """
    #Z = X * Y
    m = rows.shape[0]
    assert(m == cols.shape[0])
    assert(X.shape == Y.shape)
    xrows = X[rows,:]
    yrows = Y[cols,:]
    Z = torch.sum(xrows * yrows, dim=1)
    return(Z)