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
