def dataset_fn(ctx=None):
    """Returns tf.data.Dataset for distributed BERT pretraining."""
    batch_size = ctx.get_per_replica_batch_size(
        global_batch_size) if ctx else global_batch_size
    dataset = create_training_dataset(input_file, batch_size, config)
    return dataset
Returns tf.data.Dataset for distributed BERT pretraining.
language/compgen/csl/model/training/input_utils.py