def test_returned_tensorflow_executes_correctly_with_no_unbound_refs(self): concrete_int_type = computation_types.TensorType(tf.int32) concrete_int = building_block_factory.create_tensorflow_constant( concrete_int_type, 1) first_tf_id_type = computation_types.TensorType(tf.int32) first_tf_id = building_block_factory.create_compiled_identity( first_tf_id_type) called_tf_id = building_blocks.Call(first_tf_id, concrete_int) ref_to_call = building_blocks.Reference('call', called_tf_id.type_signature) second_tf_id_type = computation_types.TensorType(tf.int32) second_tf_id = building_block_factory.create_compiled_identity( second_tf_id_type) second_called = building_blocks.Call(second_tf_id, ref_to_call) ref_to_second_call = building_blocks.Reference( 'second_call', called_tf_id.type_signature) block_locals = [('call', called_tf_id), ('second_call', second_called)] block = building_blocks.Block( block_locals, building_blocks.Tuple([ref_to_second_call, ref_to_second_call])) tf_representing_block, _ = transformations.create_tensorflow_representing_block( block) result = test_utils.run_tensorflow( tf_representing_block.function.proto) self.assertAllEqual(result, [1, 1])
def test_returns_correct_structure_with_no_unbound_references(self): concrete_int_type = computation_types.TensorType(tf.int32) concrete_int = building_block_factory.create_tensorflow_constant( concrete_int_type, 1) first_tf_id_type = computation_types.TensorType(tf.int32) first_tf_id = building_block_factory.create_compiled_identity( first_tf_id_type) called_tf_id = building_blocks.Call(first_tf_id, concrete_int) ref_to_call = building_blocks.Reference('call', called_tf_id.type_signature) second_tf_id_type = computation_types.TensorType(tf.int32) second_tf_id = building_block_factory.create_compiled_identity( second_tf_id_type) second_called = building_blocks.Call(second_tf_id, ref_to_call) ref_to_second_call = building_blocks.Reference( 'second_call', called_tf_id.type_signature) block_locals = [('call', called_tf_id), ('second_call', second_called)] block = building_blocks.Block( block_locals, building_blocks.Tuple([ref_to_second_call, ref_to_second_call])) tf_representing_block, _ = transformations.create_tensorflow_representing_block( block) self.assertEqual(tf_representing_block.type_signature, block.type_signature) self.assertIsInstance(tf_representing_block, building_blocks.Call) self.assertIsInstance(tf_representing_block.function, building_blocks.CompiledComputation) self.assertIsNone(tf_representing_block.argument)
def test_returns_called_tf_computation_with_truct(self): constant_tuple_type = computation_types.StructType( [tf.int32, tf.float32]) constant_tuple = building_block_factory.create_tensorflow_constant( constant_tuple_type, 1) sel = building_blocks.Selection(source=constant_tuple, index=0) tup = building_blocks.Struct([sel, sel, sel]) self.assert_compiles_to_tensorflow(tup)
def test_returns_tf_computation_block_with_compiled_comp(self): concrete_int_type = computation_types.TensorType(tf.int32) tf_identity = building_block_factory.create_compiled_identity( concrete_int_type) unused_int = building_block_factory.create_tensorflow_constant( concrete_int_type, 1) block_to_id = building_blocks.Block([('x', unused_int)], tf_identity) self.assert_compiles_to_tensorflow(block_to_id)
def test_returns_tf_computation_with_functional_type_block_to_lambda_no_block( self): concrete_int_type = computation_types.TensorType(tf.int32) param = building_blocks.Reference('x', tf.float32) lam = building_blocks.Lambda(param.name, param.type_signature, param) unused_int = building_block_factory.create_tensorflow_constant( concrete_int_type, 1) blk_to_lam = building_blocks.Block([('y', unused_int)], lam) self.assert_compiles_to_tensorflow(blk_to_lam)
def test_reduces_federated_value_at_server_to_equivalent_noarg_function(self): zero = building_block_factory.create_tensorflow_constant( computation_types.TensorType(tf.int32, shape=[]), 0) federated_value = building_block_factory.create_federated_value( zero, placements.SERVER) extracted_tf = transformations.consolidate_and_extract_local_processing( federated_value, DEFAULT_GRAPPLER_CONFIG) executable_tf = computation_wrapper_instances.building_block_to_computation( extracted_tf) self.assertEqual(executable_tf(), 0)
def test_generates_tf_with_block(self): ref_to_x = building_blocks.Reference( 'x', computation_types.StructType([tf.int32, tf.float32])) identity_lambda = building_blocks.Lambda(ref_to_x.name, ref_to_x.type_signature, ref_to_x) tf_zero = building_block_factory.create_tensorflow_constant( computation_types.StructType([tf.int32, tf.float32]), 0) ref_to_z = building_blocks.Reference('z', [tf.int32, tf.float32]) called_lambda_on_z = building_blocks.Call(identity_lambda, ref_to_z) blk = building_blocks.Block([('z', tf_zero)], called_lambda_on_z) self.assert_compiles_to_tensorflow(blk)
def test_reduces_federated_value_at_clients_to_equivalent_noarg_function( self): zero = building_block_factory.create_tensorflow_constant( computation_types.TensorType(tf.int32, shape=[]), 0) federated_value = building_block_factory.create_federated_value( zero, placements.CLIENTS) federated_value_func = building_blocks.Lambda(None, None, federated_value) extracted_tf = compiler.consolidate_and_extract_local_processing( federated_value_func, DEFAULT_GRAPPLER_CONFIG) executable_tf = computation_impl.ConcreteComputation.from_building_block( extracted_tf) self.assertEqual(executable_tf(), 0)
def _construct_inlined_tuple(k): concrete_int = building_block_factory.create_tensorflow_constant( tf.int32, 1) first_tf_fn = building_block_factory.create_tensorflow_binary_operator( concrete_int.type_signature, tf.add) call = building_blocks.Call( first_tf_fn, building_blocks.Tuple([concrete_int, concrete_int])) for _ in range(k): # Simulating large TF computation call = building_blocks.Call( first_tf_fn, building_blocks.Tuple([call, call])) return building_blocks.Tuple([call, call])
def test_returns_called_tf_computation_with_non_functional_type(self): constant_tuple = building_block_factory.create_tensorflow_constant( [tf.int32, tf.float32], 1) sel = building_blocks.Selection(source=constant_tuple, index=0) tup = building_blocks.Tuple([sel, sel, sel]) transformed, modified_indicator = compiler_transformations.remove_duplicate_called_graphs( tup) self.assertTrue(modified_indicator) self.assertEqual(transformed.type_signature, tup.type_signature) self.assertIsInstance(transformed, building_blocks.Call) self.assertIsInstance(transformed.function, building_blocks.CompiledComputation) self.assertIsNone(transformed.argument)
def test_generated_tensorflow_executes_correctly_tuple_parameter(self): param = building_blocks.Reference('x', [tf.int32, tf.float32]) body = building_blocks.Tuple([ building_blocks.Selection(param, index=1), building_blocks.Selection(param, index=0) ]) int_constant = building_block_factory.create_tensorflow_constant( [tf.int32, tf.float32], 1) tf_block = compiler_transformations.construct_tensorflow_calling_lambda_on_concrete_arg( param, body, int_constant) result = test_utils.run_tensorflow(tf_block.function.proto) self.assertLen(result, 2) self.assertEqual(result[0], 1.) self.assertEqual(result[1], 1)
def test_generated_tensorflow_executes_correctly_int_parameter(self): param = building_blocks.Reference('x', tf.int32) body = building_blocks.Tuple([ building_blocks.Reference('x', tf.int32), building_blocks.Reference('x', tf.int32) ]) int_constant = building_block_factory.create_tensorflow_constant( tf.int32, 0) tf_block = transformations.construct_tensorflow_calling_lambda_on_concrete_arg( param, body, int_constant) result = test_utils.run_tensorflow(tf_block.function.proto) self.assertLen(result, 2) self.assertEqual(result[0], 0) self.assertEqual(result[1], 0)
def _construct_block_and_inlined_tuple(k): concrete_int = building_block_factory.create_tensorflow_constant( tf.int32, 1) first_tf_id = building_block_factory.create_compiled_identity(tf.int32) called_tf_id = building_blocks.Call(first_tf_id, concrete_int) for _ in range(k): # Simulating large TF computation called_tf_id = building_blocks.Call(first_tf_id, called_tf_id) ref_to_call = building_blocks.Reference('call', called_tf_id.type_signature) block_locals = [('call', called_tf_id)] block = building_blocks.Block( block_locals, building_blocks.Tuple([ref_to_call, ref_to_call])) inlined_tuple = building_blocks.Tuple([called_tf_id, called_tf_id]) return block, inlined_tuple
async def embed_tf_scalar_constant(executor, type_spec, val): """Embeds a constant `val` of TFF type `type_spec` in `executor`. Args: executor: An instance of `tff.framework.Executor`. type_spec: An instance of `tff.Type`. val: A scalar value. Returns: An instance of `tff.framework.ExecutorValue` containing an embedded value. """ py_typecheck.check_type(executor, executor_base.Executor) fn_building_block = (building_block_factory.create_tensorflow_constant( type_spec, val)) embedded_val = await executor.create_call(await executor.create_value( fn_building_block.function.proto, fn_building_block.function.type_signature)) type_utils.check_equivalent_types(embedded_val.type_signature, type_spec) return embedded_val
async def _embed_tf_scalar_constant(executor, type_spec, val): """Embeds a constant `val` of TFF type `type_spec` in `executor`. Args: executor: An instance of `tff.framework.Executor`. type_spec: An instance of `tff.Type`. val: A scalar value. Returns: An instance of `tff.framework.ExecutorValue` containing an embedded value. """ # TODO(b/134543154): Perhaps graduate this and the function below it into a # separate library, so that it can be used in other places. py_typecheck.check_type(executor, executor_base.Executor) fn_building_block = (building_block_factory.create_tensorflow_constant( type_spec, val)) embedded_val = await executor.create_call(await executor.create_value( fn_building_block.function.proto, fn_building_block.function.type_signature)) type_utils.check_equivalent_types(embedded_val.type_signature, type_spec) return embedded_val