def test_returns_comps_with_federated_aggregate_no_unbound_references(
self):
federated_aggregate = test_utils.create_dummy_called_federated_aggregate(
accumulate_parameter_name='a',
merge_parameter_name='b',
report_parameter_name='c')
tup = building_blocks.Tuple([
federated_aggregate,
federated_aggregate,
])
comp = building_blocks.Lambda('d', tf.int32, tup)
uri = intrinsic_defs.FEDERATED_AGGREGATE.uri
before, after = mapreduce_transformations.force_align_and_split_by_intrinsic(
comp, uri)
def _predicate(comp):
return building_block_analysis.is_called_intrinsic(comp, uri)
self.assertIsInstance(comp, building_blocks.Lambda)
self.assertGreater(tree_analysis.count(comp, _predicate), 0)
self.assertIsInstance(before, building_blocks.Lambda)
self.assertEqual(tree_analysis.count(before, _predicate), 0)
self.assertEqual(before.parameter_type, comp.parameter_type)
self.assertIsInstance(after, building_blocks.Lambda)
self.assertEqual(tree_analysis.count(after, _predicate), 0)
self.assertEqual(after.result.type_signature,
comp.result.type_signature)
def _count_intrinsics(comp, uri):
def _predicate(comp):
return (isinstance(comp, building_blocks.Intrinsic) and uri is not None and
comp.uri == uri)
return tree_analysis.count(comp, _predicate)
def test_returns_comps_with_federated_aggregate(self):
iterative_process = test_utils.construct_example_training_comp()
comp = test_utils.computation_to_building_block(iterative_process.next)
uri = intrinsic_defs.FEDERATED_AGGREGATE.uri
before, after = mapreduce_transformations.force_align_and_split_by_intrinsic(
comp, uri)
def _predicate(comp):
return building_block_analysis.is_called_intrinsic(comp, uri)
self.assertIsInstance(comp, building_blocks.Lambda)
self.assertGreater(tree_analysis.count(comp, _predicate), 0)
self.assertIsInstance(before, building_blocks.Lambda)
self.assertEqual(tree_analysis.count(before, _predicate), 0)
self.assertEqual(before.parameter_type, comp.parameter_type)
self.assertIsInstance(after, building_blocks.Lambda)
self.assertEqual(tree_analysis.count(after, _predicate), 0)
self.assertEqual(after.result.type_signature, comp.result.type_signature)
def test_adds_all_nodes_to_set_with_constant_true_predicate(self):
nested_tree = building_block_test_utils.create_nested_syntax_tree()
all_nodes = tree_analysis.extract_nodes_consuming(
nested_tree, lambda x: True)
node_count = tree_analysis.count(nested_tree)
self.assertLen(all_nodes, node_count)
def _count_called_intrinsics(comp, uri=None):
def _predicate(comp):
return building_block_analysis.is_called_intrinsic(comp, uri)
return tree_analysis.count(comp, _predicate)
def _count_broadcasts(comp):
"""Returns the number of called federated broadcasts found in `comp`."""
uri = intrinsic_defs.FEDERATED_BROADCAST.uri
predicate = lambda x: building_block_analysis.is_called_intrinsic(x, uri)
return tree_analysis.count(comp, predicate)