'We can specify a batch size,
independent of the object initialization.'

def next(self, batch_size=(-1)):

    if (not self.batch_iter):
        self.batch_iter = self.get_homogenous_batch_iter(batch_size)
    try:
        batch = next(self.batch_iter)
        batch = add_random_variables_to_batch(self.state, self.rng, batch, self.prev_batch, self.evaluate_mode)
        self.prev_batch = batch
    except StopIteration:
        return None
    return batch
