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)