async def _compute_intrinsic_federated_sum(self, arg):
py_typecheck.check_type(arg.type_signature,
computation_types.FederatedType)
zero_building_block = (
computation_constructing_utils.construct_tensorflow_constant(
arg.type_signature.member, 0))
zero = await self.create_call(await self.create_value(
zero_building_block.function.proto,
zero_building_block.function.type_signature))
type_utils.check_equivalent_types(arg.type_signature.member,
zero.type_signature)
# TODO(b/134543154): There is an opportunity here to import something more
# in line with the usage (no building block wrapping, etc.)
plus_building_block = (computation_constructing_utils.
construct_tensorflow_binary_operator(
zero.type_signature, tf.add))
plus = await self.create_value(plus_building_block.proto,
plus_building_block.type_signature)
type_utils.check_equivalent_types(
plus.type_signature,
type_constructors.binary_op(zero.type_signature))
return await self._compute_intrinsic_federated_reduce(
FederatedExecutorValue(
anonymous_tuple.AnonymousTuple([
(None, arg.internal_representation),
(None, zero.internal_representation),
(None, plus.internal_representation)
]),
computation_types.NamedTupleType([
arg.type_signature, zero.type_signature,
plus.type_signature
])))
def test_construct_integer_type_signature(self):
ref = computation_building_blocks.Reference('x', [tf.int32, tf.int32])
multiplier = intrinsic_utils.construct_binary_operator_with_upcast(
ref.type_signature, tf.multiply)
self.assertEqual(
multiplier.type_signature,
type_constructors.binary_op(computation_types.to_type(tf.int32)))
async def _compute_intrinsic_federated_aggregate(self, arg):
py_typecheck.check_type(arg.type_signature,
computation_types.NamedTupleType)
py_typecheck.check_type(arg.internal_representation,
anonymous_tuple.AnonymousTuple)
if len(arg.internal_representation) != 5:
raise ValueError(
'Expected 5 elements in the `federated_aggregate()` argument tuple, '
'found {}.'.format(len(arg.internal_representation)))
val_type = arg.type_signature[0]
py_typecheck.check_type(val_type, computation_types.FederatedType)
item_type = val_type.member
zero_type = arg.type_signature[1]
accumulate_type = arg.type_signature[2]
type_utils.check_equivalent_types(
accumulate_type,
type_constructors.reduction_op(zero_type, item_type))
merge_type = arg.type_signature[3]
type_utils.check_equivalent_types(
merge_type, type_constructors.binary_op(zero_type))
report_type = arg.type_signature[4]
py_typecheck.check_type(report_type, computation_types.FunctionType)
type_utils.check_equivalent_types(report_type.parameter, zero_type)
# NOTE: This is a simple initial implementation that simply forwards this
# to `federated_reduce()`. The more complete implementation would be able
# to take advantage of the parallelism afforded by `merge` to reduce the
# cost from liner (with respect to the number of clients) to sub-linear.
# TODO(b/134543154): Expand this implementation to take advantage of the
# parallelism afforded by `merge`.
val = arg.internal_representation[0]
zero = arg.internal_representation[1]
accumulate = arg.internal_representation[2]
pre_report = await self._compute_intrinsic_federated_reduce(
FederatedExecutorValue(
anonymous_tuple.AnonymousTuple([(None, val), (None, zero),
(None, accumulate)]),
computation_types.NamedTupleType(
[val_type, zero_type, accumulate_type])))
py_typecheck.check_type(pre_report.type_signature,
computation_types.FederatedType)
type_utils.check_equivalent_types(pre_report.type_signature.member,
report_type.parameter)
report = arg.internal_representation[4]
return await self._compute_intrinsic_federated_apply(
FederatedExecutorValue(
anonymous_tuple.AnonymousTuple([
(None, report), (None, pre_report.internal_representation)
]),
computation_types.NamedTupleType(
[report_type, pre_report.type_signature])))
def plus_for(type_spec, context_stack):
"""Constructs PLUS intrinsic that operates on values of TFF type `type_spec`.
Args:
type_spec: An instance of `types.Type` or something convertible to it.
intrinsic.
context_stack: The context stack to use.
Returns:
The `PLUS` intrinsic of type `<T,T> -> T`, where `T` represents `type_spec`.
"""
type_spec = computation_types.to_type(type_spec)
py_typecheck.check_type(context_stack, context_stack_base.ContextStack)
return value_impl.ValueImpl(
computation_building_blocks.Intrinsic(
intrinsic_defs.GENERIC_PLUS.uri,
type_constructors.binary_op(type_spec)), context_stack)
async def _embed_tf_binary_operator(executor, type_spec, op):
"""Embeds a binary operator `op` on `type_spec`-typed values in `executor`.
Args:
executor: An instance of `tff.framework.Executor`.
type_spec: An instance of `tff.Type` of the type of values that the binary
operator accepts as input and returns as output.
op: An operator function (such as `tf.add` or `tf.multiply`) to apply to the
tensor-level constituents of the values, pointwise.
Returns:
An instance of `tff.framework.ExecutorValue` representing the operator in
a form embedded into the executor.
"""
# TODO(b/134543154): There is an opportunity here to import something more
# in line with the usage (no building block wrapping, etc.)
fn_building_block = (
computation_constructing_utils.create_tensorflow_binary_operator(
type_spec, op))
embedded_val = await executor.create_value(
fn_building_block.proto, fn_building_block.type_signature)
type_utils.check_equivalent_types(embedded_val.type_signature,
type_constructors.binary_op(type_spec))
return embedded_val
# a = generic_partial_reduce(x, zero, accumulate, INTERMEDIATE_AGGREGATORS)
# b = generic_reduce(a, zero, merge, SERVER)
# c = generic_map(report, b)
# return c
#
# Actual implementations might vary.
#
# Type signature: <{T}@CLIENTS,U,(<U,T>->U),(<U,U>->U),(U->R)> -> R@SERVER
FEDERATED_AGGREGATE = IntrinsicDef(
'FEDERATED_AGGREGATE', 'federated_aggregate',
computation_types.FunctionType(parameter=[
type_constructors.at_clients(computation_types.AbstractType('T')),
computation_types.AbstractType('U'),
type_constructors.reduction_op(computation_types.AbstractType('U'),
computation_types.AbstractType('T')),
type_constructors.binary_op(computation_types.AbstractType('U')),
computation_types.FunctionType(computation_types.AbstractType('U'),
computation_types.AbstractType('R'))
],
result=type_constructors.at_server(
computation_types.AbstractType('R'))))
# Applies a given function to a value on the server.
#
# Type signature: <(T->U),T@SERVER> -> U@SERVER
FEDERATED_APPLY = IntrinsicDef(
'FEDERATED_APPLY', 'federated_apply',
computation_types.FunctionType(parameter=[
computation_types.FunctionType(computation_types.AbstractType('T'),
computation_types.AbstractType('U')),
type_constructors.at_server(computation_types.AbstractType('T')),
def test_binary_op(self):
self.assertEqual(str(type_constructors.binary_op(tf.bool)),
'(<bool,bool> -> bool)')
