| 748 | repeat=(None, 1, 2), |
| 749 | ) |
| 750 | def test_repeat(self, *, is_training, repeat): |
| 751 | ds = _text_ds(["a", "b", "c"]) |
| 752 | ds = batch( |
| 753 | global_batch_size=2, |
| 754 | is_training=is_training, |
| 755 | pad_example_fn=default_pad_example_fn, |
| 756 | repeat=repeat, |
| 757 | )(ds) |
| 758 | batch_index = 0 |
| 759 | for input_batch in ds: |
| 760 | if is_training or batch_index % 2 == 0: |
| 761 | self.assertSequenceEqual(input_batch["text"].numpy().tolist(), [b"a", b"b"]) |
| 762 | self.assertSequenceEqual(input_batch["index"].numpy().tolist(), [0, 1]) |
| 763 | else: |
| 764 | # The eval dataset will be padded by empty examples. |
| 765 | self.assertSequenceEqual(input_batch["text"].numpy().tolist(), [b"c", b""]) |
| 766 | self.assertSequenceEqual(input_batch["index"].numpy().tolist(), [2, 0]) |
| 767 | batch_index += 1 |
| 768 | if batch_index >= 10: |
| 769 | break |
| 770 | if repeat is None: |
| 771 | # Repeat indefinitely if is_training, otherwise do not repeat |
| 772 | # (hence 2 batches after padding). |
| 773 | self.assertEqual(10 if is_training else 2, batch_index) |
| 774 | else: |
| 775 | # If is_training, we discard remaining examples, hence one batch per epoch. |
| 776 | # Otherwise we have two batches per epoch. |
| 777 | self.assertEqual(repeat if is_training else 2 * repeat, batch_index) |
| 778 | |
| 779 | |
| 780 | class UnpackTest(test_utils.TestCase): |