import random


def get_conditional_sampler(x, y):
    """ Returns a function which allows randomly sampling images conditioned on the class. """
    x_split = [x[y == i].float() for i in range(10)]

    def sampler(digit):
        assert 0 <= digit <= 9, f"{digit} is not 0-9"
        idx = random.randint(0, x_split[digit].shape[0]-1)
        return x_split[digit][idx]
    return sampler