def METHOD_NAME(dataset, num_samples):
    if num_samples in _sampled_datasets:
        return _sampled_datasets[num_samples]
    with dataset._graph.as_default():
        SHUFFLE_BUFFER_SIZE = 300 # NOTE: Adjust the buffer size as necessary.
        SHUFFLE_SEED = 22222
        dataset = dataset.shuffle(buffer_size=SHUFFLE_BUFFER_SIZE, seed=SHUFFLE_SEED)\
                         .take(num_samples)\
                         .batch(BATCH_SIZE)
        _sampled_datasets[num_samples] = dataset
        return dataset