import torch


def stack_feats(feats):
    """
    Args:
        feats: [B, (n, c)]
    Returns:
        feats: (B, N, c)
        valid: (B, N)
    """
    max_num = max([len(f) for f in feats])
    s0 = feats[0].shape
    for f in feats:
        assert f.shape[1:] == s0[1:], f"{f.shape} vs {s0}"

    shape = (max_num,) + feats[0].shape[1:]
    new_feats = []
    valids = []
    for feat in feats:
        new_feat = torch.zeros(shape, dtype=feat.dtype).to(feat.device)
        valid = torch.zeros(max_num, dtype=torch.bool).to(feat.device)
        new_feat[:len(feat)] = feat
        valid[: len(feat)] = 1
        new_feats.append(new_feat)
        valids.append(valid)
    return torch.stack(new_feats), torch.stack(valids)