def testCombineOnBatchSimple(self): batch_1 = [np.ones((2, 6))] batch_2 = [np.ones((1, 6))] out = [3 for _ in range(6)] analyzer = impl._CombineFnWrapper( analyzers._NumPyCombinerSpec(np.sum, reduce_instance_dims=False)) self.assertCombine(analyzer, [batch_1, batch_2], out)
def testCombineOnBatchLotsOfData(self): shards = [[np.ones(3)] for _ in range( beam_impl._DEFAULT_DESIRED_BATCH_SIZE * 2)] out = [1 for _ in range(3)] analyzer = impl._CombineFnWrapper( analyzers._NumPyCombinerSpec(np.min, reduce_instance_dims=False)) self.assertCombine(analyzer, shards, out)
def testCombineOnBatchLotsOfData(self): shards = [[np.ones((1, 3))] for _ in range(2000)] out = [1 for _ in range(3)] analyzer = impl._CombineFnWrapper( analyzers._NumPyCombinerSpec(np.min, reduce_instance_dims=False, output_dtypes=[np.int64])) self.assertCombine(analyzer, shards, out)
def testCombineOnBatchWithBeamPipeline(self): # Test with a real Beam pipeline instead of calling the Combiner methods # directly. This surfaces bugs that only occur within a Beam pipeline, e.g. # due to Beam passing iterators to merge_accumulators instead of lists. with beam.Pipeline() as p: batch_1 = [np.ones((2, 6), dtype=np.int)] batch_2 = [np.ones((1, 6), dtype=np.int)] expected_output = np.ones(6) * 3 def assert_equals_expected(outputs): output, = outputs # Expect exactly one analyzer output return np.array_equal(output, expected_output) analyzer = impl._CombineFnWrapper( analyzers._NumPyCombinerSpec(np.sum, reduce_instance_dims=False)) assert_that(p | beam.Create([batch_1, batch_2]) | beam.CombineGlobally(analyzer) | beam.Map(assert_equals_expected), equal_to([True]))
def testCombineOnBatchAllEmptyRow(self): analyzer = impl._CombineFnWrapper( analyzers._NumPyCombinerSpec(np.sum, reduce_instance_dims=False)) self.assertCombine(analyzer, [[[[]]], [[[]]], [[[]]]], [])