def test_with_federated_map(self): eager_ex = eager_executor.EagerExecutor() federated_ex = federated_executor.FederatedExecutor({ None: eager_ex, placement_literals.SERVER: eager_ex }) ex = lambda_executor.LambdaExecutor(federated_ex) loop = asyncio.get_event_loop() @computations.tf_computation(tf.int32) def add_one(x): return x + 1 @computations.federated_computation(type_factory.at_server(tf.int32)) def comp(x): return intrinsics.federated_map(add_one, x) v1 = loop.run_until_complete(ex.create_value(comp)) v2 = loop.run_until_complete( ex.create_value(10, type_factory.at_server(tf.int32))) v3 = loop.run_until_complete(ex.create_call(v1, v2)) result = loop.run_until_complete(v3.compute()) self.assertEqual(result.numpy(), 11)
def test_with_federated_map_and_broadcast(self): eager_ex = eager_tf_executor.EagerTFExecutor() federated_ex = federating_executor.FederatingExecutor({ None: eager_ex, placement_literals.SERVER: eager_ex, placement_literals.CLIENTS: [eager_ex for _ in range(3)] }) ex = reference_resolving_executor.ReferenceResolvingExecutor( federated_ex) loop = asyncio.get_event_loop() @computations.tf_computation(tf.int32) def add_one(x): return x + 1 @computations.federated_computation(type_factory.at_server(tf.int32)) def comp(x): return intrinsics.federated_map(add_one, intrinsics.federated_broadcast(x)) v1 = loop.run_until_complete(ex.create_value(comp)) v2 = loop.run_until_complete( ex.create_value(10, type_factory.at_server(tf.int32))) v3 = loop.run_until_complete(ex.create_call(v1, v2)) result = loop.run_until_complete(v3.compute()) self.assertCountEqual([x.numpy() for x in result], [11, 11, 11])
def test_raises_with_closure(self): eager_ex = eager_tf_executor.EagerTFExecutor() federated_ex = federating_executor.FederatingExecutor({ None: eager_ex, placement_literals.SERVER: eager_ex, }) ex = reference_resolving_executor.ReferenceResolvingExecutor( federated_ex) loop = asyncio.get_event_loop() @computations.federated_computation(tf.int32, type_factory.at_server(tf.int32)) def foo(x, y): @computations.federated_computation(tf.int32) def bar(z): del z return x return intrinsics.federated_map(bar, y) v1 = loop.run_until_complete(ex.create_value(foo)) v2 = loop.run_until_complete( ex.create_value( [0, 0], [tf.int32, type_factory.at_server(tf.int32)])) with self.assertRaisesRegex( RuntimeError, 'lambda passed to intrinsic contains references to captured variables' ): loop.run_until_complete(ex.create_call(v1, v2))
def create_dummy_intrinsic_def_federated_apply(): value = intrinsic_defs.FEDERATED_APPLY type_signature = computation_types.FunctionType([ type_factory.unary_op(tf.float32), type_factory.at_server(tf.float32), ], type_factory.at_server(tf.float32)) return value, type_signature
async def _compute_intrinsic_federated_aggregate(self, arg): 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._child_executors)) identity_report = _create_lambda_identity_comp(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.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 = tuple(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._child_executors, val)]) parent_vals = await asyncio.gather( *[self._parent_executor.create_value(v, zero_type) for v in vals]) parent_merge, parent_report = tuple(await asyncio.gather( self._parent_executor.create_value(merge, merge_type), self._parent_executor.create_value(report, report_type))) merge_result = parent_vals[0] for next_val in parent_vals[1:]: merge_result = await self._parent_executor.create_call( parent_merge, await self._parent_executor.create_tuple([merge_result, next_val])) return CompositeValue( await self._parent_executor.create_call(parent_report, merge_result), type_factory.at_server(report_type.result))
def test_executor_create_value_with_valid_intrinsic_def(self): loop = asyncio.get_event_loop() ex = _make_test_executor() val = loop.run_until_complete( ex.create_value( intrinsic_defs.FEDERATED_APPLY, computation_types.FunctionType([ type_factory.unary_op(tf.int32), type_factory.at_server(tf.int32) ], type_factory.at_server(tf.int32)))) self.assertIsInstance(val, federated_executor.FederatedExecutorValue) self.assertEqual(str(val.type_signature), '(<(int32 -> int32),int32@SERVER> -> int32@SERVER)') self.assertIs(val.internal_representation, intrinsic_defs.FEDERATED_APPLY)
def create_dummy_intrinsic_def(): """Returns a `intrinsic_defs.IntrinsicDef` and type.""" value = intrinsic_defs.FEDERATED_EVAL_AT_SERVER type_signature = computation_types.FunctionType( computation_types.FunctionType(None, tf.int32), type_factory.at_server(tf.int32)) return value, type_signature
def test_with_temperature_sensor_example(self, executor): @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_factory.at_clients(computation_types.SequenceType( tf.float32)), type_factory.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)) with executor_test_utils.install_executor(executor): 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)
class CreateIdentityTest(parameterized.TestCase): # pyformat: disable @parameterized.named_parameters( ('int', computation_types.TensorType(tf.int32), 10), ('unnamed_tuple', computation_types.NamedTupleType([tf.int32, tf.float32]), anonymous_tuple.AnonymousTuple([(None, 10), (None, 10.0)])), ('named_tuple', computation_types.NamedTupleType([ ('a', tf.int32), ('b', tf.float32) ]), anonymous_tuple.AnonymousTuple([('a', 10), ('b', 10.0)])), ('sequence', computation_types.SequenceType(tf.int32), [10] * 3), ) # pyformat: enable 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) @parameterized.named_parameters( ('none', None), ('federated_type', type_factory.at_server(tf.int32)), ) def test_raises_type_error(self, type_signature): with self.assertRaises(TypeError): tensorflow_computation_factory.create_identity(type_signature)
def test_get_size_info(self, num_clients): @computations.federated_computation( type_factory.at_clients(computation_types.SequenceType(tf.float32)), type_factory.at_server(tf.float32)) def comp(temperatures, threshold): client_data = [temperatures, intrinsics.federated_broadcast(threshold)] result_map = intrinsics.federated_map( count_over, intrinsics.federated_zip(client_data)) count_map = intrinsics.federated_map(count_total, temperatures) return intrinsics.federated_mean(result_map, count_map) factory = executor_stacks.sizing_executor_factory(num_clients=num_clients) default_executor.set_default_executor(factory) to_float = lambda x: tf.cast(x, tf.float32) temperatures = [tf.data.Dataset.range(10).map(to_float)] * num_clients threshold = 15.0 comp(temperatures, threshold) # Each client receives a tf.float32 and uploads two tf.float32 values. expected_broadcast_bits = num_clients * 32 expected_aggregate_bits = expected_broadcast_bits * 2 expected = ({ (('CLIENTS', num_clients),): [[1, tf.float32]] * num_clients }, { (('CLIENTS', num_clients),): [[1, tf.float32]] * num_clients * 2 }, [expected_broadcast_bits], [expected_aggregate_bits]) self.assertEqual(expected, factory.get_size_info())
def _temperature_sensor_example_next_fn(): @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_factory.at_clients(computation_types.SequenceType(tf.float32)), type_factory.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)) return comp
def create_dummy_intrinsic_def_federated_secure_sum(): value = intrinsic_defs.FEDERATED_SECURE_SUM type_signature = computation_types.FunctionType([ type_factory.at_clients(tf.float32), tf.float32, ], type_factory.at_server(tf.float32)) return value, type_signature
async def _get_cardinalities_helper(): """Helper function which does the actual work of fetching cardinalities.""" one_type = type_factory.at_clients(tf.int32, all_equal=True) sum_type = computation_types.FunctionType( type_factory.at_clients(tf.int32), type_factory.at_server(tf.int32)) sum_comp = executor_utils.create_intrinsic_comp( intrinsic_defs.FEDERATED_SUM, sum_type) async def _count_leaf_executors(ex): """Counts the total number of leaf executors under `ex`.""" one_fut = ex.create_value(1, one_type) sum_comp_fut = ex.create_value(sum_comp, sum_type) one_val, sum_comp_val = tuple(await asyncio.gather( one_fut, sum_comp_fut)) sum_result = await (await ex.create_call(sum_comp_val, one_val)).compute() if isinstance(sum_result, tf.Tensor): return sum_result.numpy() else: return sum_result return await asyncio.gather( *[_count_leaf_executors(c) for c in self._child_executors])
def test_returns_value_with_federated_type_server(self): value = [eager_tf_executor.EagerValue(10.0, None, tf.float32)] type_signature = type_factory.at_server(tf.float32) value = federating_executor.FederatingExecutorValue(value, type_signature) result = self.run_sync(value.compute()) self.assertEqual(result, 10.0)
def create_dummy_intrinsic_def_federated_reduce(): value = intrinsic_defs.FEDERATED_REDUCE type_signature = computation_types.FunctionType([ type_factory.at_clients(tf.float32), tf.float32, type_factory.reduction_op(tf.float32, tf.float32), ], type_factory.at_server(tf.float32)) return value, type_signature
class CreateBinaryOperatorTest(parameterized.TestCase): # pyformat: disable @parameterized.named_parameters( ('add_int', tf.math.add, computation_types.TensorType( tf.int32), [1, 2], 3), ('add_float', tf.math.add, computation_types.TensorType( tf.float32), [1.0, 2.25], 3.25), ('add_unnamed_tuple', tf.math.add, computation_types.NamedTupleType([tf.int32, tf.float32]), [ [1, 1.0], [2, 2.25] ], anonymous_tuple.AnonymousTuple([(None, 3), (None, 3.25)])), ('add_named_tuple', tf.math.add, computation_types.NamedTupleType([ ('a', tf.int32), ('b', tf.float32) ]), [[1, 1.0], [2, 2.25] ], anonymous_tuple.AnonymousTuple([('a', 3), ('b', 3.25)])), ('multiply_int', tf.math.multiply, computation_types.TensorType(tf.int32), [2, 2], 4), ('multiply_float', tf.math.multiply, computation_types.TensorType(tf.float32), [2.0, 2.25], 4.5), ('divide_int', tf.math.divide, computation_types.TensorType( tf.int32), [4, 2], 2.0), ('divide_float', tf.math.divide, computation_types.TensorType(tf.float32), [4.0, 2.0], 2.0), ('divide_inf', tf.math.divide, computation_types.TensorType( tf.int32), [1, 0], np.inf), ) # pyformat: enable def test_returns_computation(self, operator, type_signature, operands, expected_result): proto = tensorflow_computation_factory.create_binary_operator( operator, type_signature) self.assertIsInstance(proto, pb.Computation) actual_type = type_serialization.deserialize_type(proto.type) self.assertIsInstance(actual_type, computation_types.FunctionType) # Note: It is only useful to test the parameter type; the result type # depends on the `operator` used, not the implemenation # `create_binary_operator`. expected_parameter_type = computation_types.NamedTupleType( [type_signature, type_signature]) self.assertEqual(actual_type.parameter, expected_parameter_type) actual_result = test_utils.run_tensorflow(proto, operands) self.assertEqual(actual_result, expected_result) @parameterized.named_parameters( ('non_callable_operator', 1, computation_types.TensorType(tf.int32)), ('none_type', tf.math.add, None), ('federated_type', tf.math.add, type_factory.at_server(tf.int32)), ('sequence_type', tf.math.add, computation_types.SequenceType( tf.int32)), ) def test_raises_type_error(self, operator, type_signature): with self.assertRaises(TypeError): tensorflow_computation_factory.create_binary_operator( operator, type_signature)
def _federated_mean(self, arg): type_utils.check_federated_type(arg.type_signature, None, placements.CLIENTS, False) py_typecheck.check_type(arg.value, list) server_sum = self._federated_sum(arg) unplaced_avg = multiply_by_scalar( ComputedValue(server_sum.value, server_sum.type_signature.member), 1.0 / float(len(arg.value))) return ComputedValue(unplaced_avg.value, type_factory.at_server(unplaced_avg.type_signature))
def test_serialize_deserialize_federated_at_server(self): x = 10 x_type = type_factory.at_server(tf.int32) value_proto, value_type = executor_service_utils.serialize_value( x, x_type) self.assertIsInstance(value_proto, executor_pb2.Value) self.assertEqual(str(value_type), 'int32@SERVER') y, type_spec = executor_service_utils.deserialize_value(value_proto) self.assertEqual(str(type_spec), str(x_type)) self.assertEqual(y, 10)
def test_with_removal_of_identity_mapping(self): @computations.federated_computation(type_factory.at_server(tf.int32)) def comp(x): return intrinsics.federated_map(_identity, x) def transformation_fn(x): x, _ = tree_transformations.remove_mapped_or_applied_identity(x) return x self.assertEqual(_test_create_value(comp, transformation_fn), '(FEDERATED_arg -> FEDERATED_arg)')
def test_executor_create_value_with_server_int(self): loop = asyncio.get_event_loop() ex = _make_test_executor() val = loop.run_until_complete( ex.create_value(10, type_factory.at_server(tf.int32))) self.assertIsInstance(val, federated_executor.FederatedExecutorValue) self.assertEqual(str(val.type_signature), 'int32@SERVER') self.assertIsInstance(val.internal_representation, list) self.assertLen(val.internal_representation, 1) self.assertIsInstance(val.internal_representation[0], eager_executor.EagerValue) self.assertEqual( val.internal_representation[0].internal_representation.numpy(), 10)
def _federated_zip_at_server(self, arg): py_typecheck.check_type(arg.type_signature, computation_types.NamedTupleType) for idx in range(len(arg.type_signature)): type_utils.check_federated_type(arg.type_signature[idx], None, placements.SERVER, True) return ComputedValue( arg.value, type_factory.at_server( computation_types.NamedTupleType([ (k, v.member) if k else v.member for k, v in anonymous_tuple.iter_elements(arg.type_signature) ])))
def test_returns_value_with_unplaced_type_and_server(self, executor): value, type_signature = executor_test_utils.create_dummy_value_unplaced( ) value = self.run_sync(executor.create_value(value, type_signature)) result = self.run_sync( executor_utils.compute_intrinsic_federated_value( executor, value, placement_literals.SERVER)) self.assertIsInstance(result, executor_value_base.ExecutorValue) expected_type = type_factory.at_server(type_signature) self.assertEqual(result.type_signature.compact_representation(), expected_type.compact_representation()) actual_result = self.run_sync(result.compute()) self.assertEqual(actual_result, 10.0)
def test_with_inlining_of_blocks(self): @computations.federated_computation(type_factory.at_server(tf.int32)) def comp(x): return intrinsics.federated_zip([x, x]) # TODO(b/134543154): Slide in something more powerful so that this test # doesn't break when the implementation changes; for now, this will do. def transformation_fn(x): x, _ = transformations.remove_mapped_or_applied_identity(x) x, _ = transformations.inline_block_locals(x) x, _ = transformations.replace_selection_from_tuple_with_element(x) return x self.assertIn('federated_zip_at_server(<FEDERATED_arg,FEDERATED_arg>)', _test_create_value(comp, transformation_fn))
class CreateConstantTest(parameterized.TestCase): # pyformat: disable @parameterized.named_parameters( ('int', 10, computation_types.TensorType(tf.int32, [3]), [10] * 3), ('float', 10.0, computation_types.TensorType(tf.float32, [3]), [10.0] * 3), ('unnamed_tuple', 10, computation_types.NamedTupleType( [tf.int32] * 3), anonymous_tuple.AnonymousTuple([(None, 10)] * 3)), ('named_tuple', 10, computation_types.NamedTupleType([ ('a', tf.int32), ('b', tf.int32), ('c', tf.int32) ]), anonymous_tuple.AnonymousTuple([('a', 10), ('b', 10), ('c', 10)])), ('nested_tuple', 10, computation_types.NamedTupleType([[tf.int32] * 3] * 3), anonymous_tuple.AnonymousTuple( [(None, anonymous_tuple.AnonymousTuple([(None, 10)] * 3))] * 3)), ) # pyformat: enable def test_returns_computation(self, value, type_signature, expected_result): proto = tensorflow_computation_factory.create_constant( value, type_signature) self.assertIsInstance(proto, pb.Computation) actual_type = type_serialization.deserialize_type(proto.type) expected_type = computation_types.FunctionType(None, type_signature) self.assertEqual(actual_type, expected_type) actual_result = test_utils.run_tensorflow(proto) if isinstance(expected_result, list): self.assertCountEqual(actual_result, expected_result) else: self.assertEqual(actual_result, expected_result) @parameterized.named_parameters( ('non_scalar_value', np.zeros( [1]), computation_types.TensorType(tf.int32)), ('none_type', 10, None), ('federated_type', 10, type_factory.at_server(tf.int32)), ('bad_type', 10.0, computation_types.TensorType(tf.int32)), ) def test_raises_type_error(self, value, type_signature): with self.assertRaises(TypeError): tensorflow_computation_factory.create_constant( value, type_signature)