def augmentation_pipeline(*augment_funcs):
    """Takes an arbitrary number of functions with input (X, y, epoch, rng) and output (X, y),
    and returns a single function with the same signature which applies the input
    functions in order."""
    def augmentation_func(train_x, train_y, epoch, rng):
        for func in augment_funcs:
            train_x, train_y = func(train_x, train_y, epoch=epoch, rng=rng)
        return train_x, train_y

    return augmentation_func