import torch


def get_p_star(samples, y, beta):
    """Get the values for the p_star distribution

    Parameters:
    samples: tensor of dim (batch_size x num_samples) with the samples indexes
    y: tensor of dim (batch_size) with the true label for each example
    beta: the strength of the feedback

    Returns:
    p_star: tensor of dim(batch_size x num_samples ) with +beta for correct samples and -beta for incorrect ones. 
   """    
    #Convert y to (batch_size x 1) dimensionality for comparison with samples
    y = y.unsqueeze(-1)
    positives = (samples==y).float()*beta
    negatives = torch.logical_not(samples==y).float()*-beta
    p_star = positives + negatives
    return p_star