'and `TensorType`; you have attempted to create one '
'with the type {}.'.format(parameter_type))
ctx_stack = context_stack_impl.context_stack
tf_serializer = tensorflow_serialization.tf_computation_serializer(
parameter_type, ctx_stack)
arg = next(tf_serializer)
try:
result = yield arg
except Exception as e: # pylint: disable=broad-except
tf_serializer.throw(e)
comp_pb, extra_type_spec = tf_serializer.send(result)
yield computation_impl.ComputationImpl(comp_pb, ctx_stack, extra_type_spec)
tensorflow_wrapper = computation_wrapper.ComputationWrapper(
computation_wrapper.PythonTracingStrategy(_tf_wrapper_fn))
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,
def invoke(self, zero_traced_fn, arg):
return Result(arg=arg,
arg_type=zero_traced_fn.type_signature.parameter,
zero_result=zero_traced_fn.zero_result)
@attr.s
class Result:
arg = attr.ib()
arg_type = attr.ib()
zero_result = attr.ib()
test_wrap = computation_wrapper.ComputationWrapper(
computation_wrapper.PythonTracingStrategy(_zero_tracer))
class ComputationWrapperTest(test_case.TestCase):
# Note: Many tests below silence certain linter warnings. These warnings are
# not applicable, since it's the wrapper code, not not the whimsy functions
# that are being tested, so whether the specific function declarations used
# here follow good practices is not really relevant. The purpose of the test
# is to exercise various corner cases that the wrapper needs to be able to
# correctly handle.
def test_as_decorator_with_kwargs(self):
with self.assertRaises(TypeError):
@test_wrap(foo=1)