import torch


def kumaraswamy_sample(conc1, conc0, batch_shape):
    """
    Sample from the Kumaraswamy distribution given the concentrations

    :param conc1: torch tensor: the a concentration of the distribution
    :param conc0: torch tensor: the b concentration of the distribution
    :param batch_shape: scalar: the batch shape for the samples

    :return: torch tensor: a sample from the Kumaraswamy distribution
    """

    device = torch.device("cuda" if conc1.is_cuda else "cpu")

    x = (0.01 - 0.99) * torch.rand(batch_shape, conc1.size(0)).to(device) + 0.99
    q_u = (1-(1-x)**(1./conc0))**(1./conc1)

    return q_u