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)