def METHOD_NAME(input_context=None, batch_size=None):
    inputs, output = get_numpy()
    dataset = tf.data.Dataset.from_tensor_slices((inputs, output))
    if input_context:
        dataset = dataset.shard(
            input_context.num_input_pipelines, input_context.input_pipeline_id
        )
    if batch_size is None:
        batch_size = BATCH_SIZE
    dataset = dataset.batch(batch_size).repeat(200)
    return dataset