import torch


def sample_flexible(positive, negative, batch_size, pos_percent):
    """
    Sample positives up to pos_percent, if positive num is not enough,
    sample more negatives to make it up to a batch_size
    """
    expected_pos_num = int(batch_size * pos_percent)
    if positive.numel() > expected_pos_num:
        shuffle_pos = torch.randperm(positive.numel(), device=positive.device)[:expected_pos_num]
        positive = positive[shuffle_pos]

    expected_neg_num = batch_size - positive.numel()
    if negative.numel() > expected_neg_num:
        shuffle_neg = torch.randperm(negative.numel(), device=negative.device)[:expected_neg_num]
        negative = negative[shuffle_neg]
    return positive, negative