def test_federated_sum(self):
bodies = intrinsic_bodies.get_intrinsic_bodies(
context_stack_impl.context_stack)
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
return bodies[intrinsic_defs.FEDERATED_SUM.uri](x)
self.assertEqual(
str(foo.type_signature), '({int32}@CLIENTS -> int32@SERVER)')
body_string = (r'\(FEDERATED_arg -> '
r'federated_aggregate\(<FEDERATED_arg,generic_zero,'
r'\(binary_operator_arg'
r' -> comp#[a-z0-9]+\(<binary_operator_arg\[0\],'
r'binary_operator_arg\[1\]>\)\),\(binary_operator_arg'
r' -> comp#[a-z0-9]+\(<binary_operator_arg\[0\],'
r'binary_operator_arg\[1\]>\)\),comp#[a-z0-9]+>\)\)')
self.assertRegexMatch(_body_str(foo), [body_string])
self.assertEqual(foo([1]), 1)
self.assertEqual(foo([1, 2, 3]), 6)
def create_dummy_called_federated_map_all_equal(parameter_name,
parameter_type=tf.int32):
r"""Returns a dummy called federated map.
Call
/ \
federated_map_all_equal Tuple
|
[Lambda(x), data]
|
Ref(x)
Args:
parameter_name: The name of the parameter.
parameter_type: The type of the parameter.
"""
fn = create_identity_function(parameter_name, parameter_type)
arg_type = computation_types.FederatedType(parameter_type,
placements.CLIENTS,
all_equal=True)
arg = computation_building_blocks.Data('data', arg_type)
return computation_constructing_utils.create_federated_map_all_equal(
fn, arg)
def test_federated_generic_add_with_named_tuples(self):
bodies = intrinsic_bodies.get_intrinsic_bodies(
context_stack_impl.context_stack)
@computations.federated_computation(
computation_types.FederatedType([('a', tf.int32),
('b', tf.float32)],
placement_literals.CLIENTS))
def foo(x):
return bodies[intrinsic_defs.GENERIC_PLUS.uri]([x, x])
self.assertEqual(
str(foo.type_signature),
'({<a=int32,b=float32>}@CLIENTS -> {<a=int32,b=float32>}@CLIENTS)')
self.assertEqual(foo(
[[1,
1.]]), [anonymous_tuple.AnonymousTuple([('a', 2), ('b', 2.)])])
self.assertEqual(foo([[1, 1.], [1, 2.], [1, 3.]]), [
anonymous_tuple.AnonymousTuple([('a', 2), ('b', 2.)]),
anonymous_tuple.AnonymousTuple([('a', 2), ('b', 4.)]),
anonymous_tuple.AnonymousTuple([('a', 2), ('b', 6.)])
])
def test_federated_generic_multiply_with_unnamed_tuples(self):
bodies = intrinsic_bodies.get_intrinsic_bodies(
context_stack_impl.context_stack)
@computations.federated_computation(
computation_types.FederatedType([tf.int32, tf.float32],
placements.CLIENTS))
def foo(x):
return bodies[intrinsic_defs.GENERIC_MULTIPLY.uri]([x, x])
self.assertEqual(
str(foo.type_signature),
'({<int32,float32>}@CLIENTS -> {<int32,float32>}@CLIENTS)')
self.assertEqual(
foo([[1, 1.]]),
[anonymous_tuple.AnonymousTuple([(None, 1.), (None, 1.)])])
self.assertEqual(
foo([[1, 1.], [1, 2.], [1, 3.]]), [
anonymous_tuple.AnonymousTuple([(None, 1), (None, 1.)]),
anonymous_tuple.AnonymousTuple([(None, 1), (None, 4.)]),
anonymous_tuple.AnonymousTuple([(None, 1), (None, 9.)])
])
def test_create_federated_named_tuple_one(self):
tuple_type = [('a', computation_types.TensorType(tf.float32, [2, 2])),
('b', computation_types.TensorType(tf.float32, [2, 2]))]
fed_type = computation_types.FederatedType(tuple_type,
placement_literals.SERVER)
fed_zero = intrinsic_utils.construct_generic_constant(fed_type, 1)
self.assertEqual(fed_zero.type_signature.member, fed_type.member)
self.assertEqual(fed_zero.type_signature.placement, fed_type.placement)
self.assertTrue(fed_zero.type_signature.all_equal)
self.assertIsInstance(fed_zero, computation_building_blocks.Call)
self.assertIsInstance(fed_zero.function,
computation_building_blocks.Intrinsic)
self.assertEqual(fed_zero.function.uri,
intrinsic_defs.FEDERATED_VALUE_AT_SERVER.uri)
self.assertIsInstance(fed_zero.argument,
computation_building_blocks.Call)
executable_unplaced_fn = computation_wrapper_instances.building_block_to_computation(
fed_zero.argument.function)
self.assertLen(executable_unplaced_fn(), 2)
self.assertTrue(
np.array_equal(executable_unplaced_fn().a, np.ones([2, 2])))
self.assertTrue(
np.array_equal(executable_unplaced_fn().b, np.ones([2, 2])))
def _transform(comp):
"""Returns a new transformed computation or `comp`."""
if not _should_transform(comp):
return comp, False
named_comps = anonymous_tuple.to_elements(comp)
names = [name for name, _ in named_comps]
elements = _get_comps(comp)
comps = _transform_comps(names, elements)
arg = computation_building_blocks.Tuple(comps)
first_comp = comp[0]
parameter_type = computation_types.to_type(arg.type_signature)
type_signature = [
(name, call.type_signature.member) for name, call in named_comps
]
result_type = computation_types.FederatedType(
type_signature, first_comp.type_signature.placement)
intrinsic_type = computation_types.FunctionType(parameter_type, result_type)
intrinsic = computation_building_blocks.Intrinsic(first_comp.function.uri,
intrinsic_type)
call = computation_building_blocks.Call(intrinsic, arg)
transformed_comp = computation_constructing_utils.create_federated_unzip(
call)
return transformed_comp, True
def test_federated_sum_reduces_to_aggregate(self):
uri = intrinsic_defs.FEDERATED_SUM.uri
@computations.federated_computation(
computation_types.FederatedType(tf.float32, placement_literals.CLIENTS))
def foo(x):
return intrinsics.federated_sum(x)
foo_building_block = building_blocks.ComputationBuildingBlock.from_proto(
foo._computation_proto)
count_sum_before_reduction = _count_intrinsics(foo_building_block, uri)
reduced, modified = intrinsic_reductions.replace_intrinsics_with_bodies(
foo_building_block)
count_sum_after_reduction = _count_intrinsics(reduced, uri)
count_aggregations = _count_intrinsics(
reduced, intrinsic_defs.FEDERATED_AGGREGATE.uri)
self.assertTrue(modified)
self.assert_types_identical(foo_building_block.type_signature,
reduced.type_signature)
self.assertGreater(count_sum_before_reduction, 0)
self.assertEqual(count_sum_after_reduction, 0)
self.assertGreater(count_aggregations, 0)
def test_federated_sum_named_tuples(self):
bodies = intrinsic_bodies.get_intrinsic_bodies(
context_stack_impl.context_stack)
@computations.federated_computation(
computation_types.FederatedType([('a', tf.int32),
('b', tf.float32)],
placement_literals.CLIENTS))
def foo(x):
return bodies[intrinsic_defs.FEDERATED_SUM.uri](x)
self.assertEqual(
str(foo.type_signature),
'({<a=int32,b=float32>}@CLIENTS -> <a=int32,b=float32>@SERVER)')
self.assertDictEqual(anonymous_tuple.to_odict(foo([[1, 2.]])), {
'a': 1,
'b': 2.
})
self.assertDictEqual(anonymous_tuple.to_odict(foo([[1, 2.], [3, 4.]])),
{
'a': 4,
'b': 6.
})
class IsStructureOfIntegersTest(parameterized.TestCase):
@parameterized.named_parameters(
('int', computation_types.TensorType(tf.int32)),
('ints', computation_types.NamedTupleType([tf.int32, tf.int32])),
('federated_int_at_clients',
computation_types.FederatedType(tf.int32,
placement_literals.CLIENTS)),
)
def test_returns_true(self, type_spec):
self.assertTrue(type_analysis.is_structure_of_integers(type_spec))
@parameterized.named_parameters(
('bool', computation_types.TensorType(tf.bool)),
('string', computation_types.TensorType(tf.string)),
('int_and_bool', computation_types.NamedTupleType([tf.int32, tf.bool
])),
('sequence_of_ints', computation_types.SequenceType(tf.int32)),
('placement', computation_types.PlacementType()),
('function', computation_types.FunctionType(tf.int32, tf.int32)),
('abstract', computation_types.AbstractType('T')),
)
def test_returns_false(self, type_spec):
self.assertFalse(type_analysis.is_structure_of_integers(type_spec))
def test_federated_aggregate_with_unknown_dimension(self):
Accumulator = collections.namedtuple('Accumulator', ['samples']) # pylint: disable=invalid-name
accumulator_type = computation_types.StructType(
Accumulator(samples=computation_types.TensorType(dtype=tf.int32,
shape=[None])))
@computations.tf_computation()
def build_empty_accumulator():
return Accumulator(samples=tf.zeros(shape=[0], dtype=tf.int32))
# The operator to use during the first stage simply adds an element to the
# tensor, increasing its size.
@computations.tf_computation([accumulator_type, tf.int32])
def accumulate(accu, elem):
return Accumulator(samples=tf.concat(
[accu.samples, tf.expand_dims(elem, axis=0)], axis=0))
# The operator to use during the second stage simply adds total and count.
@computations.tf_computation([accumulator_type, accumulator_type])
def merge(x, y):
return Accumulator(
samples=tf.concat([x.samples, y.samples], axis=0))
# The operator to use during the final stage simply computes the ratio.
@computations.tf_computation(accumulator_type)
def report(accu):
return accu
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
val = intrinsics.federated_aggregate(x, build_empty_accumulator(),
accumulate, merge, report)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '({int32}@CLIENTS -> <samples=int32[?]>@SERVER)')
def create_dummy_called_federated_aggregate(accumulate_parameter_name,
merge_parameter_name,
report_parameter_name,
value_type=tf.int32):
r"""Returns a dummy called federated aggregate.
Call
/ \
federated_aggregate Tuple
|
[data, data, Lambda(x), Lambda(x), Lambda(x)]
| | |
data data data
Args:
accumulate_parameter_name: The name of the accumulate parameter.
merge_parameter_name: The name of the merge parameter.
report_parameter_name: The name of the report parameter.
value_type: The TFF type of the value to be aggregated, placed at
CLIENTS.
"""
federated_value_type = computation_types.FederatedType(
value_type, placements.CLIENTS)
value = building_blocks.Data('data', federated_value_type)
zero = building_blocks.Data('data', tf.float32)
accumulate_type = computation_types.NamedTupleType((tf.float32, value_type))
accumulate_result = building_blocks.Data('data', tf.float32)
accumulate = building_blocks.Lambda(accumulate_parameter_name,
accumulate_type, accumulate_result)
merge_type = computation_types.NamedTupleType((tf.float32, tf.float32))
merge_result = building_blocks.Data('data', tf.float32)
merge = building_blocks.Lambda(merge_parameter_name, merge_type, merge_result)
report_result = building_blocks.Data('data', tf.bool)
report = building_blocks.Lambda(report_parameter_name, tf.float32,
report_result)
return building_block_factory.create_federated_aggregate(
value, zero, accumulate, merge, report)
def create_unweighted(
self,
value_type: factory.ValueType) -> aggregation_process.AggregationProcess:
py_typecheck.check_type(value_type, factory.ValueType.__args__)
if not all([t.dtype.is_floating for t in structure.flatten(value_type)]):
raise TypeError(f'All values in provided value_type must be of floating '
f'dtype. Provided value_type: {value_type}')
value_sum_process = self._value_sum_factory.create_unweighted(value_type)
@computations.federated_computation()
def init_fn():
state = collections.OrderedDict(
value_sum_process=value_sum_process.initialize())
return intrinsics.federated_zip(state)
@computations.federated_computation(init_fn.type_signature.result,
computation_types.FederatedType(
value_type, placements.CLIENTS))
def next_fn(state, value):
value_sum_output = value_sum_process.next(state['value_sum_process'],
value)
count = intrinsics.federated_sum(
intrinsics.federated_value(1, placements.CLIENTS))
mean_value = intrinsics.federated_map(_div,
(value_sum_output.result, count))
state = collections.OrderedDict(value_sum_process=value_sum_output.state)
measurements = collections.OrderedDict(
value_sum_process=value_sum_output.measurements)
return measured_process.MeasuredProcessOutput(
intrinsics.federated_zip(state), mean_value,
intrinsics.federated_zip(measurements))
return aggregation_process.AggregationProcess(init_fn, next_fn)
def test_computation_with_int_sequence_raises(self):
batch_size = 1
ds1_shape = tf.TensorShape([batch_size])
sequence_type = computation_types.SequenceType(
computation_types.TensorType(tf.int32, ds1_shape))
federated_type = computation_types.FederatedType(
sequence_type, placements.CLIENTS)
@computations.tf_computation(sequence_type)
def foo(z):
value1 = z.reduce(0, lambda x, y: x + tf.reduce_sum(y))
return value1
@computations.federated_computation(federated_type)
def bar(x):
return intrinsics.federated_map(foo, x)
ds1 = tf.data.Dataset.from_tensor_slices([10, 20]).batch(batch_size)
ds2 = tf.data.Dataset.from_tensor_slices([30, 40]).batch(batch_size)
with self.assertRaisesRegexp(ValueError, 'Please pass a list'):
bar(ds1)
with self.assertRaisesRegexp(ValueError, 'Please pass a list'):
bar(ds2)
class IsSumCompatibleTest(parameterized.TestCase):
@parameterized.named_parameters([
('tensor_type', computation_types.TensorType(tf.int32)),
('tuple_type_int', computation_types.StructType([tf.int32, tf.int32],)),
('tuple_type_float',
computation_types.StructType([tf.complex128, tf.float32, tf.float64])),
('federated_type',
computation_types.FederatedType(tf.int32, placement_literals.CLIENTS)),
])
def test_positive_examples(self, type_spec):
self.assertTrue(type_analysis.is_sum_compatible(type_spec))
@parameterized.named_parameters([
('tensor_type_bool', computation_types.TensorType(tf.bool)),
('tensor_type_string', computation_types.TensorType(tf.string)),
('tuple_type', computation_types.StructType([tf.int32, tf.bool])),
('sequence_type', computation_types.SequenceType(tf.int32)),
('placement_type', computation_types.PlacementType()),
('function_type', computation_types.FunctionType(tf.int32, tf.int32)),
('abstract_type', computation_types.AbstractType('T')),
])
def test_negative_examples(self, type_spec):
self.assertFalse(type_analysis.is_sum_compatible(type_spec))
def test_generic_add_federated_named_tuple_by_tensor(self):
bodies = intrinsic_bodies.get_intrinsic_bodies(
context_stack_impl.context_stack)
@computations.federated_computation(
computation_types.FederatedType([[('a', tf.float32),
('b', tf.float32)], tf.float32],
placement_literals.CLIENTS))
def foo(x):
return bodies[intrinsic_defs.GENERIC_PLUS.uri]([x[0], x[1]])
self.assertEqual(
str(foo.type_signature),
'({<<a=float32,b=float32>,float32>}@CLIENTS -> {<a=float32,b=float32>}@CLIENTS)'
)
self.assertEqual(
foo([[[1., 1.], 1.]]), [structure.Struct([('a', 2.), ('b', 2.)])])
self.assertEqual(
foo([[[1., 1.], 1.], [[1., 2.], 2.], [[1., 4.], 4.]]), [
structure.Struct([('a', 2.), ('b', 2.)]),
structure.Struct([('a', 3.), ('b', 4.)]),
structure.Struct([('a', 5.), ('b', 8.)])
])
def test_batching_namedtuple_dataset(self):
batch_type = collections.namedtuple('Batch', ['x', 'y'])
federated_sequence_type = computation_types.FederatedType(
computation_types.SequenceType(
batch_type(
x=computation_types.TensorType(tf.float32, [None, 2]),
y=computation_types.TensorType(tf.float32, [None, 1]))),
placements.CLIENTS,
all_equal=False)
@computations.tf_computation(federated_sequence_type.member)
def test_batch_select_and_reduce(z):
i = z.map(lambda x: x.y)
return i.reduce(0., lambda x, y: x + tf.reduce_sum(y))
@computations.federated_computation(federated_sequence_type)
def map_y_sum(x):
return intrinsics.federated_map(test_batch_select_and_reduce, x)
ds = tf.data.Dataset.from_tensor_slices({
'x': [[1., 2.], [3., 4.]],
'y': [[5.], [6.]]
}).batch(1)
self.assertEqual(map_y_sum([ds] * 5), [np.array([[11.]])] * 5)
def test_federated_aggregate_with_client_int(self):
# The representation used during the aggregation process will be a named
# tuple with 2 elements - the integer 'total' that represents the sum of
# elements encountered, and the integer element 'count'.
# pylint: disable=invalid-name
Accumulator = collections.namedtuple('Accumulator', 'total count')
# pylint: enable=invalid-name
accumulator_type = computation_types.NamedTupleType(
Accumulator(tf.int32, tf.int32))
# The operator to use during the first stage simply adds an element to the
# total and updates the count.
@computations.tf_computation([accumulator_type, tf.int32])
def accumulate(accu, elem):
return Accumulator(accu.total + elem, accu.count + 1)
# The operator to use during the second stage simply adds total and count.
@computations.tf_computation([accumulator_type, accumulator_type])
def merge(x, y):
return Accumulator(x.total + y.total, x.count + y.count)
# The operator to use during the final stage simply computes the ratio.
@computations.tf_computation(accumulator_type)
def report(accu):
return tf.cast(accu.total, tf.float32) / tf.cast(
accu.count, tf.float32)
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
val = intrinsics.federated_aggregate(x, Accumulator(0, 0),
accumulate, merge, report)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '({int32}@CLIENTS -> float32@SERVER)')
async def compute_intrinsic_federated_broadcast(
executor: executor_base.Executor, arg: executor_value_base.ExecutorValue
) -> executor_value_base.ExecutorValue:
"""Computes a federated broadcast on the given `executor`.
Args:
executor: The executor to use.
arg: The value to broadcast. Expected to be embedded in the `executor` and
have federated type placed at `tff.SERVER` with all_equal of `True`.
Returns:
The result embedded in `executor`.
Raises:
TypeError: If the arguments are of the wrong types.
"""
py_typecheck.check_type(executor, executor_base.Executor)
py_typecheck.check_type(arg, executor_value_base.ExecutorValue)
type_analysis.check_federated_type(
arg.type_signature, placement=placement_literals.SERVER, all_equal=True)
value = await arg.compute()
type_signature = computation_types.FederatedType(
arg.type_signature.member, placement_literals.CLIENTS, all_equal=True)
return await executor.create_value(value, type_signature)
def _create_stateless_int_dataset_reduction_iterative_process():
@computations.tf_computation()
def make_zero():
return tf.cast(0, tf.int64)
@computations.federated_computation()
def init():
return intrinsics.federated_eval(make_zero, placements.SERVER)
@computations.tf_computation(computation_types.SequenceType(tf.int64))
def reduce_dataset(x):
return x.reduce(tf.cast(0, tf.int64), lambda x, y: x + y)
@computations.federated_computation(
(init.type_signature.result,
computation_types.FederatedType(
computation_types.SequenceType(tf.int64), placements.CLIENTS)))
def next_fn(server_state, client_data):
del server_state # Unused
return intrinsics.federated_sum(
intrinsics.federated_map(reduce_dataset, client_data))
return iterative_process.IterativeProcess(initialize_fn=init,
next_fn=next_fn)
async def compute_intrinsic_federated_value(
executor: executor_base.Executor, arg: executor_value_base.ExecutorValue,
placement: placement_literals.PlacementLiteral
) -> executor_value_base.ExecutorValue:
"""Computes a federated value on the given `executor`.
Args:
executor: The executor to use.
arg: The value to place.
placement: The new placement of the value.
Returns:
The result embedded in `executor`.
Raises:
TypeError: If the arguments are of the wrong types.
"""
py_typecheck.check_type(executor, executor_base.Executor)
py_typecheck.check_type(arg, executor_value_base.ExecutorValue)
py_typecheck.check_type(placement, placement_literals.PlacementLiteral)
value = await arg.compute()
type_signature = computation_types.FederatedType(
arg.type_signature, placement, all_equal=True)
return await executor.create_value(value, type_signature)
def test_federated_mean_with_tuples(self):
@computations.federated_computation(
computation_types.FederatedType([('A', tf.float32),
('B', tf.float32)],
placements.CLIENTS))
def foo(x):
return intrinsics.federated_mean(x)
self.assertEqual(
str(foo.type_signature),
'({<A=float32,B=float32>}@CLIENTS -> <A=float32,B=float32>@SERVER)'
)
self.assertEqual(
str(
foo([{
'A': 1.0,
'B': 5.0
}, {
'A': 2.0,
'B': 6.0
}, {
'A': 3.0,
'B': 7.0
}])), '<A=2.0,B=6.0>')
async def _map(self, arg, all_equal=None):
py_typecheck.check_type(arg.internal_representation,
anonymous_tuple.AnonymousTuple)
py_typecheck.check_len(arg.internal_representation, 2)
fn_type = arg.type_signature[0]
py_typecheck.check_type(fn_type, computation_types.FunctionType)
val_type = arg.type_signature[1]
py_typecheck.check_type(val_type, computation_types.FederatedType)
if all_equal is None:
all_equal = val_type.all_equal
elif all_equal and not val_type.all_equal:
raise ValueError(
'Cannot map a non-all_equal argument into an all_equal result.')
fn = arg.internal_representation[0]
py_typecheck.check_type(fn, pb.Computation)
val = arg.internal_representation[1]
py_typecheck.check_type(val, list)
map_type = computation_types.FunctionType(
[fn_type, type_factory.at_clients(fn_type.parameter)],
type_factory.at_clients(fn_type.result))
map_comp = executor_utils.create_intrinsic_comp(
intrinsic_defs.FEDERATED_MAP, map_type)
async def _child_fn(ex, v):
py_typecheck.check_type(v, executor_value_base.ExecutorValue)
fn_val = await ex.create_value(fn, fn_type)
map_val, map_arg = tuple(await asyncio.gather(
ex.create_value(map_comp, map_type), ex.create_tuple([fn_val, v])))
return await ex.create_call(map_val, map_arg)
result_vals = await asyncio.gather(
*[_child_fn(c, v) for c, v in zip(self._child_executors, val)])
federated_type = computation_types.FederatedType(
fn_type.result, val_type.placement, all_equal=all_equal)
return CompositeValue(result_vals, federated_type)
def test_fails_with_bad_types(self):
function = computation_types.FunctionType(
None, computation_types.TensorType(tf.int32))
federated = computation_types.FederatedType(tf.int32,
placement_literals.CLIENTS)
tuple_on_function = computation_types.StructType([federated, function])
def foo(x): # pylint: disable=unused-variable
del x # Unused.
with self.assertRaisesRegex(
TypeError,
r'you have attempted to create one with the type {int32}@CLIENTS'):
computation_wrapper_instances.tensorflow_wrapper(foo, federated)
# pylint: disable=anomalous-backslash-in-string
with self.assertRaisesRegex(
TypeError,
r'you have attempted to create one with the type \( -> int32\)'):
computation_wrapper_instances.tensorflow_wrapper(foo, function)
with self.assertRaisesRegex(
TypeError, r'you have attempted to create one with the type placement'):
computation_wrapper_instances.tensorflow_wrapper(
foo, computation_types.PlacementType())
with self.assertRaisesRegex(
TypeError, r'you have attempted to create one with the type T'):
computation_wrapper_instances.tensorflow_wrapper(
foo, computation_types.AbstractType('T'))
with self.assertRaisesRegex(
TypeError,
r'you have attempted to create one with the type <{int32}@CLIENTS,\( '
'-> int32\)>'):
computation_wrapper_instances.tensorflow_wrapper(foo, tuple_on_function)
def test_returns_federated_aggregate(self):
value_type = computation_types.FederatedType(tf.int32,
placements.CLIENTS, False)
value = computation_building_blocks.Data('v', value_type)
zero = computation_building_blocks.Data('z', tf.int32)
accumulate_type = computation_types.NamedTupleType(
(tf.int32, tf.int32))
accumulate_result = computation_building_blocks.Data('a', tf.int32)
accumulate = computation_building_blocks.Lambda(
'x', accumulate_type, accumulate_result)
merge_type = computation_types.NamedTupleType((tf.int32, tf.int32))
merge_result = computation_building_blocks.Data('m', tf.int32)
merge = computation_building_blocks.Lambda('x', merge_type,
merge_result)
report_ref = computation_building_blocks.Reference('r', tf.int32)
report = computation_building_blocks.Lambda(report_ref.name,
report_ref.type_signature,
report_ref)
comp = computation_constructing_utils.create_federated_aggregate(
value, zero, accumulate, merge, report)
self.assertEqual(
comp.tff_repr,
'federated_aggregate(<v,z,(x -> a),(x -> m),(r -> r)>)')
self.assertEqual(str(comp.type_signature), 'int32@SERVER')
def test_replace_chained_federated_maps_replaces_federated_maps(self):
map_arg_type = computation_types.FederatedType(tf.int32,
placements.CLIENTS)
map_arg = computation_building_blocks.Reference('arg', map_arg_type)
inner_lambda = _create_lambda_to_add_one(map_arg.type_signature.member)
inner_call = _create_call_to_federated_map(inner_lambda, map_arg)
outer_lambda = _create_lambda_to_add_one(
inner_call.function.type_signature.result.member)
outer_call = _create_call_to_federated_map(outer_lambda, inner_call)
map_lambda = computation_building_blocks.Lambda(
map_arg.name, map_arg.type_signature, outer_call)
comp = map_lambda
uri = intrinsic_defs.FEDERATED_MAP.uri
self.assertEqual(_get_number_of_intrinsics(comp, uri), 2)
comp_impl = _to_comp(comp)
self.assertEqual(comp_impl([(1)]), [3])
transformed_comp = transformations.replace_chained_federated_maps_with_federated_map(
comp)
self.assertEqual(_get_number_of_intrinsics(transformed_comp, uri), 1)
transformed_comp_impl = _to_comp(transformed_comp)
self.assertEqual(transformed_comp_impl([(1)]), [3])
def test_handles_federated_broadcasts_nested_in_tuple(self):
first_broadcast = compiler_test_utils.create_whimsy_called_federated_broadcast(
)
packed_broadcast = building_blocks.Struct([
building_blocks.Data(
'a',
computation_types.FederatedType(
computation_types.TensorType(tf.int32), placements.SERVER)),
first_broadcast
])
sel = building_blocks.Selection(packed_broadcast, index=0)
second_broadcast = building_block_factory.create_federated_broadcast(sel)
result, _ = compiler_transformations.transform_to_call_dominant(
second_broadcast)
comp = building_blocks.Lambda('a', tf.int32, result)
uri = [intrinsic_defs.FEDERATED_BROADCAST.uri]
before, after = transformations.force_align_and_split_by_intrinsics(
comp, uri)
self.assertIsInstance(before, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(before, uri))
self.assertIsInstance(after, building_blocks.Lambda)
self.assertFalse(tree_analysis.contains_called_intrinsic(after, uri))
def test_two_tuple_zip_fails_bad_args(self):
server_test_ref = computation_building_blocks.Reference(
'test',
computation_types.NamedTupleType([
computation_types.FederatedType(tf.int32, placements.CLIENTS,
True),
computation_types.FederatedType(tf.bool, placements.SERVER,
True)
]))
with six.assertRaisesRegex(self, TypeError,
'should be placed at CLIENTS'):
_ = value_utils.zip_two_tuple(
value_impl.to_value(server_test_ref, None, _context_stack),
_context_stack)
client_test_ref = computation_building_blocks.Reference(
'test',
computation_types.NamedTupleType([
computation_types.FederatedType(tf.int32, placements.CLIENTS,
True),
computation_types.FederatedType(tf.bool, placements.CLIENTS,
True)
]))
with six.assertRaisesRegex(self, TypeError, '(Expected).*(Value)'):
_ = value_utils.zip_two_tuple(client_test_ref, _context_stack)
three_tuple_test_ref = computation_building_blocks.Reference(
'three_tuple_test',
computation_types.NamedTupleType([
computation_types.FederatedType(tf.int32, placements.CLIENTS,
True),
computation_types.FederatedType(tf.int32, placements.CLIENTS,
True),
computation_types.FederatedType(tf.int32, placements.CLIENTS,
True)
]))
with six.assertRaisesRegex(self, ValueError, 'must be a 2-tuple'):
_ = value_utils.zip_two_tuple(
value_impl.to_value(three_tuple_test_ref, None,
_context_stack), _context_stack)
class IntrinsicsTest(parameterized.TestCase):
def assert_type(self, value, type_string):
self.assertEqual(value.type_signature.compact_representation(),
type_string)
def test_constant_to_value_raises_outside_decorator(self):
with self.assertRaises(context_base.ContextError):
intrinsics.federated_value(2, placements.SERVER)
def test_intrinsic_construction_raises_context_error_outside_decorator(
self):
@computations.tf_computation()
def return_2():
return 2
with self.assertRaises(context_base.ContextError):
intrinsics.federated_eval(return_2, placements.SERVER)
def test_federated_broadcast_with_server_all_equal_int(self):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.SERVER))
def foo(x):
val = intrinsics.federated_broadcast(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(int32@SERVER -> int32@CLIENTS)')
def test_federated_broadcast_with_server_non_all_equal_int(self):
with self.assertRaises(TypeError):
@computations.federated_computation(
computation_types.FederatedType(tf.int32,
placements.SERVER,
all_equal=False))
def _(x):
return intrinsics.federated_broadcast(x)
def test_federated_broadcast_with_client_int(self):
with self.assertRaises(TypeError):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS,
True))
def _(x):
return intrinsics.federated_broadcast(x)
def test_federated_broadcast_with_non_federated_val(self):
with self.assertRaises(TypeError):
@computations.federated_computation(tf.int32)
def _(x):
return intrinsics.federated_broadcast(x)
def test_federated_eval_rand_on_clients(self):
@computations.federated_computation
def rand_on_clients():
@computations.tf_computation
def rand():
return tf.random.normal([])
val = intrinsics.federated_eval(rand, placements.CLIENTS)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(rand_on_clients, '( -> {float32}@CLIENTS)')
def test_federated_eval_rand_on_server(self):
@computations.federated_computation
def rand_on_server():
@computations.tf_computation
def rand():
return tf.random.normal([])
val = intrinsics.federated_eval(rand, placements.SERVER)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(rand_on_server, '( -> float32@SERVER)')
def test_federated_map_with_client_all_equal_int(self):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS,
True))
def foo(x):
val = intrinsics.federated_map(
computations.tf_computation(lambda x: x > 10), x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(int32@CLIENTS -> {bool}@CLIENTS)')
def test_federated_map_with_client_non_all_equal_int(self):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
val = intrinsics.federated_map(
computations.tf_computation(lambda x: x > 10), x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '({int32}@CLIENTS -> {bool}@CLIENTS)')
def test_federated_map_with_server_int(self):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.SERVER))
def foo(x):
val = intrinsics.federated_map(
computations.tf_computation(lambda x: x > 10), x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(int32@SERVER -> bool@SERVER)')
def test_federated_map_injected_zip_with_server_int(self):
@computations.federated_computation([
computation_types.FederatedType(tf.int32, placements.SERVER),
computation_types.FederatedType(tf.int32, placements.SERVER)
])
def foo(x, y):
val = intrinsics.federated_map(
computations.tf_computation(lambda x, y: x > 10,
), [x, y])
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(<int32@SERVER,int32@SERVER> -> bool@SERVER)')
def test_federated_map_injected_zip_fails_different_placements(self):
def foo(x, y):
val = intrinsics.federated_map(
computations.tf_computation(lambda x, y: x > 10,
), [x, y])
self.assertIsInstance(val, value_base.Value)
return val
with self.assertRaisesRegex(
TypeError,
'The value to be mapped must be a FederatedType or implicitly '
'convertible to a FederatedType.'):
computations.federated_computation(foo, [
computation_types.FederatedType(tf.int32, placements.SERVER),
computation_types.FederatedType(tf.int32, placements.CLIENTS)
])
def test_federated_map_with_non_federated_val(self):
with self.assertRaises(TypeError):
@computations.federated_computation(tf.int32)
def _(x):
return intrinsics.federated_map(
computations.tf_computation(lambda x: x > 10), x)
def test_federated_sum_with_client_int(self):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
val = intrinsics.federated_sum(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '({int32}@CLIENTS -> int32@SERVER)')
def test_federated_sum_with_client_string(self):
with self.assertRaises(TypeError):
@computations.federated_computation(
computation_types.FederatedType(tf.string, placements.CLIENTS))
def _(x):
return intrinsics.federated_sum(x)
def test_federated_sum_with_server_int(self):
with self.assertRaises(TypeError):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.SERVER))
def _(x):
return intrinsics.federated_sum(x)
def test_federated_zip_with_client_non_all_equal_int_and_bool(self):
@computations.federated_computation([
computation_types.FederatedType(tf.int32, placements.CLIENTS),
computation_types.FederatedType(tf.bool, placements.CLIENTS, True)
])
def foo(x, y):
val = intrinsics.federated_zip([x, y])
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(
foo, '(<{int32}@CLIENTS,bool@CLIENTS> -> {<int32,bool>}@CLIENTS)')
def test_federated_zip_with_single_unnamed_int_client(self):
@computations.federated_computation([
computation_types.FederatedType(tf.int32, placements.CLIENTS),
])
def foo(x):
val = intrinsics.federated_zip(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(<{int32}@CLIENTS> -> {<int32>}@CLIENTS)')
def test_federated_zip_with_single_unnamed_int_server(self):
@computations.federated_computation([
computation_types.FederatedType(tf.int32, placements.SERVER),
])
def foo(x):
val = intrinsics.federated_zip(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(<int32@SERVER> -> <int32>@SERVER)')
def test_federated_zip_with_single_named_bool_clients(self):
@computations.federated_computation([
('a', computation_types.FederatedType(tf.bool,
placements.CLIENTS)),
])
def foo(x):
val = intrinsics.federated_zip(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(<a={bool}@CLIENTS> -> {<a=bool>}@CLIENTS)')
def test_federated_zip_with_single_named_bool_server(self):
@computations.federated_computation([
('a', computation_types.FederatedType(tf.bool, placements.SERVER)),
])
def foo(x):
val = intrinsics.federated_zip(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(<a=bool@SERVER> -> <a=bool>@SERVER)')
def test_federated_zip_with_names_client_non_all_equal_int_and_bool(self):
@computations.federated_computation([
computation_types.FederatedType(tf.int32, placements.CLIENTS),
computation_types.FederatedType(tf.bool, placements.CLIENTS, True)
])
def foo(x, y):
a = {'x': x, 'y': y}
val = intrinsics.federated_zip(a)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(
foo,
'(<{int32}@CLIENTS,bool@CLIENTS> -> {<x=int32,y=bool>}@CLIENTS)')
def test_federated_zip_with_client_all_equal_int_and_bool(self):
@computations.federated_computation([
computation_types.FederatedType(tf.int32, placements.CLIENTS,
True),
computation_types.FederatedType(tf.bool, placements.CLIENTS, True)
])
def foo(x, y):
val = intrinsics.federated_zip([x, y])
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(
foo, '(<int32@CLIENTS,bool@CLIENTS> -> {<int32,bool>}@CLIENTS)')
def test_federated_zip_with_names_client_all_equal_int_and_bool(self):
@computations.federated_computation([
computation_types.FederatedType(tf.int32, placements.CLIENTS,
True),
computation_types.FederatedType(tf.bool, placements.CLIENTS, True)
])
def foo(arg):
a = {'x': arg[0], 'y': arg[1]}
val = intrinsics.federated_zip(a)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(
foo,
'(<int32@CLIENTS,bool@CLIENTS> -> {<x=int32,y=bool>}@CLIENTS)')
def test_federated_zip_with_server_int_and_bool(self):
@computations.federated_computation([
computation_types.FederatedType(tf.int32, placements.SERVER),
computation_types.FederatedType(tf.bool, placements.SERVER)
])
def foo(x, y):
val = intrinsics.federated_zip([x, y])
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(
foo, '(<int32@SERVER,bool@SERVER> -> <int32,bool>@SERVER)')
def test_federated_zip_with_names_server_int_and_bool(self):
@computations.federated_computation([
('a', computation_types.FederatedType(tf.int32,
placements.SERVER)),
('b', computation_types.FederatedType(tf.bool, placements.SERVER)),
])
def foo(arg):
val = intrinsics.federated_zip(arg)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(
foo, '(<a=int32@SERVER,b=bool@SERVER> -> <a=int32,b=bool>@SERVER)')
def test_federated_zip_error_different_placements(self):
with self.assertRaises(TypeError):
@computations.federated_computation([
('a',
computation_types.FederatedType(tf.int32, placements.SERVER)),
('b',
computation_types.FederatedType(tf.bool, placements.CLIENTS)),
])
def _(arg):
return intrinsics.federated_zip(arg)
def test_federated_collect_with_client_int(self):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
val = intrinsics.federated_collect(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '({int32}@CLIENTS -> int32*@SERVER)')
def test_federated_collect_with_server_int_fails(self):
with self.assertRaises(TypeError):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.SERVER))
def _(x):
return intrinsics.federated_collect(x)
def test_federated_mean_with_client_float32_without_weight(self):
@computations.federated_computation(
computation_types.FederatedType(tf.float32, placements.CLIENTS))
def foo(x):
val = intrinsics.federated_mean(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '({float32}@CLIENTS -> float32@SERVER)')
def test_federated_mean_with_all_equal_client_float32_without_weight(self):
federated_all_equal_float = computation_types.FederatedType(
tf.float32, placements.CLIENTS, all_equal=True)
@computations.federated_computation(federated_all_equal_float)
def foo(x):
val = intrinsics.federated_mean(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(float32@CLIENTS -> float32@SERVER)')
def test_federated_mean_with_all_equal_client_float32_with_weight(self):
federated_all_equal_float = computation_types.FederatedType(
tf.float32, placements.CLIENTS, all_equal=True)
@computations.federated_computation(federated_all_equal_float)
def foo(x):
val = intrinsics.federated_mean(x, x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(float32@CLIENTS -> float32@SERVER)')
def test_federated_mean_with_client_tuple_with_int32_weight(self):
@computations.federated_computation([
computation_types.FederatedType([('x', tf.float64),
('y', tf.float64)],
placements.CLIENTS),
computation_types.FederatedType(tf.int32, placements.CLIENTS)
])
def foo(x, y):
val = intrinsics.federated_mean(x, y)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(
foo, '(<{<x=float64,y=float64>}@CLIENTS,{int32}@CLIENTS> '
'-> <x=float64,y=float64>@SERVER)')
def test_federated_mean_with_client_int32_fails(self):
with self.assertRaises(TypeError):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def _(x):
return intrinsics.federated_mean(x)
def test_federated_mean_with_string_weight_fails(self):
with self.assertRaises(TypeError):
@computations.federated_computation([
computation_types.FederatedType(tf.float32,
placements.CLIENTS),
computation_types.FederatedType(tf.string, placements.CLIENTS)
])
def _(x, y):
return intrinsics.federated_mean(x, y)
def test_federated_aggregate_with_client_int(self):
# The representation used during the aggregation process will be a named
# tuple with 2 elements - the integer 'total' that represents the sum of
# elements encountered, and the integer element 'count'.
# pylint: disable=invalid-name
Accumulator = collections.namedtuple('Accumulator', 'total count')
# pylint: enable=invalid-name
# The operator to use during the first stage simply adds an element to the
# total and updates the count.
@computations.tf_computation
def accumulate(accu, elem):
return Accumulator(accu.total + elem, accu.count + 1)
# The operator to use during the second stage simply adds total and count.
@computations.tf_computation
def merge(x, y):
return Accumulator(x.total + y.total, x.count + y.count)
# The operator to use during the final stage simply computes the ratio.
@computations.tf_computation
def report(accu):
return tf.cast(accu.total, tf.float32) / tf.cast(
accu.count, tf.float32)
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
val = intrinsics.federated_aggregate(x, Accumulator(0, 0),
accumulate, merge, report)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '({int32}@CLIENTS -> float32@SERVER)')
def test_federated_aggregate_with_federated_zero_fails(self):
@computations.federated_computation()
def build_federated_zero():
val = intrinsics.federated_value(0, placements.SERVER)
self.assertIsInstance(val, value_base.Value)
return val
@computations.tf_computation([tf.int32, tf.int32])
def accumulate(accu, elem):
return accu + elem
# The operator to use during the second stage simply adds total and count.
@computations.tf_computation([tf.int32, tf.int32])
def merge(x, y):
return x + y
# The operator to use during the final stage simply computes the ratio.
@computations.tf_computation(tf.int32)
def report(accu):
return accu
def foo(x):
return intrinsics.federated_aggregate(x, build_federated_zero(),
accumulate, merge, report)
with self.assertRaisesRegex(
TypeError, 'Expected `zero` to be assignable to type int32, '
'but was of incompatible type int32@SERVER'):
computations.federated_computation(
foo,
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def test_federated_aggregate_with_unknown_dimension(self):
Accumulator = collections.namedtuple('Accumulator', ['samples']) # pylint: disable=invalid-name
accumulator_type = computation_types.StructType(
Accumulator(samples=computation_types.TensorType(dtype=tf.int32,
shape=[None])))
@computations.tf_computation()
def build_empty_accumulator():
return Accumulator(samples=tf.zeros(shape=[0], dtype=tf.int32))
# The operator to use during the first stage simply adds an element to the
# tensor, increasing its size.
@computations.tf_computation([accumulator_type, tf.int32])
def accumulate(accu, elem):
return Accumulator(samples=tf.concat(
[accu.samples, tf.expand_dims(elem, axis=0)], axis=0))
# The operator to use during the second stage simply adds total and count.
@computations.tf_computation([accumulator_type, accumulator_type])
def merge(x, y):
return Accumulator(
samples=tf.concat([x.samples, y.samples], axis=0))
# The operator to use during the final stage simply computes the ratio.
@computations.tf_computation(accumulator_type)
def report(accu):
return accu
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
val = intrinsics.federated_aggregate(x, build_empty_accumulator(),
accumulate, merge, report)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '({int32}@CLIENTS -> <samples=int32[?]>@SERVER)')
def test_federated_reduce_with_tf_add_raw_constant(self):
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.CLIENTS))
def foo(x):
plus = computations.tf_computation(tf.add)
val = intrinsics.federated_reduce(x, 0, plus)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '({int32}@CLIENTS -> int32@SERVER)')
def test_num_over_temperature_threshold_example(self):
@computations.federated_computation([
computation_types.FederatedType(tf.float32, placements.CLIENTS),
computation_types.FederatedType(tf.float32, placements.SERVER)
])
def foo(temperatures, threshold):
val = intrinsics.federated_sum(
intrinsics.federated_map(
computations.tf_computation(
lambda x, y: tf.cast(tf.greater(x, y), tf.int32)),
[temperatures,
intrinsics.federated_broadcast(threshold)]))
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(
foo, '(<{float32}@CLIENTS,float32@SERVER> -> int32@SERVER)')
@parameterized.named_parameters(('test_n_2', 2), ('test_n_3', 3),
('test_n_5', 5))
def test_n_tuple_federated_zip_tensor_args(self, n):
fed_type = computation_types.FederatedType(tf.int32,
placements.CLIENTS)
initial_tuple_type = computation_types.StructType([fed_type] * n)
final_fed_type = computation_types.FederatedType([tf.int32] * n,
placements.CLIENTS)
function_type = computation_types.FunctionType(initial_tuple_type,
final_fed_type)
@computations.federated_computation(
[computation_types.FederatedType(tf.int32, placements.CLIENTS)] * n
)
def foo(x):
val = intrinsics.federated_zip(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, function_type.compact_representation())
@parameterized.named_parameters(
('test_n_2_int', 2,
computation_types.FederatedType(tf.int32, placements.CLIENTS)),
('test_n_3_int', 3,
computation_types.FederatedType(tf.int32, placements.CLIENTS)),
('test_n_5_int', 5,
computation_types.FederatedType(tf.int32, placements.CLIENTS)),
('test_n_2_tuple', 2,
computation_types.FederatedType([tf.int32, tf.int32],
placements.CLIENTS)),
('test_n_3_tuple', 3,
computation_types.FederatedType([tf.int32, tf.int32],
placements.CLIENTS)),
('test_n_5_tuple', 5,
computation_types.FederatedType([tf.int32, tf.int32],
placements.CLIENTS)))
def test_named_n_tuple_federated_zip(self, n, fed_type):
initial_tuple_type = computation_types.StructType([fed_type] * n)
named_fed_type = computation_types.FederatedType(
[(str(k), fed_type.member) for k in range(n)], placements.CLIENTS)
mixed_fed_type = computation_types.FederatedType(
[(str(k), fed_type.member) if k % 2 == 0 else fed_type.member
for k in range(n)], placements.CLIENTS)
named_function_type = computation_types.FunctionType(
initial_tuple_type, named_fed_type)
mixed_function_type = computation_types.FunctionType(
initial_tuple_type, mixed_fed_type)
@computations.federated_computation([fed_type] * n)
def foo(x):
arg = {str(k): x[k] for k in range(n)}
val = intrinsics.federated_zip(arg)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, named_function_type.compact_representation())
def _make_test_tuple(x, k):
"""Make a test tuple with a name if k is even, otherwise unnamed."""
if k % 2 == 0:
return str(k), x[k]
else:
return None, x[k]
@computations.federated_computation([fed_type] * n)
def bar(x):
arg = structure.Struct(_make_test_tuple(x, k) for k in range(n))
val = intrinsics.federated_zip(arg)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(bar, mixed_function_type.compact_representation())
@parameterized.named_parameters([
('test_n_' + str(n) + '_m_' + str(m), n, m)
for n, m in itertools.product([1, 2, 3], [1, 2, 3])
])
def test_n_tuple_federated_zip_mixed_args(self, n, m):
tuple_fed_type = computation_types.FederatedType([tf.int32, tf.int32],
placements.CLIENTS)
single_fed_type = computation_types.FederatedType(
tf.int32, placements.CLIENTS)
initial_tuple_type = computation_types.StructType(
[tuple_fed_type] * n + [single_fed_type] * m)
final_fed_type = computation_types.FederatedType(
[[tf.int32, tf.int32]] * n + [tf.int32] * m, placements.CLIENTS)
function_type = computation_types.FunctionType(initial_tuple_type,
final_fed_type)
@computations.federated_computation(
[
computation_types.FederatedType(
computation_types.StructType([tf.int32, tf.int32]),
placements.CLIENTS)
] * n +
[computation_types.FederatedType(tf.int32, placements.CLIENTS)] * m
)
def baz(x):
val = intrinsics.federated_zip(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(baz, function_type.compact_representation())
def test_federated_apply_raises_warning(self):
with warnings.catch_warnings(record=True) as w:
warnings.simplefilter('always')
@computations.federated_computation(
computation_types.FederatedType(tf.int32, placements.SERVER))
def foo(x):
val = intrinsics.federated_apply(
computations.tf_computation(lambda x: x * x), x)
self.assertIsInstance(val, value_base.Value)
return val
self.assertLen(w, 1)
self.assertIsInstance(w[0].category(), DeprecationWarning)
self.assertIn('tff.federated_apply() is deprecated',
str(w[0].message))
self.assert_type(foo, '(int32@SERVER -> int32@SERVER)')
def test_federated_value_with_bool_on_clients(self):
@computations.federated_computation(tf.bool)
def foo(x):
val = intrinsics.federated_value(x, placements.CLIENTS)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(bool -> bool@CLIENTS)')
def test_federated_value_raw_np_scalar(self):
@computations.federated_computation
def foo():
floatv = np.float64(0)
tff_float = intrinsics.federated_value(floatv, placements.SERVER)
self.assertIsInstance(tff_float, value_base.Value)
self.assert_type(tff_float, 'float64@SERVER')
intv = np.int64(0)
tff_int = intrinsics.federated_value(intv, placements.SERVER)
self.assertIsInstance(tff_int, value_base.Value)
self.assert_type(tff_int, 'int64@SERVER')
return (tff_float, tff_int)
self.assert_type(foo, '( -> <float64@SERVER,int64@SERVER>)')
def test_federated_value_raw_tf_scalar_variable(self):
v = tf.Variable(initial_value=0., name='test_var')
with self.assertRaisesRegex(
TypeError, 'TensorFlow construct (.*) has been '
'encountered in a federated context.'):
@computations.federated_computation()
def _():
return intrinsics.federated_value(v, placements.SERVER)
def test_federated_value_with_bool_on_server(self):
@computations.federated_computation(tf.bool)
def foo(x):
val = intrinsics.federated_value(x, placements.SERVER)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo, '(bool -> bool@SERVER)')
def test_sequence_sum(self):
@computations.federated_computation(
computation_types.SequenceType(tf.int32))
def foo1(x):
val = intrinsics.sequence_sum(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo1, '(int32* -> int32)')
@computations.federated_computation(
computation_types.FederatedType(
computation_types.SequenceType(tf.int32), placements.SERVER))
def foo2(x):
val = intrinsics.sequence_sum(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo2, '(int32*@SERVER -> int32@SERVER)')
@computations.federated_computation(
computation_types.FederatedType(
computation_types.SequenceType(tf.int32), placements.CLIENTS))
def foo3(x):
val = intrinsics.sequence_sum(x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo3, '({int32*}@CLIENTS -> {int32}@CLIENTS)')
def test_sequence_map(self):
@computations.tf_computation(tf.int32)
def over_threshold(x):
return x > 10
@computations.federated_computation(
computation_types.SequenceType(tf.int32))
def foo1(x):
val = intrinsics.sequence_map(over_threshold, x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo1, '(int32* -> bool*)')
@computations.federated_computation(
computation_types.FederatedType(
computation_types.SequenceType(tf.int32), placements.SERVER))
def foo2(x):
val = intrinsics.sequence_map(over_threshold, x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo2, '(int32*@SERVER -> bool*@SERVER)')
@computations.federated_computation(
computation_types.FederatedType(
computation_types.SequenceType(tf.int32), placements.CLIENTS))
def foo3(x):
val = intrinsics.sequence_map(over_threshold, x)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo3, '({int32*}@CLIENTS -> {bool*}@CLIENTS)')
def test_sequence_reduce(self):
add_numbers = computations.tf_computation(tf.add, [tf.int32, tf.int32])
@computations.federated_computation(
computation_types.SequenceType(tf.int32))
def foo1(x):
val = intrinsics.sequence_reduce(x, 0, add_numbers)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo1, '(int32* -> int32)')
@computations.federated_computation(
computation_types.FederatedType(
computation_types.SequenceType(tf.int32), placements.SERVER))
def foo2(x):
val = intrinsics.sequence_reduce(x, 0, add_numbers)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo2, '(int32*@SERVER -> int32@SERVER)')
@computations.federated_computation(
computation_types.FederatedType(
computation_types.SequenceType(tf.int32), placements.CLIENTS))
def foo3(x):
val = intrinsics.sequence_reduce(x, 0, add_numbers)
self.assertIsInstance(val, value_base.Value)
return val
self.assert_type(foo3, '({int32*}@CLIENTS -> {int32}@CLIENTS)')
def foo():
return intrinsic_utils.zero_for(
computation_types.FederatedType(tf.int32, placements.SERVER, True),
context_stack_impl.context_stack)
@computations.tf_computation(tf.string)
def float_dataset_computation(x):
del x # Unused
return tf.data.Dataset.range(5, output_type=tf.float32)
@computations.tf_computation(tf.int32)
def int_identity(x):
return x
@computations.federated_computation(
tf.int32,
computation_types.FederatedType(computation_types.SequenceType(tf.int64),
placements.CLIENTS),
tf.float32,
)
def test_int64_sequence_struct_computation(a, dataset, b):
return a, dataset, b
@computations.federated_computation(
computation_types.FederatedType(computation_types.SequenceType(tf.int64),
placements.CLIENTS))
def test_int64_sequence_computation(dataset):
del dataset
return intrinsics.federated_value(5, placements.SERVER)
class ConstructDatasetsOnClientsComputationTest(absltest.TestCase):
