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)