def batch(self, batch_size):
"""Combines consecutive elements of this dataset into batches.
Args:
batch_size: A `tf.int64` scalar `tf.Tensor`, representing the number of
consecutive elements of this dataset to combine in a single batch.
Returns:
A `Dataset`.
"""
return Dataset(dataset_ops.BatchDataset(self._dataset, batch_size))
def _apply_fn(dataset):
ragged_dataset = _DenseToRaggedDataset(dataset, row_splits_dtype)
return dataset_ops.BatchDataset(ragged_dataset,
batch_size=batch_size,
drop_remainder=drop_remainder)
