import torch


def make_A(As, ns, device=torch.device("cpu"), dtype=torch.float32):
    """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.
        device: torch device (default cpu).
        dtype: torch dtype (default float32).
    """
    A = torch.zeros(len(ns), ns.max(), As[0].shape[0], device=device, dtype=dtype)
    for j, n_j in enumerate(ns):
        # Fill A[j] with A_j'.
        A[j, :n_j, :] = As[j].t()
    return A