def _federated_computation_wrapper_fn(target_fn, parameter_type, name=None):
"""Wrapper function to plug orchestration logic in to TFF framework."""
ctx_stack = context_stack_impl.context_stack
target_lambda = (
federated_computation_utils.zero_or_one_arg_fn_to_building_block(
target_fn,
'arg' if parameter_type else None,
parameter_type,
ctx_stack,
suggested_name=name))
return computation_impl.ComputationImpl(target_lambda.proto, ctx_stack)
def test_py_container_args(self, fn, parameter_type, exepcted_result_type):
parameter_name = 'foo'
fn = function_utils.wrap_as_zero_or_one_arg_callable(fn, parameter_type)
_, type_signature = federated_computation_utils.zero_or_one_arg_fn_to_building_block(
fn, parameter_name, parameter_type, context_stack_impl.context_stack)
self.assertIs(type(type_signature.result), type(exepcted_result_type))
self.assertIs(
computation_types.NamedTupleTypeWithPyContainerType.get_container_type(
type_signature.result),
computation_types.NamedTupleTypeWithPyContainerType.get_container_type(
exepcted_result_type))
self.assertEqual(type_signature.result, exepcted_result_type)
def _transform(comp):
"""Internal transform function."""
if not _should_transform(comp):
return comp
# We need 'wrapped_body' to accept exactly one argument.
wrapped_body = lambda x: body(x) # pylint: disable=unnecessary-lambda
return federated_computation_utils.zero_or_one_arg_fn_to_building_block(
wrapped_body,
'arg',
comp.type_signature.parameter,
context_stack,
suggested_name=uri)
def test_py_container_args(self, fn, parameter_type, result_type):
parameter_name = 'foo'
parameter_type = computation_types.to_type(parameter_type)
fn = function_utils.wrap_as_zero_or_one_arg_callable(
fn, parameter_type)
_, annotated_type = federated_computation_utils.zero_or_one_arg_fn_to_building_block(
fn, parameter_name, parameter_type,
context_stack_impl.context_stack)
self.assertIs(type(annotated_type.result), type(result_type))
self.assertIs(
NamedTupleTypeWithPyContainerType.get_container_type(
annotated_type.result),
NamedTupleTypeWithPyContainerType.get_container_type(result_type))
self.assertEqual(annotated_type.result, result_type)
def _transformation_fn(comp, uri, body):
"""Internal function to replace occurrences of an intrinsic."""
if not isinstance(comp, computation_building_blocks.Intrinsic):
return comp
elif comp.uri != uri:
return comp
else:
py_typecheck.check_type(comp.type_signature,
computation_types.FunctionType)
# We need 'wrapped_body' to accept exactly one argument.
wrapped_body = lambda x: body(x) # pylint: disable=unnecessary-lambda
return federated_computation_utils.zero_or_one_arg_fn_to_building_block(
wrapped_body,
'arg',
comp.type_signature.parameter,
context_stack,
suggested_name=uri)
def _federated_computation_wrapper_fn(target_fn,
parameter_type,
unpack,
name=None):
"""Wrapper function to plug orchestration logic into the TFF framework.
This function is passed through `computation_wrapper.ComputationWrapper`.
Documentation its arguments can be found inside the definition of that class.
"""
target_fn = function_utils.wrap_as_zero_or_one_arg_callable(
target_fn, parameter_type, unpack)
ctx_stack = context_stack_impl.context_stack
target_lambda = (
federated_computation_utils.zero_or_one_arg_fn_to_building_block(
target_fn,
'arg' if parameter_type else None,
parameter_type,
ctx_stack,
suggested_name=name))
return computation_impl.ComputationImpl(target_lambda.proto, ctx_stack)
