import torch


def getBatch(batch_size, correlations, anticorrelations):
    """Make inputs and outputs for a single output model.

    Arguments:

        batch_size (int)               : Number of elements in a batch.
        correlations     (torch.tensor): Probabilities of each input being 1 when the target output
                                         is 1.
        anticorrelations (torch.tensor): Probabilities of each input being 1 when the target output
                                         is 0. Must be the same length as correlations.
    """
    assert len(correlations) == len(anticorrelations)
    # Make some 0s and 1s as the output
    outputs = torch.empty((batch_size, 1)).random_(0, 2).cuda()

    # Randomly select which inputs would be active if the output is 0 or 1
    probs = correlations.expand(batch_size, -1) * outputs + anticorrelations.expand(batch_size, -1) * (1. - outputs)
    return torch.bernoulli(probs), outputs