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 test_wrap_as_zero_or_one_arg_callable(self, unused_index, fn,
parameter_type, unpack, arg,
expected_result):
parameter_type = computation_types.to_type(parameter_type)
unpack_arguments = function_utils.create_argument_unpacking_fn(
fn, parameter_type, unpack)
args, kwargs = unpack_arguments(arg)
actual_result = fn(*args, **kwargs)
self.assertEqual(actual_result, expected_result)
def _polymorphic_wrapper(parameter_type: computation_types.Type,
unpack: Optional[bool]):
unpack_arguments_fn = function_utils.create_argument_unpacking_fn(
fn_to_wrap, parameter_type, unpack=unpack)
wrapped_fn_generator = _wrap_concrete(fn_name,
self._wrapper_fn,
parameter_type)
args, kwargs = unpack_arguments_fn(next(wrapped_fn_generator))
result = fn_to_wrap(*args, **kwargs)
if result is None:
raise ComputationReturnedNoneError(fn_to_wrap)
return wrapped_fn_generator.send(result)
def __init__(self, fn, parameter_type, unpack, name=None):
del name
unpack_args = function_utils.create_argument_unpacking_fn(
fn, parameter_type, unpack)
def _wrapped_fn(arg):
args, kwargs = unpack_args(arg)
return fn(*args, **kwargs)
self._fn = _wrapped_fn
super().__init__(
computation_types.FunctionType(parameter_type, tf.string),
context_stack_impl.context_stack)
def test_nested_structure_type_signature_roundtrip(self):
def traced_fn(x):
return x[0][0]
param_type = computation_types.to_type([(np.int32,)])
arg_fn = function_utils.create_argument_unpacking_fn(traced_fn, param_type)
ctx_stack = context_stack_impl.context_stack
comp_pb = jax_serialization.serialize_jax_computation(
traced_fn, arg_fn, param_type, ctx_stack)
self.assertIsInstance(comp_pb, pb.Computation)
self.assertEqual(comp_pb.WhichOneof('computation'), 'xla')
type_spec = type_serialization.deserialize_type(comp_pb.type)
self.assertEqual(str(type_spec), '(<<int32>> -> int32)')
def test_arg_ordering(self):
param_type = computation_types.to_type(
(computation_types.TensorType(np.int32, 10),
computation_types.TensorType(np.int32)))
def traced_fn(b, a):
return jax.numpy.add(a, jax.numpy.sum(b))
arg_fn = function_utils.create_argument_unpacking_fn(traced_fn, param_type)
ctx_stack = context_stack_impl.context_stack
comp_pb = jax_serialization.serialize_jax_computation(
traced_fn, arg_fn, param_type, ctx_stack)
self.assertIsInstance(comp_pb, pb.Computation)
self.assertEqual(comp_pb.WhichOneof('computation'), 'xla')
type_spec = type_serialization.deserialize_type(comp_pb.type)
self.assertEqual(str(type_spec), '(<int32[10],int32> -> int32)')
def test_serialize_jax_with_int32_to_int32(self):
def traced_fn(x):
return x + 10
param_type = computation_types.to_type(np.int32)
arg_fn = function_utils.create_argument_unpacking_fn(traced_fn, param_type)
ctx_stack = context_stack_impl.context_stack
comp_pb = jax_serialization.serialize_jax_computation(
traced_fn, arg_fn, param_type, ctx_stack)
self.assertIsInstance(comp_pb, pb.Computation)
self.assertEqual(comp_pb.WhichOneof('computation'), 'xla')
type_spec = type_serialization.deserialize_type(comp_pb.type)
self.assertEqual(str(type_spec), '(int32 -> int32)')
xla_comp = xla_serialization.unpack_xla_computation(comp_pb.xla.hlo_module)
self.assertIn('ROOT tuple.6 = (s32[]) tuple(add.5)', xla_comp.as_hlo_text())
self.assertEqual(str(comp_pb.xla.result), str(comp_pb.xla.parameter))
self.assertEqual(str(comp_pb.xla.result), 'tensor {\n' ' index: 0\n' '}\n')
def test_serialize_jax_with_2xint32_to_2xint32(self):
def traced_fn(x):
return collections.OrderedDict([('sum', x['foo'] + x['bar']),
('difference', x['bar'] - x['foo'])])
param_type = computation_types.to_type(
collections.OrderedDict([('foo', np.int32), ('bar', np.int32)]))
arg_fn = function_utils.create_argument_unpacking_fn(traced_fn, param_type)
ctx_stack = context_stack_impl.context_stack
comp_pb = jax_serialization.serialize_jax_computation(
traced_fn, arg_fn, param_type, ctx_stack)
self.assertIsInstance(comp_pb, pb.Computation)
self.assertEqual(comp_pb.WhichOneof('computation'), 'xla')
type_spec = type_serialization.deserialize_type(comp_pb.type)
self.assertEqual(
str(type_spec),
'(<foo=int32,bar=int32> -> <sum=int32,difference=int32>)')
xla_comp = xla_serialization.unpack_xla_computation(comp_pb.xla.hlo_module)
self.assertIn(
# pylint: disable=line-too-long
' constant.4 = pred[] constant(false)\n'
' parameter.1 = (s32[], s32[]) parameter(0)\n'
' get-tuple-element.2 = s32[] get-tuple-element(parameter.1), index=0\n'
' get-tuple-element.3 = s32[] get-tuple-element(parameter.1), index=1\n'
' add.5 = s32[] add(get-tuple-element.2, get-tuple-element.3)\n'
' subtract.6 = s32[] subtract(get-tuple-element.3, get-tuple-element.2)\n'
' ROOT tuple.7 = (s32[], s32[]) tuple(add.5, subtract.6)\n',
xla_comp.as_hlo_text())
self.assertEqual(str(comp_pb.xla.result), str(comp_pb.xla.parameter))
self.assertEqual(
str(comp_pb.xla.parameter), 'struct {\n'
' element {\n'
' tensor {\n'
' index: 0\n'
' }\n'
' }\n'
' element {\n'
' tensor {\n'
' index: 1\n'
' }\n'
' }\n'
'}\n')
def _tf_wrapper_fn(target_fn, parameter_type, unpack, name=None):
"""Wrapper function to plug Tensorflow 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.
"""
del name # Unused.
unpack_arguments = function_utils.create_argument_unpacking_fn(
target_fn, parameter_type, unpack)
if not type_analysis.is_tensorflow_compatible_type(parameter_type):
raise TypeError('`tf_computation`s can accept only parameter types with '
'constituents `SequenceType`, `StructType` '
'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)
args, kwargs = unpack_arguments(next(tf_serializer))
result = target_fn(*args, **kwargs)
comp_pb, extra_type_spec = tf_serializer.send(result)
return computation_impl.ComputationImpl(comp_pb, ctx_stack, extra_type_spec)
def _jax_strategy_fn(fn_to_wrap, fn_name, parameter_type, unpack):
"""Serializes a Python function containing JAX code as a TFF computation.
Args:
fn_to_wrap: The Python function containing JAX code to be serialized as a
computation containing XLA.
fn_name: The name for the constructed computation (currently ignored).
parameter_type: An instance of `computation_types.Type` that represents the
TFF type of the computation parameter, or `None` if there's none.
unpack: See `unpack` in `function_utils.create_argument_unpacking_fn`.
Returns:
An instance of `computation_impl.ComputationImpl` with the constructed
computation.
"""
del fn_name # Unused.
unpack_arguments_fn = function_utils.create_argument_unpacking_fn(
fn_to_wrap, parameter_type, unpack=unpack)
ctx_stack = context_stack_impl.context_stack
comp_pb = jax_serialization.serialize_jax_computation(
fn_to_wrap, unpack_arguments_fn, parameter_type, ctx_stack)
return computation_impl.ComputationImpl(comp_pb, ctx_stack)
def test_serialize_jax_with_two_args(self):
def traced_fn(x, y):
return x + y
param_type = computation_types.to_type(
collections.OrderedDict([('x', np.int32), ('y', np.int32)]))
arg_fn = function_utils.create_argument_unpacking_fn(traced_fn, param_type)
ctx_stack = context_stack_impl.context_stack
comp_pb = jax_serialization.serialize_jax_computation(
traced_fn, arg_fn, param_type, ctx_stack)
self.assertIsInstance(comp_pb, pb.Computation)
self.assertEqual(comp_pb.WhichOneof('computation'), 'xla')
type_spec = type_serialization.deserialize_type(comp_pb.type)
self.assertEqual(str(type_spec), '(<x=int32,y=int32> -> int32)')
xla_comp = xla_serialization.unpack_xla_computation(comp_pb.xla.hlo_module)
self.assertIn(
# pylint: disable=line-too-long
' constant.4 = pred[] constant(false)\n'
' parameter.1 = (s32[], s32[]) parameter(0)\n'
' get-tuple-element.2 = s32[] get-tuple-element(parameter.1), index=0\n'
' get-tuple-element.3 = s32[] get-tuple-element(parameter.1), index=1\n'
' add.5 = s32[] add(get-tuple-element.2, get-tuple-element.3)\n'
' ROOT tuple.6 = (s32[]) tuple(add.5)\n',
xla_comp.as_hlo_text())
self.assertEqual(
str(comp_pb.xla.parameter), 'struct {\n'
' element {\n'
' tensor {\n'
' index: 0\n'
' }\n'
' }\n'
' element {\n'
' tensor {\n'
' index: 1\n'
' }\n'
' }\n'
'}\n')
self.assertEqual(str(comp_pb.xla.result), 'tensor {\n' ' index: 0\n' '}\n')
def __call__(self, fn_to_wrap, fn_name, parameter_type, unpack):
unpack_arguments_fn = function_utils.create_argument_unpacking_fn(
fn_to_wrap, parameter_type, unpack=unpack)
wrapped_fn_generator = _wrap_concrete(fn_name, self._wrapper_fn,
parameter_type)
packed_args = next(wrapped_fn_generator)
try:
args, kwargs = unpack_arguments_fn(packed_args)
result = fn_to_wrap(*args, **kwargs)
if result is None:
raise ComputationReturnedNoneError(fn_to_wrap)
except Exception:
# Give nested generators an opportunity to clean up, then
# re-raise the original error without extra context.
# We don't want to simply pass the error into the generators,
# as that would result in the whole generator stack being added
# to the error message.
try:
wrapped_fn_generator.throw(_TracingError())
except _TracingError:
pass
raise
return wrapped_fn_generator.send(result)
def __call__(self, *args, tff_internal_types=None):
"""Handles the different modes of usage of the decorator/wrapper.
Args:
*args: Positional arguments (the decorator at this point does not accept
keyword arguments, although that might change in the future).
tff_internal_types: TFF internal usage only. This argument should be
considered private.
Returns:
Either a result of wrapping, or a callable that expects a function,
method, or a tf.function and performs wrapping on it, depending on
specific usage pattern.
Raises:
TypeError: if the arguments are of the wrong types.
ValueError: if the function to wrap returns `None`.
"""
if not args or not is_function(args[0]):
# If invoked as a decorator, and with an empty argument list as "@xyz()"
# applied to a function definition, expect the Python function being
# decorated to be passed in the subsequent call, and potentially create
# a polymorphic callable. The parameter type is unspecified.
# Deliberate wrapping with a lambda to prevent the caller from being able
# to accidentally specify parameter type as a second argument.
# The tricky partial recursion is needed to inline the logic in the
# "success" case below.
if tff_internal_types is not None:
raise TypeError(f'Expected a function to wrap, found {args}.')
provided_types = tuple(map(computation_types.to_type, args))
return functools.partial(self.__call__,
tff_internal_types=provided_types)
# If the first argument on the list is a Python function, instance method,
# or a tf.function, this is the one that's being wrapped. This is the case
# of either a decorator invocation without arguments as "@xyz" applied to
# a function definition, of an inline invocation as
# `... = xyz(lambda....).`
# Any of the following arguments, if present, are the arguments to the
# wrapper that are to be interpreted as the type specification.
fn_to_wrap = args[0]
if not tff_internal_types:
tff_internal_types = tuple(map(computation_types.to_type,
args[1:]))
else:
if len(args) > 1:
raise TypeError(
f'Expected no further arguments, found {args[1:]}.')
parameter_types = tff_internal_types
parameters = _parameters(fn_to_wrap)
# NOTE: many of the properties checked here are only necessary for
# non-polymorphic computations whose type signatures must be resolved
# prior to use. However, we continue to enforce these requirements even
# in the polymorphic case in order to avoid creating an inconsistency.
_check_parameters(parameters)
try:
fn_name = fn_to_wrap.__name__
except AttributeError:
fn_name = None
if (not parameter_types) and parameters:
# There is no TFF type specification, and the function/tf.function
# declares parameters. Create a polymorphic template.
def _polymorphic_wrapper(parameter_type: computation_types.Type,
unpack: Optional[bool]):
unpack_arguments_fn = function_utils.create_argument_unpacking_fn(
fn_to_wrap, parameter_type, unpack=unpack)
wrapped_fn_generator = _wrap_concrete(fn_name,
self._wrapper_fn,
parameter_type)
args, kwargs = unpack_arguments_fn(next(wrapped_fn_generator))
result = fn_to_wrap(*args, **kwargs)
if result is None:
raise ComputationReturnedNoneError(fn_to_wrap)
return wrapped_fn_generator.send(result)
wrapped_func = function_utils.PolymorphicFunction(
_polymorphic_wrapper)
else:
# Either we have a concrete parameter type, or this is no-arg function.
parameter_type = _parameter_type(parameters, parameter_types)
unpack_arguments_fn = function_utils.create_argument_unpacking_fn(
fn_to_wrap, parameter_type, unpack=None)
wrapped_fn_generator = _wrap_concrete(fn_name, self._wrapper_fn,
parameter_type)
args, kwargs = unpack_arguments_fn(next(wrapped_fn_generator))
result = fn_to_wrap(*args, **kwargs)
if result is None:
raise ComputationReturnedNoneError(fn_to_wrap)
wrapped_func = wrapped_fn_generator.send(result)
# Copy the __doc__ attribute with the documentation in triple-quotes from
# the decorated function.
wrapped_func.__doc__ = getattr(fn_to_wrap, '__doc__', None)
return wrapped_func