def test_batch_fn(self):
dataset = test_dataset_ints([32])
dataset = dataset.repeat(10)
batches = inputs.batch_fn(
dataset, True, ([None], [None]), [], 1, batch_size=10)
count = 0
for example in tfds.as_numpy(batches):
count += 1
self.assertEqual(example[0].shape[0], 10) # Batch size = 10.
self.assertEqual(count, 1) # Just one batch here.
def test_batch_fn_n_devices(self):
dataset = test_dataset_ints([32])
dataset = dataset.repeat(9)
batches = inputs.batch_fn(
dataset, True, ([None], [None]), [], 9, batch_size=10)
count = 0
for example in tfds.as_numpy(batches):
count += 1
# Batch size adjusted to be divisible by n_devices.
self.assertEqual(example[0].shape[0], 9)
self.assertEqual(count, 1) # Just one batch here.