def test_fit_argument(self): old_arg = reference_executor.ComputedValue( anonymous_tuple.AnonymousTuple([('A', 10)]), [('A', type_constructors.at_clients(tf.int32, all_equal=True))]) new_arg = reference_executor.fit_argument( old_arg, [('A', type_constructors.at_clients(tf.int32))], reference_executor.ComputationContext( cardinalities={placements.CLIENTS: 3})) self.assertEqual(str(new_arg.type_signature), '<A={int32}@CLIENTS>') self.assertEqual(new_arg.value.A, [10, 10, 10])
def test_with_federated_value_as_a_non_py_list(self, val): loop = asyncio.get_event_loop() ex = _make_test_executor(num_clients=4) v = loop.run_until_complete( ex.create_value(val, type_constructors.at_clients(tf.int32))) self.assertEqual(str(v.type_signature), '{int32}@CLIENTS') result = tf.nest.map_structure(lambda x: x.numpy(), loop.run_until_complete(v.compute())) self.assertCountEqual(result, [1, 2, 3, 4])
def test_executor_create_value_with_all_equal_client_int(self): loop = asyncio.get_event_loop() ex = _make_test_executor(3) val = loop.run_until_complete( ex.create_value( 10, type_constructors.at_clients(tf.int32, all_equal=True))) self.assertIsInstance(val, federated_executor.FederatedExecutorValue) self.assertEqual(str(val.type_signature), 'int32@CLIENTS') self.assertIsInstance(val.internal_representation, list) self.assertLen(val.internal_representation, 3) for v in val.internal_representation: self.assertIsInstance(v, eager_executor.EagerValue) self.assertEqual(v.internal_representation.numpy(), 10)
def test_federated_mean_with_floats(self): loop = asyncio.get_event_loop() ex = _make_test_executor(4) v1 = loop.run_until_complete( ex.create_value([1.0, 2.0, 3.0, 4.0], type_constructors.at_clients(tf.float32))) self.assertEqual(str(v1.type_signature), '{float32}@CLIENTS') v2 = loop.run_until_complete( ex.create_value( intrinsic_defs.FEDERATED_MEAN, computation_types.FunctionType( type_constructors.at_clients(tf.float32), type_constructors.at_server(tf.float32)))) self.assertEqual(str(v2.type_signature), '({float32}@CLIENTS -> float32@SERVER)') v3 = loop.run_until_complete(ex.create_call(v2, v1)) self.assertEqual(str(v3.type_signature), 'float32@SERVER') result = loop.run_until_complete(v3.compute()) self.assertEqual(result.numpy(), 2.5)
def test_federated_weighted_mean_with_floats(self): loop = asyncio.get_event_loop() ex = _make_test_executor(num_clients=4, use_lambda_executor=True) v1 = loop.run_until_complete( ex.create_value([1.0, 2.0, 3.0, 4.0], type_constructors.at_clients(tf.float32))) self.assertEqual(str(v1.type_signature), '{float32}@CLIENTS') v2 = loop.run_until_complete( ex.create_value([5.0, 10.0, 3.0, 2.0], type_constructors.at_clients(tf.float32))) self.assertEqual(str(v2.type_signature), '{float32}@CLIENTS') v3 = loop.run_until_complete( ex.create_tuple( anonymous_tuple.AnonymousTuple([(None, v1), (None, v2)]))) self.assertEqual(str(v3.type_signature), '<{float32}@CLIENTS,{float32}@CLIENTS>') v4 = loop.run_until_complete( ex.create_value( intrinsic_defs.FEDERATED_WEIGHTED_MEAN, computation_types.FunctionType([ type_constructors.at_clients(tf.float32), type_constructors.at_clients(tf.float32) ], type_constructors.at_server(tf.float32)))) self.assertEqual( str(v4.type_signature), '(<{float32}@CLIENTS,{float32}@CLIENTS> -> float32@SERVER)') v5 = loop.run_until_complete(ex.create_call(v4, v3)) self.assertEqual(str(v5.type_signature), 'float32@SERVER') result = loop.run_until_complete(v5.compute()) self.assertAlmostEqual(result.numpy(), 2.1, places=3)
async def _compute_intrinsic_federated_broadcast(self, arg): py_typecheck.check_type(arg.type_signature, computation_types.FederatedType) py_typecheck.check_type(arg.internal_representation, list) if not arg.type_signature.all_equal: raise ValueError('Cannot broadcast a non all_equal value.') if len(arg.internal_representation) != 1: raise ValueError( 'Cannot broadcast a with a non-singleton representation.') val = await arg.internal_representation[0].compute() return FederatedExecutorValue( await asyncio.gather(*[ c.create_value(val, arg.type_signature.member) for c in self._target_executors[placement_literals.CLIENTS] ]), type_constructors.at_clients(arg.type_signature.member, all_equal=True))
def test_with_temperature_sensor_example(self): @computations.tf_computation(computation_types.SequenceType( tf.float32), tf.float32) def count_over(ds, t): return ds.reduce( np.float32(0), lambda n, x: n + tf.cast(tf.greater(x, t), tf.float32)) @computations.tf_computation(computation_types.SequenceType(tf.float32) ) def count_total(ds): return ds.reduce(np.float32(0.0), lambda n, _: n + 1.0) @computations.federated_computation( type_constructors.at_clients( computation_types.SequenceType(tf.float32)), type_constructors.at_server(tf.float32)) def comp(temperatures, threshold): return intrinsics.federated_mean( intrinsics.federated_map( count_over, intrinsics.federated_zip([ temperatures, intrinsics.federated_broadcast(threshold) ])), intrinsics.federated_map(count_total, temperatures)) num_clients = 3 set_default_executor.set_default_executor( executor_stacks.create_local_executor(num_clients)) to_float = lambda x: tf.cast(x, tf.float32) temperatures = [ tf.data.Dataset.range(10).map(to_float), tf.data.Dataset.range(20).map(to_float), tf.data.Dataset.range(30).map(to_float) ] threshold = 15.0 result = comp(temperatures, threshold) self.assertAlmostEqual(result, 8.333, places=3) set_default_executor.set_default_executor()
def test_with_incomplete_temperature_sensor_example(self): @computations.federated_computation( type_constructors.at_clients( computation_types.SequenceType(tf.float32)), type_constructors.at_server(tf.float32)) def comp(temperatures, threshold): @computations.tf_computation( computation_types.SequenceType(tf.float32), tf.float32) def count(ds, t): return ds.reduce( np.int32(0), lambda n, x: n + tf.cast(tf.greater(x, t), tf.int32)) return intrinsics.federated_map( count, intrinsics.federated_zip( [temperatures, intrinsics.federated_broadcast(threshold)])) num_clients = 10 set_default_executor.set_default_executor( executor_stacks.create_local_executor(num_clients)) temperatures = [ tf.data.Dataset.range(1000).map(lambda x: tf.cast(x, tf.float32)) for _ in range(num_clients) ] threshold = 100.0 result = comp(temperatures, threshold) self.assertCountEqual([x.numpy() for x in result], [899 for _ in range(num_clients)]) set_default_executor.set_default_executor()
# intrinsics defined above, as follows. # # @federated_computation # def federated_aggregate(x, zero, accumulate, merge, report): # 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',
def test_at_clients(self): self.assertEqual(str(type_constructors.at_clients(tf.bool)), '{bool}@CLIENTS')