def build_encoded_mean(values, encoders):
"""Builds `StatefulAggregateFn` for `values`, to be encoded by `encoders`.
Args:
values: Values to be encoded by the `StatefulAggregateFn`. Must be
convertible to `tff.Value`.
encoders: A collection of `GatherEncoder` objects to be used for encoding
`values`. Must have the same structure as `values`.
Returns:
A `StatefulAggregateFn` of which `next_fn` encodes the input at
`tff.CLIENTS`, and computes their mean at `tff.SERVER`, automatically
splitting the decoding part based on its commutativity with sum.
Raises:
ValueError: If `values` and `encoders` do not have the same structure.
TypeError: If `encoders` are not instances of `GatherEncoder`, or if
`values` are not compatible with the expected input of the `encoders`.
"""
tf.nest.assert_same_structure(values, encoders)
tf.nest.map_structure(
lambda e, v: _validate_encoder(e, v, te.core.GatherEncoder), encoders,
values)
value_type = tff_framework.type_from_tensors(values)
initial_state_fn = _build_initial_state_tf_computation(encoders)
state_type = initial_state_fn.type_signature.result
nest_encoder = _build_tf_computations_for_gather(state_type, value_type,
encoders)
encoded_sum_fn = _build_encoded_sum_fn(nest_encoder)
@tff.tf_computation(value_type, tff.to_type(tf.float32))
def multiply_fn(value, weight):
return tf.nest.map_structure(lambda v: v * tf.cast(weight, v.dtype),
value)
@tff.tf_computation(value_type, tff.to_type(tf.float32))
def divide_fn(value, denominator):
return tf.nest.map_structure(
lambda v: v / tf.cast(denominator, v.dtype), value)
def encoded_mean_fn(state, values, weight):
weighted_values = tff.federated_map(multiply_fn, [values, weight])
updated_state, summed_decoded_values = encoded_sum_fn(
state, weighted_values)
summed_weights = tff.federated_sum(weight)
decoded_values = tff.federated_map(
divide_fn, [summed_decoded_values, summed_weights])
return updated_state, decoded_values
return StatefulAggregateFn(initialize_fn=initial_state_fn,
next_fn=encoded_mean_fn)
def build_encoded_broadcast(values, encoders):
"""Builds `StatefulBroadcastFn` for `values`, to be encoded by `encoders`.
Args:
values: Values to be broadcasted by the `StatefulBroadcastFn`. Must be
convertible to `tff.Value`.
encoders: A collection of `SimpleEncoder` objects to be used for encoding
`values`. Must have the same structure as `values`.
Returns:
A `StatefulBroadcastFn` of which `next_fn` encodes the input at
`tff.SERVER`, broadcasts the encoded representation and decodes the encoded
representation at `tff.CLIENTS`.
Raises:
ValueError: If `values` and `encoders` do not have the same structure.
TypeError: If `encoders` are not instances of `SimpleEncoder`, or if
`values` are not compatible with the expected input of the `encoders`.
"""
def validate_encoder(encoder, value):
if not isinstance(encoder, te.core.SimpleEncoder):
raise TypeError(
'Provided encoder must be an instance of SimpleEncoder.')
if not encoder.input_tensorspec.is_compatible_with(
tf.TensorSpec(value.shape, value.dtype)):
raise TypeError('Provided encoder and value are not compatible.')
tf.nest.assert_same_structure(values, encoders)
tf.nest.map_structure(validate_encoder, encoders, values)
initial_state_fn = _build_initial_state_tf_computation(encoders)
fed_state_type = tff.FederatedType(initial_state_fn.type_signature.result,
tff.SERVER)
fed_value_type = tff.FederatedType(tff_framework.type_from_tensors(values),
tff.SERVER)
@tff.federated_computation(fed_state_type, fed_value_type)
def encoded_broadcast_fn(state, value):
"""Broadcast function, to be wrapped as federated_computation."""
state_type = state.type_signature.member
value_type = value.type_signature.member
encode_fn, decode_fn = _build_encode_decode_tf_computations_for_broadcast(
state_type, value_type, encoders)
new_state, encoded_value = tff.federated_apply(encode_fn,
(state, value))
client_encoded_value = tff.federated_broadcast(encoded_value)
client_value = tff.federated_map(decode_fn, client_encoded_value)
return new_state, client_value
return StatefulBroadcastFn(initialize_fn=initial_state_fn,
next_fn=encoded_broadcast_fn)
def wrap_aggregate_fn(dp_aggregate_fn, sample_value):
tff_types = tff_framework.type_from_tensors(sample_value)
@tff.federated_computation
def run_initialize():
return tff.federated_value(dp_aggregate_fn.initialize(), tff.SERVER)
@tff.federated_computation(run_initialize.type_signature.result,
tff.FederatedType(tff_types, tff.CLIENTS))
def run_aggregate(global_state, client_values):
return dp_aggregate_fn(global_state, client_values)
return run_initialize, run_aggregate
def build_encoded_broadcast(values, encoders):
"""Builds `StatefulBroadcastFn` for `values`, to be encoded by `encoders`.
Args:
values: Values to be broadcasted by the `StatefulBroadcastFn`. Must be
convertible to `tff.Value`.
encoders: A collection of `SimpleEncoder` objects to be used for encoding
`values`. Must have the same structure as `values`.
Returns:
A `StatefulBroadcastFn` of which `next_fn` encodes the input at
`tff.SERVER`, broadcasts the encoded representation and decodes the encoded
representation at `tff.CLIENTS`.
Raises:
ValueError: If `values` and `encoders` do not have the same structure.
TypeError: If `encoders` are not instances of `SimpleEncoder`, or if
`values` are not compatible with the expected input of the `encoders`.
"""
tf.nest.assert_same_structure(values, encoders)
tf.nest.map_structure(
lambda e, v: _validate_encoder(e, v, te.core.SimpleEncoder), encoders,
values)
value_type = tff_framework.type_from_tensors(values)
initial_state_fn = _build_initial_state_tf_computation(encoders)
state_type = initial_state_fn.type_signature.result
encode_fn, decode_fn = _build_encode_decode_tf_computations_for_broadcast(
state_type, value_type, encoders)
def encoded_broadcast_fn(state, value):
"""Encoded broadcast federated_computation."""
new_state, encoded_value = tff.federated_apply(encode_fn,
(state, value))
client_encoded_value = tff.federated_broadcast(encoded_value)
client_value = tff.federated_map(decode_fn, client_encoded_value)
return new_state, client_value
return StatefulBroadcastFn(initialize_fn=initial_state_fn,
next_fn=encoded_broadcast_fn)
def build_encoded_sum(values, encoders):
"""Builds `StatefulAggregateFn` for `values`, to be encoded by `encoders`.
Args:
values: Values to be encoded by the `StatefulAggregateFn`. Must be
convertible to `tff.Value`.
encoders: A collection of `GatherEncoder` objects to be used for encoding
`values`. Must have the same structure as `values`.
Returns:
A `StatefulAggregateFn` of which `next_fn` encodes the input at
`tff.CLIENTS`, and computes their sum at `tff.SERVER`, automatically
splitting the decoding part based on its commutativity with sum.
Raises:
ValueError: If `values` and `encoders` do not have the same structure.
TypeError: If `encoders` are not instances of `GatherEncoder`, or if
`values` are not compatible with the expected input of the `encoders`.
"""
tf.nest.assert_same_structure(values, encoders)
tf.nest.map_structure(
lambda e, v: _validate_encoder(e, v, te.core.GatherEncoder), encoders,
values)
value_type = tff_framework.type_from_tensors(values)
initial_state_fn = _build_initial_state_tf_computation(encoders)
state_type = initial_state_fn.type_signature.result
nest_encoder = _build_tf_computations_for_gather(state_type, value_type,
encoders)
encoded_sum_fn = _build_encoded_sum_fn(nest_encoder)
return StatefulAggregateFn(initialize_fn=initial_state_fn,
next_fn=encoded_sum_fn)