def test_replace_intrinsic_plus_reduce_lambdas(self):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.SERVER, True))
def foo(x):
return intrinsics.federated_sum(intrinsics.federated_broadcast(x))
comp = _to_building_block(foo)
self.assertEqual(
str(comp),
'(foo_arg -> federated_sum(federated_broadcast(foo_arg)))')
bodies = intrinsic_bodies.get_intrinsic_bodies(
context_stack_impl.context_stack)
transformed_comp = transformations.replace_intrinsic(
comp, intrinsic_defs.FEDERATED_SUM.uri, bodies['federated_sum'],
context_stack_impl.context_stack)
self.assertEqual(
str(transformed_comp),
'(foo_arg -> (federated_sum_arg -> federated_reduce('
'<federated_sum_arg,generic_zero,generic_plus>))'
'(federated_broadcast(foo_arg)))')
reduced_lambda_comp = transformations.replace_called_lambdas_with_block(
transformed_comp)
self.assertEqual(
str(reduced_lambda_comp),
'(foo_arg -> (let federated_sum_arg=federated_broadcast(foo_arg) in '
'federated_reduce(<federated_sum_arg,generic_zero,generic_plus>)))'
)
def test_simple_reduce_lambda(self):
x = computation_building_blocks.Reference('x', [tf.int32])
l = computation_building_blocks.Lambda('x', [tf.int32], x)
input_val = computation_building_blocks.Tuple(
[computation_building_blocks.Data('test', tf.int32)])
called = computation_building_blocks.Call(l, input_val)
self.assertEqual(str(called), '(x -> x)(<test>)')
reduced = transformations.replace_called_lambdas_with_block(called)
self.assertEqual(str(reduced), '(let x=<test> in x)')
def test_no_reduce_lambda_without_call(self):
@computations.federated_computation(tf.int32)
def foo(x):
return x
comp = _to_building_block(foo)
py_typecheck.check_type(comp, computation_building_blocks.Lambda)
lambda_reduced_comp = transformations.replace_called_lambdas_with_block(
comp)
self.assertEqual(str(comp), '(foo_arg -> foo_arg)')
self.assertEqual(str(comp), str(lambda_reduced_comp))
def test_no_reduce_separated_lambda_and_call(self):
@computations.federated_computation(tf.int32)
def foo(x):
return x
comp = _to_building_block(foo)
block_wrapped_comp = computation_building_blocks.Block([], comp)
test_arg = computation_building_blocks.Data('test', tf.int32)
called_block = computation_building_blocks.Call(
block_wrapped_comp, test_arg)
lambda_reduced_comp = transformations.replace_called_lambdas_with_block(
called_block)
self.assertEqual(str(called_block),
'(let in (foo_arg -> foo_arg))(test)')
self.assertEqual(str(called_block), str(lambda_reduced_comp))
def test_nested_reduce_lambda(self):
comp = computation_building_blocks.Tuple(
[computation_building_blocks.Data('test', tf.int32)])
input1 = computation_building_blocks.Reference('input1',
comp.type_signature)
first_level_call = computation_building_blocks.Call(
computation_building_blocks.Lambda('input1', input1.type_signature,
input1), comp)
input2 = computation_building_blocks.Reference(
'input2', first_level_call.type_signature)
second_level_call = computation_building_blocks.Call(
computation_building_blocks.Lambda('input2', input2.type_signature,
input2), first_level_call)
lambda_reduced_comp = transformations.replace_called_lambdas_with_block(
second_level_call)
self.assertEqual(str(second_level_call),
'(input2 -> input2)((input1 -> input1)(<test>))')
self.assertEqual(
str(lambda_reduced_comp),
'(let input2=(let input1=<test> in input1) in input2)')