Inverse transform sampling

sample.py

class InverseTransform:
    def __init__(self, probs, rng=None):
        if np.sum(probs) < 1.:
            raise ValueError('Must sum up to 1')
        
        self.rng = rng
        if rng is None:
            self.rng = np.random.RandomState()
        self.cdf = np.cumsum(probs)
        
    def __call__(self, size=None):
        return np.argmax(np.greater(self.cdf, self.rng.rand(size)[:, np.newaxis]), axis=1)   

You should use searchsorted rather than argmax in line 12! It works because cdfs are sorted. That will give you a runtime of size * log |probs| rather than size * |probs|!