async def compute_federated_aggregate(
self, arg: FederatedComposingStrategyValue
) -> FederatedComposingStrategyValue:
value_type, zero_type, accumulate_type, merge_type, report_type = (
executor_utils.parse_federated_aggregate_argument_types(
arg.type_signature))
py_typecheck.check_type(arg.internal_representation, structure.Struct)
py_typecheck.check_len(arg.internal_representation, 5)
val = arg.internal_representation[0]
py_typecheck.check_type(val, list)
py_typecheck.check_len(val, len(self._target_executors))
identity_report, identity_report_type = tensorflow_computation_factory.create_identity(
zero_type)
aggr_type = computation_types.FunctionType(
computation_types.StructType([
value_type, zero_type, accumulate_type, merge_type,
identity_report_type
]), computation_types.at_server(zero_type))
aggr_comp = executor_utils.create_intrinsic_comp(
intrinsic_defs.FEDERATED_AGGREGATE, aggr_type)
zero = await (await
self._executor.create_selection(arg, index=1)).compute()
accumulate = arg.internal_representation[2]
merge = arg.internal_representation[3]
report = arg.internal_representation[4]
async def _child_fn(ex, v):
py_typecheck.check_type(v, executor_value_base.ExecutorValue)
arg_values = [
ex.create_value(zero, zero_type),
ex.create_value(accumulate, accumulate_type),
ex.create_value(merge, merge_type),
ex.create_value(identity_report, identity_report_type)
]
aggr_func, aggr_args = await asyncio.gather(
ex.create_value(aggr_comp, aggr_type),
ex.create_struct([v] +
list(await asyncio.gather(*arg_values))))
child_result = await (await ex.create_call(aggr_func,
aggr_args)).compute()
result_at_server = await self._server_executor.create_value(
child_result, zero_type)
return result_at_server
val_futures = asyncio.as_completed(
[_child_fn(c, v) for c, v in zip(self._target_executors, val)])
parent_merge, parent_report = await asyncio.gather(
self._server_executor.create_value(merge, merge_type),
self._server_executor.create_value(report, report_type))
merge_result = await next(val_futures)
for next_val_future in val_futures:
next_val = await next_val_future
merge_arg = await self._server_executor.create_struct(
[merge_result, next_val])
merge_result = await self._server_executor.create_call(
parent_merge, merge_arg)
report_result = await self._server_executor.create_call(
parent_report, merge_result)
return FederatedComposingStrategyValue(
report_result, computation_types.at_server(report_type.result))
def test_returns_computation(self, type_signature, value):
proto = tensorflow_computation_factory.create_identity(type_signature)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
expected_type = type_factory.unary_op(type_signature)
self.assertEqual(actual_type, expected_type)
actual_result = test_utils.run_tensorflow(proto, value)
self.assertEqual(actual_result, value)
async def compute_federated_aggregate(
self, arg: FederatedComposingStrategyValue
) -> FederatedComposingStrategyValue:
value_type, zero_type, accumulate_type, merge_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)
val = arg.internal_representation[0]
py_typecheck.check_type(val, list)
py_typecheck.check_len(val, len(self._target_executors))
identity_report = tensorflow_computation_factory.create_identity(
zero_type)
identity_report_type = type_factory.unary_op(zero_type)
aggr_type = computation_types.FunctionType(
computation_types.NamedTupleType([
value_type, zero_type, accumulate_type, merge_type,
identity_report_type
]), type_factory.at_server(zero_type))
aggr_comp = executor_utils.create_intrinsic_comp(
intrinsic_defs.FEDERATED_AGGREGATE, aggr_type)
zero = await (await
self._executor.create_selection(arg, index=1)).compute()
accumulate = arg.internal_representation[2]
merge = arg.internal_representation[3]
report = arg.internal_representation[4]
async def _child_fn(ex, v):
py_typecheck.check_type(v, executor_value_base.ExecutorValue)
aggr_func, aggr_args = await asyncio.gather(
ex.create_value(aggr_comp, aggr_type),
ex.create_tuple([v] + list(await asyncio.gather(
ex.create_value(zero, zero_type),
ex.create_value(accumulate, accumulate_type),
ex.create_value(merge, merge_type),
ex.create_value(identity_report, identity_report_type)))))
return await (await ex.create_call(aggr_func, aggr_args)).compute()
vals = await asyncio.gather(
*[_child_fn(c, v) for c, v in zip(self._target_executors, val)])
parent_vals = await asyncio.gather(
*[self._server_executor.create_value(v, zero_type) for v in vals])
parent_merge, parent_report = await asyncio.gather(
self._server_executor.create_value(merge, merge_type),
self._server_executor.create_value(report, report_type))
merge_result = parent_vals[0]
for next_val in parent_vals[1:]:
merge_result = await self._server_executor.create_call(
parent_merge, await
self._server_executor.create_tuple([merge_result, next_val]))
return FederatedComposingStrategyValue(
await self._server_executor.create_call(parent_report,
merge_result),
type_factory.at_server(report_type.result))
def test_returns_computation_sequence(self):
type_signature = computation_types.SequenceType(tf.int32)
proto = tensorflow_computation_factory.create_identity(type_signature)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
expected_type = type_factory.unary_op(type_signature)
self.assertEqual(actual_type, expected_type)
expected_value = [10] * 3
actual_value = test_utils.run_tensorflow(proto, expected_value)
self.assertEqual(actual_value, expected_value)
def test_federated_aggregate_with_unknown_dimension(self):
Accumulator = collections.namedtuple('Accumulator', ['samples']) # pylint: disable=invalid-name
accumulator_type = computation_types.to_type(
Accumulator(samples=computation_types.TensorType(dtype=tf.int32,
shape=[None])))
x = _mock_data_of_type(computation_types.at_clients(tf.int32))
def initialize_fn():
return Accumulator(samples=tf.zeros(shape=[0], dtype=tf.int32))
initialize_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
initialize_fn, None)
initialize = computation_impl.ConcreteComputation(
initialize_proto, context_stack_impl.context_stack)
zero = initialize()
# The operator to use during the first stage simply adds an element to the
# tensor, increasing its size.
def _accumulate(arg):
return Accumulator(samples=tf.concat(
[arg[0].samples,
tf.expand_dims(arg[1], axis=0)], axis=0))
accumulate_type = computation_types.StructType(
[accumulator_type, tf.int32])
accumulate_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_accumulate, accumulate_type)
accumulate = computation_impl.ConcreteComputation(
accumulate_proto, context_stack_impl.context_stack)
# The operator to use during the second stage simply adds total and count.
def _merge(arg):
return Accumulator(
samples=tf.concat([arg[0].samples, arg[1].samples], axis=0))
merge_type = computation_types.StructType(
[accumulator_type, accumulator_type])
merge_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_merge, merge_type)
merge = computation_impl.ConcreteComputation(
merge_proto, context_stack_impl.context_stack)
# The operator to use during the final stage simply computes the ratio.
report_proto, _ = tensorflow_computation_factory.create_identity(
accumulator_type)
report = computation_impl.ConcreteComputation(
report_proto, context_stack_impl.context_stack)
value = intrinsics.federated_aggregate(x, zero, accumulate, merge,
report)
self.assert_value(value, '<samples=int32[?]>@SERVER')
def test_returns_computation_tuple_named(self):
type_signature = [('a', tf.int32), ('b', tf.float32)]
proto = tensorflow_computation_factory.create_identity(type_signature)
self.assertIsInstance(proto, pb.Computation)
actual_type = type_serialization.deserialize_type(proto.type)
expected_type = type_factory.unary_op(type_signature)
self.assertEqual(actual_type, expected_type)
expected_value = anonymous_tuple.AnonymousTuple([('a', 10),
('b', 10.0)])
actual_value = test_utils.run_tensorflow(proto, expected_value)
self.assertEqual(actual_value, expected_value)
def test_federated_select_keys_must_be_fixed_length(
self, federated_select):
client_keys, max_key, server_val, select_fn = (
self.basic_federated_select_args())
unshape_type = computation_types.TensorType(tf.int32, [None])
unshape_proto, _ = tensorflow_computation_factory.create_identity(
unshape_type)
unshape = computation_impl.ConcreteComputation(
unshape_proto, context_stack_impl.context_stack)
bad_client_keys = intrinsics.federated_map(unshape, client_keys)
with self.assertRaises(TypeError):
federated_select(bad_client_keys, max_key, server_val, select_fn)
async def compute_federated_sum(
self, arg: FederatedComposingStrategyValue
) -> FederatedComposingStrategyValue:
type_analysis.check_federated_type(
arg.type_signature, placement=placement_literals.CLIENTS)
id_comp, id_type = tensorflow_computation_factory.create_identity(
arg.type_signature.member)
zero, plus, identity = await asyncio.gather(
executor_utils.embed_tf_constant(self._executor,
arg.type_signature.member, 0),
executor_utils.embed_tf_binary_operator(self._executor,
arg.type_signature.member,
tf.add),
self._executor.create_value(id_comp, id_type))
aggregate_args = await self._executor.create_struct(
[arg, zero, plus, plus, identity])
return await self.compute_federated_aggregate(aggregate_args)
def test_federated_aggregate_with_federated_zero_fails(self):
x = _mock_data_of_type(computation_types.at_clients(tf.int32))
zero = intrinsics.federated_value(0, placements.SERVER)
accumulate = _create_computation_add()
# The operator to use during the second stage simply adds total and count.
merge = _create_computation_add()
# The operator to use during the final stage simply computes the ratio.
report_type = computation_types.TensorType(tf.int32)
report_proto, _ = tensorflow_computation_factory.create_identity(
report_type)
report = computation_impl.ConcreteComputation(
report_proto, context_stack_impl.context_stack)
with self.assertRaisesRegex(
TypeError, 'Expected `zero` to be assignable to type int32, '
'but was of incompatible type int32@SERVER'):
intrinsics.federated_aggregate(x, zero, accumulate, merge, report)
def test_infers_accumulate_return_as_merge_arg_merge_return_as_report_arg(
self):
type_spec = computation_types.TensorType(dtype=tf.int64, shape=[None])
x = _mock_data_of_type(computation_types.at_clients(tf.int64))
def initialize_fn():
return tf.constant([], dtype=tf.int64, shape=[0])
initialize_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
initialize_fn, None)
initialize = computation_impl.ConcreteComputation(
initialize_proto, context_stack_impl.context_stack)
zero = initialize()
def _accumulate(arg):
return tf.concat([arg[0], [arg[1]]], 0)
accumulate_type = computation_types.StructType([type_spec, tf.int64])
accumulate_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_accumulate, accumulate_type)
accumulate = computation_impl.ConcreteComputation(
accumulate_proto, context_stack_impl.context_stack)
def _merge(arg):
return tf.concat([arg[0], arg[1]], 0)
merge_type = computation_types.StructType([type_spec, type_spec])
merge_proto, _ = tensorflow_computation_factory.create_computation_for_py_fn(
_merge, merge_type)
merge = computation_impl.ConcreteComputation(
merge_proto, context_stack_impl.context_stack)
report_proto, _ = tensorflow_computation_factory.create_identity(
type_spec)
report = computation_impl.ConcreteComputation(
report_proto, context_stack_impl.context_stack)
value = intrinsics.federated_aggregate(x, zero, accumulate, merge,
report)
self.assert_value(value, 'int64[?]@SERVER')
def create_dummy_computation_tensorflow_identity():
"""Returns a tensorflow computation and type `(float32 -> float32)`."""
tensor_type = computation_types.TensorType(tf.float32)
value, type_signature = tensorflow_computation_factory.create_identity(
tensor_type)
return value, type_signature
def create_dummy_computation_tensorflow_identity():
"""Returns a tensorflow computation and type `(float32 -> float32)`."""
type_spec = tf.float32
value = tensorflow_computation_factory.create_identity(type_spec)
type_signature = computation_types.FunctionType(type_spec, type_spec)
return value, type_signature
def create_whimsy_computation_tensorflow_identity(arg_type=tf.float32):
"""Returns a tensorflow computation and type `(float32 -> float32)`."""
value, type_signature = tensorflow_computation_factory.create_identity(
computation_types.to_type(arg_type))
return value, type_signature
def test_raises_type_error(self, type_signature):
with self.assertRaises(TypeError):
tensorflow_computation_factory.create_identity(type_signature)
def test_feeds_and_fetches_different(self):
proto, _ = tensorflow_computation_factory.create_identity(
computation_types.TensorType(tf.int32))
self.assertNotEqual(proto.tensorflow.parameter,
proto.tensorflow.result)
