gelfista.py 文件源码

def make_A(As, ns, use_gpu=False):
    """Create the 3D tensor A as needed by gel_solve, given a list of feature
        matrices.

    Arguments:
        As: list of feature matrices, one per group (size mxn_j).
        ns: LongTensor of group sizes.
        use_gpu: move the final tensor to GPU.
    """
    A = torch.zeros(len(ns), ns.max(), As[0].size()[0])
    for j, n_j in enumerate(ns):
        # Fill A[j] with A_j.T
        A_j = As[j]
        A[j, :n_j, :] = A_j.t()
    if use_gpu:
        A = A.cuda()
    return A