def _federated_computation_wrapper_fn(parameter_type, name):
"""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.
"""
ctx_stack = context_stack_impl.context_stack
if parameter_type is None:
parameter_name = None
else:
parameter_name = 'arg'
fn_generator = federated_computation_utils.federated_computation_serializer(
parameter_name=parameter_name,
parameter_type=parameter_type,
context_stack=ctx_stack,
suggested_name=name)
arg = next(fn_generator)
try:
result = yield arg
except Exception as e: # pylint: disable=broad-except
fn_generator.throw(e)
target_lambda, extra_type_spec = fn_generator.send(result)
fn_generator.close()
yield computation_impl.ComputationImpl(target_lambda.proto, ctx_stack,
extra_type_spec)
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.
"""
unpack_arguments = function_utils.create_argument_unpacking_fn(
target_fn, parameter_type, unpack)
ctx_stack = context_stack_impl.context_stack
fn_generator = federated_computation_utils.federated_computation_serializer(
'arg' if parameter_type else None,
parameter_type,
ctx_stack,
suggested_name=name)
args, kwargs = unpack_arguments(next(fn_generator))
result = target_fn(*args, **kwargs)
if result is None:
line_number = target_fn.__code__.co_firstlineno
filename = target_fn.__code__.co_filename
raise ValueError(
f'The function defined on line {line_number} of file {filename} '
'returned `None`, but `federated_computation`s must return some '
'non-`None` value.')
target_lambda, extra_type_spec = fn_generator.send(result)
return computation_impl.ComputationImpl(target_lambda.proto, ctx_stack,
extra_type_spec)
def _federated_computation_serializer(fn, parameter_name, parameter_type):
unpack_arguments = function_utils.create_argument_unpacking_fn(
fn, parameter_type)
fn_gen = federated_computation_utils.federated_computation_serializer(
parameter_name, parameter_type, context_stack_impl.context_stack)
args, kwargs = unpack_arguments(next(fn_gen))
result = fn(*args, **kwargs)
return fn_gen.send(result)
def _transform(comp):
if not _should_transform(comp):
return comp, False
fn_generator = federated_computation_utils.federated_computation_serializer(
'arg', comp.type_signature.parameter, context_stack, uri)
result = body(next(fn_generator))
transformed_comp, _ = fn_generator.send(result)
return transformed_comp, True
def _federated_computation_wrapper_fn(parameter_type, name):
"""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.
"""
ctx_stack = context_stack_impl.context_stack
fn_generator = federated_computation_utils.federated_computation_serializer(
'arg' if parameter_type else None,
parameter_type,
ctx_stack,
suggested_name=name)
result = yield next(fn_generator)
target_lambda, extra_type_spec = fn_generator.send(result)
yield computation_impl.ComputationImpl(target_lambda.proto, ctx_stack,
extra_type_spec)
