def test_raises_on_selections_at_different_placements(self):
fed_at_clients = computation_types.FederatedType(tf.int32,
placements.CLIENTS)
fed_at_server = computation_types.FederatedType(tf.int32, placements.SERVER)
tuple_of_federated_types = computation_types.StructType(
[fed_at_clients, fed_at_server])
lam = building_blocks.Lambda(
'x', tuple_of_federated_types,
building_blocks.Selection(
building_blocks.Reference('x', tuple_of_federated_types), index=0))
with self.assertRaises(form_utils._MismatchedSelectionPlacementError):
form_utils._as_function_of_some_federated_subparameters(lam, [(0,), (1,)])
def test_multiple_nested_named_element_selection(self):
fed_at_clients = computation_types.FederatedType(
tf.int32, placements.CLIENTS)
fed_at_server = computation_types.FederatedType(
tf.int32, placements.SERVER)
tuple_of_federated_types = computation_types.StructType([
('a', [('a', fed_at_clients)]), ('b', fed_at_server),
('c', [('c', fed_at_clients)])
])
first_selection = building_blocks.Selection(building_blocks.Selection(
building_blocks.Reference('x',
tuple_of_federated_types), name='a'),
name='a')
second_selection = building_blocks.Selection(building_blocks.Selection(
building_blocks.Reference('x',
tuple_of_federated_types), name='c'),
name='c')
lam = building_blocks.Lambda(
'x', tuple_of_federated_types,
building_blocks.Struct([first_selection, second_selection]))
new_lam = form_utils._as_function_of_some_federated_subparameters(
lam, [(0, 0), (2, 0)])
expected_parameter_type = computation_types.at_clients(
(tf.int32, tf.int32))
type_test_utils.assert_types_equivalent(
new_lam.type_signature,
computation_types.FunctionType(expected_parameter_type,
lam.result.type_signature))
def test_single_element_selection_leaves_no_unbound_references(self):
fed_at_clients = computation_types.FederatedType(tf.int32,
placements.CLIENTS)
fed_at_server = computation_types.FederatedType(tf.int32, placements.SERVER)
tuple_of_federated_types = computation_types.StructType(
[fed_at_clients, fed_at_server])
lam = building_blocks.Lambda(
'x', tuple_of_federated_types,
building_blocks.Selection(
building_blocks.Reference('x', tuple_of_federated_types), index=0))
new_lam = form_utils._as_function_of_some_federated_subparameters(
lam, [(0,)])
unbound_references = transformation_utils.get_map_of_unbound_references(
new_lam)[new_lam]
self.assertEmpty(unbound_references)
def test_single_element_selection(self):
fed_at_clients = computation_types.FederatedType(tf.int32,
placements.CLIENTS)
fed_at_server = computation_types.FederatedType(tf.int32, placements.SERVER)
tuple_of_federated_types = computation_types.StructType(
[fed_at_clients, fed_at_server])
lam = building_blocks.Lambda(
'x', tuple_of_federated_types,
building_blocks.Selection(
building_blocks.Reference('x', tuple_of_federated_types), index=0))
new_lam = form_utils._as_function_of_some_federated_subparameters(
lam, [(0,)])
expected_parameter_type = computation_types.at_clients((tf.int32,))
self.assert_types_equivalent(
new_lam.type_signature,
computation_types.FunctionType(expected_parameter_type,
lam.result.type_signature))
def test_binding_multiple_args_results_in_unique_names(self):
fed_at_clients = computation_types.FederatedType(
tf.int32, placements.CLIENTS)
fed_at_server = computation_types.FederatedType(
tf.int32, placements.SERVER)
tuple_of_federated_types = computation_types.StructType(
[[fed_at_clients], fed_at_server, [fed_at_clients]])
first_selection = building_blocks.Selection(building_blocks.Selection(
building_blocks.Reference('x', tuple_of_federated_types), index=0),
index=0)
second_selection = building_blocks.Selection(building_blocks.Selection(
building_blocks.Reference('x', tuple_of_federated_types), index=2),
index=0)
lam = building_blocks.Lambda(
'x', tuple_of_federated_types,
building_blocks.Struct([first_selection, second_selection]))
new_lam = form_utils._as_function_of_some_federated_subparameters(
lam, [(0, 0), (2, 0)])
tree_analysis.check_has_unique_names(new_lam)
def test_raises_on_non_federated_selection(self):
lam = building_blocks.Lambda(
'x', [tf.int32], building_blocks.Reference('x', [tf.int32]))
with self.assertRaises(form_utils._NonFederatedSelectionError):
form_utils._as_function_of_some_federated_subparameters(
lam, [(0, )])