def select_coefs(model_coefs, idx):
    """
    Auxiliary function for compute_couplings().
    Used to select the desired subvector/submatrix from the vector/matrix
    of model coefficients.
    """
    if len(model_coefs.shape) == 1:  # Binomial case
        sel_coefs = model_coefs[range(idx, idx + 20)]
    else:  # Multinomial case
        sel_coefs = model_coefs[:, range(idx, idx + 20)]
    return sel_coefs, idx + 20