async def _compute_intrinsic_federated_weighted_mean(self, arg):
type_utils.check_valid_federated_weighted_mean_argument_tuple_type(
arg.type_signature)
zipped_arg = await self._compute_intrinsic_federated_zip_at_clients(arg)
# TODO(b/134543154): Replace with something that produces a section of
# plain TensorFlow code instead of constructing a lambda (so that this
# can be executed directly on top of a plain TensorFlow-based executor).
multiply_blk = intrinsic_utils.construct_binary_operator_with_upcast(
zipped_arg.type_signature.member, tf.multiply)
sum_of_products = await self._compute_intrinsic_federated_sum(
await self._compute_intrinsic_federated_map(
FederatedExecutorValue(
anonymous_tuple.AnonymousTuple([
(None, multiply_blk.proto),
(None, zipped_arg.internal_representation)
]),
computation_types.NamedTupleType(
[multiply_blk.type_signature, zipped_arg.type_signature]))))
total_weight = await self._compute_intrinsic_federated_sum(
FederatedExecutorValue(arg.internal_representation[1],
arg.type_signature[1]))
divide_arg = await self._compute_intrinsic_federated_zip_at_server(
await self.create_tuple(
anonymous_tuple.AnonymousTuple([(None, sum_of_products),
(None, total_weight)])))
divide_blk = intrinsic_utils.construct_binary_operator_with_upcast(
divide_arg.type_signature.member, tf.divide)
return await self._compute_intrinsic_federated_apply(
FederatedExecutorValue(
anonymous_tuple.AnonymousTuple([
(None, divide_blk.proto),
(None, divide_arg.internal_representation)
]),
computation_types.NamedTupleType(
[divide_blk.type_signature, divide_arg.type_signature])))
def check_valid_federated_weighted_mean_argument_tuple_type(self):
type_utils.check_valid_federated_weighted_mean_argument_tuple_type(
computation_types.to_type([type_factory.at_clients(tf.float32)] *
2))
with self.assertRaises(TypeError):
type_utils.check_valid_federated_weighted_mean_argument_tuple_type(
computation_types.to_type([type_factory.at_clients(tf.int32)] *
2))
def _federated_weighted_mean(self, arg):
type_utils.check_valid_federated_weighted_mean_argument_tuple_type(
arg.type_signature)
v_type = arg.type_signature[0].member
total = sum(arg.value[1])
products_val = [
multiply_by_scalar(ComputedValue(v, v_type), w / total).value
for v, w in zip(arg.value[0], arg.value[1])
]
return self._federated_sum(
ComputedValue(products_val, type_factory.at_clients(v_type)))
async def compute_federated_weighted_mean(executor, arg):
"""Computes a federated weighted using simpler intrinsic coroutines.
Args:
executor: The executor to use.
arg: The argument tuple value, which must be embedded in `executor`.
Returns:
The result embedded in `executor`.
"""
type_utils.check_valid_federated_weighted_mean_argument_tuple_type(
arg.type_signature)
zip1_type = computation_types.FunctionType(
computation_types.NamedTupleType([
type_factory.at_clients(arg.type_signature[0].member),
type_factory.at_clients(arg.type_signature[1].member)
]),
type_factory.at_clients(
computation_types.NamedTupleType(
[arg.type_signature[0].member, arg.type_signature[1].member])))
zip1_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_ZIP_AT_CLIENTS,
zip1_type)
zipped_arg = await executor.create_call(
await executor.create_value(zip1_comp, zip1_type), arg)
# TODO(b/134543154): Replace with something that produces a section of
# plain TensorFlow code instead of constructing a lambda (so that this
# can be executed directly on top of a plain TensorFlow-based executor).
multiply_blk = building_block_factory.create_binary_operator_with_upcast(
zipped_arg.type_signature.member, tf.multiply)
map_type = computation_types.FunctionType(
computation_types.NamedTupleType(
[multiply_blk.type_signature, zipped_arg.type_signature]),
type_factory.at_clients(multiply_blk.type_signature.result))
map_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_MAP, map_type)
products = await executor.create_call(
await executor.create_value(map_comp, map_type), await
executor.create_tuple([
await executor.create_value(multiply_blk.proto,
multiply_blk.type_signature),
zipped_arg
]))
sum1_type = computation_types.FunctionType(
type_factory.at_clients(products.type_signature.member),
type_factory.at_server(products.type_signature.member))
sum1_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM, sum1_type)
sum_of_products = await executor.create_call(
await executor.create_value(sum1_comp, sum1_type), products)
sum2_type = computation_types.FunctionType(
type_factory.at_clients(arg.type_signature[1].member),
type_factory.at_server(arg.type_signature[1].member))
sum2_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_SUM, sum2_type)
total_weight = await executor.create_call(
*(await asyncio.gather(executor.create_value(sum2_comp, sum2_type),
executor.create_selection(arg, index=1))))
zip2_type = computation_types.FunctionType(
computation_types.NamedTupleType(
[sum_of_products.type_signature, total_weight.type_signature]),
type_factory.at_server(
computation_types.NamedTupleType([
sum_of_products.type_signature.member,
total_weight.type_signature.member
])))
zip2_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_ZIP_AT_SERVER,
zip2_type)
divide_arg = await executor.create_call(*(await asyncio.gather(
executor.create_value(zip2_comp, zip2_type),
executor.create_tuple([sum_of_products, total_weight]))))
divide_blk = building_block_factory.create_binary_operator_with_upcast(
divide_arg.type_signature.member, tf.divide)
apply_type = computation_types.FunctionType(
computation_types.NamedTupleType(
[divide_blk.type_signature, divide_arg.type_signature]),
type_factory.at_server(divide_blk.type_signature.result))
apply_comp = create_intrinsic_comp(intrinsic_defs.FEDERATED_APPLY,
apply_type)
return await executor.create_call(*(await asyncio.gather(
executor.create_value(apply_comp, apply_type),
executor.create_tuple([
await executor.create_value(divide_blk.proto,
divide_blk.type_signature), divide_arg
]))))