def _get_accumulator_type(member_type):
"""Constructs a `tff.Type` for the accumulator in sample aggregation.
Args:
member_type: A `tff.Type` representing the member components of the
federated type.
Returns:
The `tff.NamedTupleType` associated with the accumulator. The tuple contains
two parts, `accumulators` and `rands`, that are parallel lists (e.g. the
i-th index in one corresponds to the i-th index in the other). These two
lists are used to sample from the accumulators with equal probability.
"""
# TODO(b/121288403): Special-casing anonymous tuple shouldn't be needed.
if isinstance(member_type, tff.NamedTupleType):
a = anonymous_tuple.map_structure(
lambda v: tff.TensorType(v.dtype, [None] + v.shape.dims),
member_type)
return tff.NamedTupleType(
collections.OrderedDict({
'accumulators':
tff.NamedTupleType(anonymous_tuple.to_odict(a, True)),
'rands':
tff.TensorType(tf.float32, shape=[None])
}))
return tff.NamedTupleType(
collections.OrderedDict({
'accumulators':
tff.TensorType(member_type.dtype,
shape=[None] + member_type.shape.dims),
'rands':
tff.TensorType(tf.float32, shape=[None])
}))
def test_federated_sample_on_nested_scalars(self):
tuple_type = tff.NamedTupleType([
('x', tf.float32),
('y', tf.float32),
])
@tff.federated_computation(tff.FederatedType(tuple_type, tff.CLIENTS))
def call_federated_sample(value):
return federated_aggregations.federated_sample(value)
x0 = 0.0
y0 = 1.0
x1 = -1.0
y1 = 5.0
test_type = collections.namedtuple('NestedScalars', ['x', 'y'])
value = call_federated_sample(
[test_type(x0, y0),
test_type(x1, y1),
test_type(2.0, -10.0)])
result = value._asdict()
i0 = list(result['x']).index(x0)
i1 = list(result['y']).index(y1)
# Assert shuffled in unison.
self.assertEqual(result['y'][i0], y0)
self.assertEqual(result['x'][i1], x1)
def test_named_n_tuple_federated_zip(self, n, fed_type):
initial_tuple_type = tff.NamedTupleType([fed_type] * n)
named_fed_type = tff.FederatedType(
[(str(k), fed_type.member) for k in range(n)], tff.CLIENTS)
mixed_fed_type = tff.FederatedType(
[(str(k), fed_type.member) if k % 2 == 0 else fed_type.member
for k in range(n)], tff.CLIENTS)
named_function_type = tff.FunctionType(initial_tuple_type, named_fed_type)
mixed_function_type = tff.FunctionType(initial_tuple_type, mixed_fed_type)
named_type_string = str(named_function_type)
mixed_type_string = str(mixed_function_type)
@tff.federated_computation([fed_type] * n)
def foo(x):
arg = {str(k): x[k] for k in range(n)}
return tff.federated_zip(arg)
self.assertEqual(str(foo.type_signature), named_type_string)
@tff.federated_computation([fed_type] * n)
def bar(x):
arg = anonymous_tuple.AnonymousTuple([
(str(k), x[k]) if k % 2 == 0 else (None, x[k]) for k in range(n)
])
return tff.federated_zip(arg)
self.assertEqual(str(bar.type_signature), mixed_type_string)
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 = tff.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.
@tff.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.
@tff.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.
@tff.tf_computation(accumulator_type)
def report(accu):
return tf.to_float(accu.total) / tf.to_float(accu.count)
@tff.federated_computation(tff.FederatedType(tf.int32, tff.CLIENTS))
def foo(x):
return tff.federated_aggregate(x, Accumulator(0, 0), accumulate, merge,
report)
self.assertEqual(
str(foo.type_signature), '({int32}@CLIENTS -> float32@SERVER)')
def test_federated_aggregate_with_unknown_dimension(self):
Accumulator = collections.namedtuple('Accumulator', ['samples']) # pylint: disable=invalid-name
accumulator_type = tff.NamedTupleType(
Accumulator(samples=tff.TensorType(dtype=tf.int32, shape=[None])))
@tff.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.
@tff.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.
@tff.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.
@tff.tf_computation(accumulator_type)
def report(accu):
return accu
@tff.federated_computation(tff.FederatedType(tf.int32, tff.CLIENTS))
def foo(x):
return tff.federated_aggregate(x, build_empty_accumulator(),
accumulate, merge, report)
self.assertEqual(str(foo.type_signature),
'({int32}@CLIENTS -> <samples=int32[?]>@SERVER)')
def test_named_n_tuple_federated_zip(self, n, fed_type):
initial_tuple_type = tff.NamedTupleType([fed_type] * n)
named_fed_type = tff.FederatedType([(str(k), fed_type.member)
for k in range(n)], tff.CLIENTS)
mixed_fed_type = tff.FederatedType(
[(str(k), fed_type.member) if k % 2 == 0 else fed_type.member
for k in range(n)], tff.CLIENTS)
named_function_type = tff.FunctionType(initial_tuple_type,
named_fed_type)
mixed_function_type = tff.FunctionType(initial_tuple_type,
mixed_fed_type)
named_type_string = str(named_function_type)
mixed_type_string = str(mixed_function_type)
@tff.federated_computation([fed_type] * n)
def foo(x):
arg = {str(k): x[k] for k in range(n)}
return tff.federated_zip(arg)
self.assertEqual(str(foo.type_signature), named_type_string)
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]
@tff.federated_computation([fed_type] * n)
def bar(x):
arg = anonymous_tuple.AnonymousTuple(
_make_test_tuple(x, k) for k in range(n))
return tff.federated_zip(arg)
self.assertEqual(str(bar.type_signature), mixed_type_string)
def test_n_tuple_federated_zip_mixed_args(self, n, m):
tuple_fed_type = tff.FederatedType([tf.int32, tf.int32], tff.CLIENTS)
single_fed_type = tff.FederatedType(tf.int32, tff.CLIENTS)
initial_tuple_type = tff.NamedTupleType([tuple_fed_type] * n +
[single_fed_type] * m)
final_fed_type = tff.FederatedType([[tf.int32, tf.int32]] * n +
[tf.int32] * m, tff.CLIENTS)
function_type = tff.FunctionType(initial_tuple_type, final_fed_type)
type_string = str(function_type)
@tff.federated_computation([
tff.FederatedType(
tff.NamedTupleType([tf.int32, tf.int32]), tff.CLIENTS)
] * n + [tff.FederatedType(tf.int32, tff.CLIENTS)] * m)
def baz(x):
return tff.federated_zip(x)
self.assertEqual(str(baz.type_signature), type_string)
def test_n_tuple_federated_zip_tensor_args(self, n):
fed_type = tff.FederatedType(tf.int32, tff.CLIENTS)
initial_tuple_type = tff.NamedTupleType([fed_type] * n)
final_fed_type = tff.FederatedType([tf.int32] * n, tff.CLIENTS)
function_type = tff.FunctionType(initial_tuple_type, final_fed_type)
type_string = str(function_type)
@tff.federated_computation([tff.FederatedType(tf.int32, tff.CLIENTS)] * n)
def foo(x):
return tff.federated_zip(x)
self.assertEqual(str(foo.type_signature), type_string)
def test_federated_max_on_nested_scalars(self):
tuple_type = tff.NamedTupleType([
('a', tf.int32),
('b', tf.int32),
])
@tff.federated_computation(tff.FederatedType(tuple_type, tff.CLIENTS))
def call_federated_max(value):
return federated_aggregations.federated_max(value)
test_type = collections.namedtuple('NestedScalars', ['a', 'b'])
value = call_federated_max(
[test_type(1, 5), test_type(2, 3),
test_type(1, 8)])
self.assertDictEqual(value._asdict(), {'a': 2, 'b': 8})
def test_federated_min_on_nested_scalars(self):
tuple_type = tff.NamedTupleType([
('x', tf.float32),
('y', tf.float32),
])
@tff.federated_computation(tff.FederatedType(tuple_type, tff.CLIENTS))
def call_federated_min(value):
return federated_aggregations.federated_min(value)
test_type = collections.namedtuple('NestedScalars', ['x', 'y'])
value = call_federated_min(
[test_type(0.0, 1.0),
test_type(-1.0, 5.0),
test_type(2.0, -10.0)])
self.assertDictEqual(value._asdict(), {'x': -1.0, 'y': -10.0})
def test_federated_max_nested_tensor_value(self):
tuple_type = tff.NamedTupleType([
('a', (tf.int32, [2])),
('b', (tf.int32, [3])),
])
@tff.federated_computation(tff.FederatedType(tuple_type, tff.CLIENTS))
def call_federated_max(value):
return federated_aggregations.federated_max(value)
test_type = collections.namedtuple('NestedScalars', ['a', 'b'])
client1 = test_type(
np.array([4, 5], dtype=np.int32), np.array([1, -2, 3], dtype=np.int32))
client2 = test_type(
np.array([9, 0], dtype=np.int32), np.array([5, 1, -2], dtype=np.int32))
value = call_federated_max([client1, client2])
self.assertCountEqual(value[0], [9, 5])
self.assertCountEqual(value[1], [5, 1, 3])
def _split_by_intrinsic(comp, uri):
"""Splits `comp` into `before` and `after` the intrinsic for the given `uri`.
This function finds the intrinsic for the given `uri` in `comp`; splits `comp`
into two computations `before` and `after` the intrinsic; and returns a Python
tuple representing the pair of `before` and `after` computations.
NOTE: This function is generally safe to call on computations that do not fit
into canonical form. It is left to the caller to determine if the resulting
computations are expected.
Args:
comp: The `tff_framework.Lambda` to transform.
uri: A URI of an intrinsic.
Returns:
A pair of `tff_framework.ComputationBuildingBlock`s representing the
computations `before` and `after` the intrinsic.
Raises:
ValueError: If `comp` is not a `tff_framework.Lambda` referencing a
`tff_framework.Block` referencing a collections of variables containing an
intrinsic with the given `uri` or if there is more than one intrinsic with
the given `uri` in `comp`.
"""
py_typecheck.check_type(comp, tff_framework.Lambda)
py_typecheck.check_type(uri, six.string_types)
py_typecheck.check_type(comp.result, tff_framework.Block)
def _get_called_intrinsic_from_block_variables(variables, uri):
for index, (name, variable) in enumerate(variables):
if tff_framework.is_called_intrinsic(variable, uri):
return index, name, variable
raise ValueError(
'Expected a lambda referencing a block referencing a collection of '
'variables containing an intrinsic with the uri: {}, found None.'
.format(uri))
index, name, variable = _get_called_intrinsic_from_block_variables(
comp.result.locals, uri)
intrinsics = _get_called_intrinsics(comp, uri)
length = len(intrinsics)
if length != 1:
raise ValueError(
'Expected a computation with exactly one intrinsic with the uri: {}, '
'found: {}.'.format(uri, length))
name_generator = tff_framework.unique_name_generator(comp)
before = tff_framework.Lambda(comp.parameter_name, comp.parameter_type,
variable.argument)
parameter_type = tff.NamedTupleType(
(comp.parameter_type, variable.type_signature))
ref_name = six.next(name_generator)
ref = tff_framework.Reference(ref_name, parameter_type)
sel_0 = tff_framework.Selection(ref, index=0)
sel_1 = tff_framework.Selection(ref, index=1)
variables = comp.result.locals
variables[index] = (name, sel_1)
variables.insert(0, (comp.parameter_name, sel_0))
block = tff_framework.Block(variables, comp.result.result)
after = tff_framework.Lambda(ref.name, ref.type_signature, block)
return before, after
