def test_splits_on_selected_intrinsic_broadcast(self):
federated_broadcast = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
called_intrinsics = building_blocks.Struct([federated_broadcast])
comp = building_blocks.Lambda('a', tf.int32, called_intrinsics)
call = building_block_factory.create_null_federated_broadcast()
self.assert_splits_on(comp, call)
def test_splits_on_intrinsic_noarg_function(self):
federated_broadcast = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
called_intrinsics = building_blocks.Struct([federated_broadcast])
comp = building_blocks.Lambda(None, None, called_intrinsics)
call = building_block_factory.create_null_federated_broadcast()
self.assert_splits_on(comp, call)
def test_splits_on_selected_intrinsic_nested_in_tuple_broadcast(self):
first_broadcast = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
packed_broadcast = building_blocks.Struct([
building_blocks.Data('a', computation_types.at_server(tf.int32)),
first_broadcast
])
sel = building_blocks.Selection(packed_broadcast, index=0)
second_broadcast = building_block_factory.create_federated_broadcast(
sel)
result = transformations.to_call_dominant(second_broadcast)
comp = building_blocks.Lambda('a', tf.int32, result)
call = building_block_factory.create_null_federated_broadcast()
self.assert_splits_on(comp, call)
def test_does_not_find_aggregate_dependent_on_broadcast(self):
broadcast = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
value_type = broadcast.type_signature
zero = building_blocks.Data('zero', value_type.member)
accumulate_result = building_blocks.Data('accumulate_result',
value_type.member)
accumulate = building_blocks.Lambda(
'accumulate_parameter', [value_type.member, value_type.member],
accumulate_result)
merge_result = building_blocks.Data('merge_result', value_type.member)
merge = building_blocks.Lambda('merge_parameter',
[value_type.member, value_type.member],
merge_result)
report_result = building_blocks.Data('report_result',
value_type.member)
report = building_blocks.Lambda('report_parameter', value_type.member,
report_result)
aggregate_dependent_on_broadcast = building_block_factory.create_federated_aggregate(
broadcast, zero, accumulate, merge, report)
tree_analysis.check_aggregate_not_dependent_on_aggregate(
aggregate_dependent_on_broadcast)
def test_returns_false_with_unmatched_called_intrinsic(self):
comp = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
uri = intrinsic_defs.FEDERATED_MAP.uri
self.assertFalse(tree_analysis.contains_called_intrinsic(comp, uri))
def test_returns_true_with_matching_uri(self):
comp = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
uri = intrinsic_defs.FEDERATED_BROADCAST.uri
self.assertTrue(tree_analysis.contains_called_intrinsic(comp, uri))
def test_returns_true_with_none_uri(self):
comp = building_block_test_utils.create_whimsy_called_federated_broadcast(
)
self.assertTrue(tree_analysis.contains_called_intrinsic(comp))
def test_returns_true_for_tuples(self):
data_1 = building_blocks.Data('data', tf.int32)
tuple_1 = building_blocks.Struct([data_1, data_1])
data_2 = building_blocks.Data('data', tf.int32)
tuple_2 = building_blocks.Struct([data_2, data_2])
self.assertTrue(tree_analysis.trees_equal(tuple_1, tuple_2))
def test_returns_true_for_identical_graphs_with_nans(self):
tf_comp1 = _create_tensorflow_graph_with_nan()
tf_comp2 = _create_tensorflow_graph_with_nan()
self.assertTrue(tree_analysis.trees_equal(tf_comp1, tf_comp2))
non_aggregation_intrinsics = building_blocks.Struct([
(None,
building_block_test_utils.create_whimsy_called_federated_broadcast()),
(None,
building_block_test_utils.create_whimsy_called_federated_value(
placements.CLIENTS))
])
unit = computation_types.StructType([])
trivial_aggregate = building_block_test_utils.create_whimsy_called_federated_aggregate(
value_type=unit)
trivial_mean = building_block_test_utils.create_whimsy_called_federated_mean(
unit)
trivial_sum = building_block_test_utils.create_whimsy_called_federated_sum(
unit)
# TODO(b/120439632) Enable once federated_mean accepts structured weights.
# trivial_weighted_mean = ...
trivial_secure_sum = building_block_test_utils.create_whimsy_called_federated_secure_sum_bitwidth(