def test_batches_stream(self): stream = iter([(0, 1), (2, 3), (4, 5), (6, 7)]) batched_stream = simple.batch_stream(stream, batch_size=2) np.testing.assert_equal(next(batched_stream), (np.array([0, 2]), np.array([1, 3]))) np.testing.assert_equal(next(batched_stream), (np.array([4, 6]), np.array([5, 7])))
def mix_and_batch(streams): (init_stream, own_stream) = streams mixed_stream = simple.mix_streams(init_stream, own_stream, mix_prob) return simple.batch_stream(mixed_stream, self._model_train_batch_size)