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 construct_tensorflow_calling_lambda_on_concrete_arg(
parameter: building_blocks.Reference,
body: building_blocks.ComputationBuildingBlock,
concrete_arg: building_blocks.ComputationBuildingBlock):
"""Generates TensorFlow for lambda invocation with given arg, body and param.
That is, generates TensorFlow block encapsulating the logic represented by
invoking a function with parameter `parameter` and body `body`, with argument
`concrete_arg`.
Via the guarantee made in `compiled_computation_transforms.TupleCalledGraphs`,
this function makes the claim that the computations which define
`concrete_arg` will be executed exactly once in the generated TenosorFlow.
Args:
parameter: Instance of `building_blocks.Reference` defining the parameter of
the function to be generated and invoked, as described above. After
calling this transformation, every instance of parameter` in `body` will
represent a reference to `concrete_arg`.
body: `building_blocks.ComputationBuildingBlock` representing the body of
the function for which we are generating TensorFlow.
concrete_arg: `building_blocks.ComputationBuildingBlock` representing the
argument to be passed to the resulting function. `concrete_arg` will then
be referred to by every occurrence of `parameter` in `body`. Therefore
`concrete_arg` must have an equivalent type signature to that of
`parameter`.
Returns:
A called `building_blocks.CompiledComputation`, as specified above.
Raises:
TypeError: If the arguments are of the wrong types, or the type signature
of `concrete_arg` does not match that of `parameter`.
"""
py_typecheck.check_type(parameter, building_blocks.Reference)
py_typecheck.check_type(body, building_blocks.ComputationBuildingBlock)
py_typecheck.check_type(concrete_arg,
building_blocks.ComputationBuildingBlock)
type_utils.check_equivalent_types(parameter.type_signature,
concrete_arg.type_signature)
def _generate_simple_tensorflow(comp):
tf_parser_callable = tree_to_cc_transformations.TFParser()
comp, _ = tree_transformations.insert_called_tf_identity_at_leaves(
comp)
comp, _ = transformation_utils.transform_postorder(
comp, tf_parser_callable)
return comp
encapsulating_lambda = _generate_simple_tensorflow(
building_blocks.Lambda(parameter.name, parameter.type_signature, body))
comp_called = _generate_simple_tensorflow(
building_blocks.Call(encapsulating_lambda, concrete_arg))
return comp_called
def federated_secure_sum(self, value, bitwidth):
"""Implements `federated_secure_sum` as defined in `api/intrinsics.py`."""
value = value_impl.to_value(value, None, self._context_stack)
value = value_utils.ensure_federated_value(value, placements.CLIENTS,
'value to be summed')
type_utils.check_is_structure_of_integers(value.type_signature)
bitwidth = value_impl.to_value(bitwidth, None, self._context_stack)
type_utils.check_equivalent_types(value.type_signature.member,
bitwidth.type_signature)
value = value_impl.ValueImpl.get_comp(value)
bitwidth = value_impl.ValueImpl.get_comp(bitwidth)
comp = building_block_factory.create_federated_secure_sum(value, bitwidth)
return value_impl.ValueImpl(comp, self._context_stack)
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])))
async def embed_tf_scalar_constant(executor, type_spec, val):
"""Embeds a constant `val` of TFF type `type_spec` in `executor`.
Args:
executor: An instance of `tff.framework.Executor`.
type_spec: An instance of `tff.Type`.
val: A scalar value.
Returns:
An instance of `tff.framework.ExecutorValue` containing an embedded value.
"""
py_typecheck.check_type(executor, executor_base.Executor)
fn_building_block = (building_block_factory.create_tensorflow_constant(
type_spec, val))
embedded_val = await executor.create_call(await executor.create_value(
fn_building_block.function.proto,
fn_building_block.function.type_signature))
type_utils.check_equivalent_types(embedded_val.type_signature, type_spec)
return embedded_val
async def _compute_intrinsic_federated_reduce(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) != 3:
raise ValueError(
'Expected 3 elements in the `federated_reduce()` 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]
op_type = arg.type_signature[2]
type_utils.check_equivalent_types(
op_type, type_factory.reduction_op(zero_type, item_type))
val = arg.internal_representation[0]
py_typecheck.check_type(val, list)
child = self._target_executors[placement_literals.SERVER][0]
async def _move(v):
return await child.create_value(await v.compute(), item_type)
items = await asyncio.gather(*[_move(v) for v in val])
zero = await child.create_value(
await (await self.create_selection(arg, index=1)).compute(),
zero_type)
op = await child.create_value(arg.internal_representation[2], op_type)
result = zero
for item in items:
result = await child.create_call(
op, await child.create_tuple(
anonymous_tuple.AnonymousTuple([(None, result),
(None, item)])))
return FederatedExecutorValue([result],
computation_types.FederatedType(
result.type_signature,
placement_literals.SERVER,
all_equal=True))
async def _embed_tf_scalar_constant(executor, type_spec, val):
"""Embeds a constant `val` of TFF type `type_spec` in `executor`.
Args:
executor: An instance of `tff.framework.Executor`.
type_spec: An instance of `tff.Type`.
val: A scalar value.
Returns:
An instance of `tff.framework.ExecutorValue` containing an embedded value.
"""
# TODO(b/134543154): Perhaps graduate this and the function below it into a
# separate library, so that it can be used in other places.
py_typecheck.check_type(executor, executor_base.Executor)
fn_building_block = (building_block_factory.create_tensorflow_constant(
type_spec, val))
embedded_val = await executor.create_call(await executor.create_value(
fn_building_block.function.proto,
fn_building_block.function.type_signature))
type_utils.check_equivalent_types(embedded_val.type_signature, type_spec)
return embedded_val
def parse_federated_aggregate_argument_types(type_spec):
"""Verifies and parses `type_spec` into constituents.
Args:
type_spec: An instance of `computation_types.NamedTupleType`.
Returns:
A tuple of (value_type, zero_type, accumulate_type, merge_type, report_type)
for the 5 type constituents.
"""
py_typecheck.check_type(type_spec, computation_types.NamedTupleType)
py_typecheck.check_len(type_spec, 5)
value_type = type_spec[0]
py_typecheck.check_type(value_type, computation_types.FederatedType)
item_type = value_type.member
zero_type = type_spec[1]
accumulate_type = type_spec[2]
type_utils.check_equivalent_types(
accumulate_type, type_factory.reduction_op(zero_type, item_type))
merge_type = type_spec[3]
type_utils.check_equivalent_types(merge_type,
type_factory.binary_op(zero_type))
report_type = type_spec[4]
py_typecheck.check_type(report_type, computation_types.FunctionType)
type_utils.check_equivalent_types(report_type.parameter, zero_type)
return value_type, zero_type, accumulate_type, merge_type, report_type
async def _compute_intrinsic_federated_aggregate(self, arg):
val_type, zero_type, accumulate_type, _, report_type = (
executor_utils.parse_federated_aggregate_argument_types(
arg.type_signature))
py_typecheck.check_type(arg.internal_representation,
anonymous_tuple.AnonymousTuple)
py_typecheck.check_len(arg.internal_representation, 5)
# 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(
FederatingExecutorValue(
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(
FederatingExecutorValue(
anonymous_tuple.AnonymousTuple([
(None, report), (None, pre_report.internal_representation)
]),
computation_types.NamedTupleType(
(report_type, pre_report.type_signature))))
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 = (
building_block_factory.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_factory.binary_op(type_spec))
return embedded_val
def test_check_equivalent_types(self):
type_utils.check_equivalent_types(tf.int32, tf.int32)
with self.assertRaises(TypeError):
type_utils.check_equivalent_types(tf.int32, tf.bool)
