def test_reduces_unplaced_lambda_to_equivalent_tf(self):
lam = building_blocks.Lambda('x', tf.int32,
building_blocks.Reference('x', tf.int32))
extracted_tf = mapreduce_transformations.consolidate_and_extract_local_processing(
lam)
executable_tf = computation_wrapper_instances.building_block_to_computation(
extracted_tf)
executable_lam = computation_wrapper_instances.building_block_to_computation(
lam)
for k in range(10):
self.assertEqual(executable_tf(k), executable_lam(k))
def test_identity_lambda_executes_as_identity(self):
lam = building_blocks.Lambda('x', tf.int32,
building_blocks.Reference('x', tf.int32))
computation_impl_lambda = computation_wrapper_instances.building_block_to_computation(
lam)
for k in range(10):
self.assertEqual(computation_impl_lambda(k), k)
def test_prune_does_not_change_exeuction(self):
def bad_fn(x):
_ = tf.constant(0)
return x
comp = _create_compiled_computation(bad_fn, tf.int32)
reduced_proto = proto_transformations.prune_tensorflow_proto(
comp.proto)
reduced_comp = building_blocks.CompiledComputation(reduced_proto)
orig_executable = computation_wrapper_instances.building_block_to_computation(
comp)
reduced_executable = computation_wrapper_instances.building_block_to_computation(
reduced_comp)
for k in range(5):
self.assertEqual(orig_executable(k), reduced_executable(k))
def test_converts_building_block_to_computation(self):
lam = building_blocks.Lambda('x', tf.int32,
building_blocks.Reference('x', tf.int32))
computation_impl_lambda = computation_wrapper_instances.building_block_to_computation(
lam)
self.assertIsInstance(computation_impl_lambda,
computation_impl.ComputationImpl)
def test_reduces_federated_value_at_clients_to_equivalent_noarg_function(
self):
federated_value = intrinsics.federated_value(0,
placements.CLIENTS)._comp
extracted_tf = mapreduce_transformations.consolidate_and_extract_local_processing(
federated_value)
executable_tf = computation_wrapper_instances.building_block_to_computation(
extracted_tf)
self.assertEqual(executable_tf(), 0)
def test_reduces_federated_apply_to_equivalent_function(self):
lam = building_blocks.Lambda('x', tf.int32,
building_blocks.Reference('x', tf.int32))
arg = building_blocks.Reference(
'arg', computation_types.FederatedType(tf.int32,
placements.CLIENTS))
mapped_fn = building_block_factory.create_federated_map_or_apply(
lam, arg)
extracted_tf = mapreduce_transformations.consolidate_and_extract_local_processing(
mapped_fn)
self.assertIsInstance(extracted_tf,
building_blocks.CompiledComputation)
executable_tf = computation_wrapper_instances.building_block_to_computation(
extracted_tf)
executable_lam = computation_wrapper_instances.building_block_to_computation(
lam)
for k in range(10):
self.assertEqual(executable_tf(k), executable_lam(k))
def test_constructs_tensor_zero(self):
tensor_type = computation_types.TensorType(tf.float32, [2, 2])
tensor_zero = intrinsic_utils.create_generic_constant(tensor_type, 0)
self.assertEqual(tensor_zero.type_signature, tensor_type)
self.assertIsInstance(tensor_zero, computation_building_blocks.Call)
executable_noarg_fn = computation_wrapper_instances.building_block_to_computation(
tensor_zero.function)
self.assertTrue(np.array_equal(executable_noarg_fn(), np.zeros([2,
2])))
def test_tensor_computation_fails_well(self):
cf = test_utils.get_temperature_sensor_example()
it = canonical_form_utils.get_iterative_process_for_canonical_form(cf)
init_result = it.initialize.type_signature.result
lam = building_blocks.Lambda(
'x', init_result, building_blocks.Reference('x', init_result))
bad_it = computation_utils.IterativeProcess(
it.initialize,
computation_wrapper_instances.building_block_to_computation(lam))
with self.assertRaisesRegex(TypeError,
'instances of `tff.NamedTupleType`.'):
canonical_form_utils.get_canonical_form_for_iterative_process(
bad_it)
def test_create_unnamed_tuple_zero(self):
tuple_type = [computation_types.TensorType(tf.float32, [2, 2])] * 2
tuple_zero = intrinsic_utils.construct_generic_constant(tuple_type, 0)
self.assertEqual(tuple_zero.type_signature,
computation_types.to_type(tuple_type))
self.assertIsInstance(tuple_zero, computation_building_blocks.Call)
executable_noarg_fn = computation_wrapper_instances.building_block_to_computation(
tuple_zero.function)
self.assertLen(executable_noarg_fn(), 2)
self.assertTrue(
np.array_equal(executable_noarg_fn()[0], np.zeros([2, 2])))
self.assertTrue(
np.array_equal(executable_noarg_fn()[1], np.zeros([2, 2])))
def test_create_named_tuple_one(self):
tuple_type = [('a', computation_types.TensorType(tf.float32, [2, 2])),
('b', computation_types.TensorType(tf.float32, [2, 2]))]
tuple_zero = intrinsic_utils.create_generic_constant(tuple_type, 1)
self.assertEqual(tuple_zero.type_signature,
computation_types.to_type(tuple_type))
self.assertIsInstance(tuple_zero, computation_building_blocks.Call)
executable_noarg_fn = computation_wrapper_instances.building_block_to_computation(
tuple_zero.function)
self.assertLen(executable_noarg_fn(), 2)
self.assertTrue(
np.array_equal(executable_noarg_fn().a, np.ones([2, 2])))
self.assertTrue(
np.array_equal(executable_noarg_fn().b, np.ones([2, 2])))
def test_create_federated_tensor_one(self):
fed_type = computation_types.FederatedType(
computation_types.TensorType(tf.float32, [2, 2]),
placement_literals.CLIENTS)
fed_zero = intrinsic_utils.construct_generic_constant(fed_type, 1)
self.assertEqual(fed_zero.type_signature.member, fed_type.member)
self.assertEqual(fed_zero.type_signature.placement, fed_type.placement)
self.assertTrue(fed_zero.type_signature.all_equal)
self.assertIsInstance(fed_zero, computation_building_blocks.Call)
self.assertIsInstance(fed_zero.function,
computation_building_blocks.Intrinsic)
self.assertEqual(fed_zero.function.uri,
intrinsic_defs.FEDERATED_VALUE_AT_CLIENTS.uri)
self.assertIsInstance(fed_zero.argument,
computation_building_blocks.Call)
executable_unplaced_fn = computation_wrapper_instances.building_block_to_computation(
fed_zero.argument.function)
self.assertTrue(
np.array_equal(executable_unplaced_fn(), np.ones([2, 2])))
def test_create_named_tuple_of_federated_tensors_zero(self):
fed_type = computation_types.FederatedType(
computation_types.TensorType(tf.float32, [2, 2]),
placement_literals.CLIENTS,
all_equal=True)
tuple_type = [('a', fed_type), ('b', fed_type)]
zero = intrinsic_utils.construct_generic_constant(tuple_type, 0)
fed_zero = zero.argument[0]
self.assertEqual(zero.type_signature,
computation_types.to_type(tuple_type))
self.assertIsInstance(fed_zero.function,
computation_building_blocks.Intrinsic)
self.assertEqual(fed_zero.function.uri,
intrinsic_defs.FEDERATED_VALUE_AT_CLIENTS.uri)
self.assertIsInstance(fed_zero.argument,
computation_building_blocks.Call)
executable_unplaced_fn = computation_wrapper_instances.building_block_to_computation(
fed_zero.argument.function)
self.assertTrue(
np.array_equal(executable_unplaced_fn(), np.zeros([2, 2])))
def test_create_federated_named_tuple_one(self):
tuple_type = [('a', computation_types.TensorType(tf.float32, [2, 2])),
('b', computation_types.TensorType(tf.float32, [2, 2]))]
fed_type = computation_types.FederatedType(tuple_type,
placement_literals.SERVER)
fed_zero = intrinsic_utils.create_generic_constant(fed_type, 1)
self.assertEqual(fed_zero.type_signature.member, fed_type.member)
self.assertEqual(fed_zero.type_signature.placement, fed_type.placement)
self.assertTrue(fed_zero.type_signature.all_equal)
self.assertIsInstance(fed_zero, computation_building_blocks.Call)
self.assertIsInstance(fed_zero.function,
computation_building_blocks.Intrinsic)
self.assertEqual(fed_zero.function.uri,
intrinsic_defs.FEDERATED_VALUE_AT_SERVER.uri)
self.assertIsInstance(fed_zero.argument,
computation_building_blocks.Call)
executable_unplaced_fn = computation_wrapper_instances.building_block_to_computation(
fed_zero.argument.function)
self.assertLen(executable_unplaced_fn(), 2)
self.assertTrue(
np.array_equal(executable_unplaced_fn().a, np.ones([2, 2])))
self.assertTrue(
np.array_equal(executable_unplaced_fn().b, np.ones([2, 2])))
def test_broadcast_dependent_on_aggregate_fails_well(self):
cf = test_utils.get_temperature_sensor_example()
it = canonical_form_utils.get_iterative_process_for_canonical_form(cf)
next_comp = test_utils.computation_to_building_block(it.next)
top_level_param = building_blocks.Reference(next_comp.parameter_name,
next_comp.parameter_type)
first_result = building_blocks.Call(next_comp, top_level_param)
middle_param = building_blocks.Tuple([
building_blocks.Selection(first_result, index=0),
building_blocks.Selection(top_level_param, index=1)
])
second_result = building_blocks.Call(next_comp, middle_param)
not_reducible = building_blocks.Lambda(next_comp.parameter_name,
next_comp.parameter_type,
second_result)
not_reducible_it = computation_utils.IterativeProcess(
it.initialize,
computation_wrapper_instances.building_block_to_computation(
not_reducible))
with self.assertRaisesRegex(ValueError,
'broadcast dependent on aggregate'):
canonical_form_utils.get_canonical_form_for_iterative_process(
not_reducible_it)
def test_raises_on_none(self):
with self.assertRaises(TypeError):
computation_wrapper_instances.building_block_to_computation(None)
