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)