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, tensor_encoding.core.GatherEncoder),
encoders, values)
value_type = type_utils.type_from_tensors(values)
initial_state_fn, state_type = _build_initial_state_tf_computation(encoders)
nest_encoder = _build_tf_computations_for_gather(state_type, value_type,
encoders)
encoded_sum_fn = _build_encoded_sum_fn(nest_encoder)
return computation_utils.StatefulAggregateFn(
initialize_fn=initial_state_fn, next_fn=encoded_sum_fn)
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, tensor_encoding.core.GatherEncoder
), encoders, values)
value_type = type_utils.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)
@computations.tf_computation(value_type,
computation_types.to_type(tf.float32))
def multiply_fn(value, weight):
return tf.nest.map_structure(lambda v: v * tf.cast(weight, v.dtype),
value)
@computations.tf_computation(value_type,
computation_types.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 = intrinsics.federated_map(multiply_fn,
[values, weight])
updated_state, summed_decoded_values = encoded_sum_fn(
state, weighted_values)
summed_weights = intrinsics.federated_sum(weight)
decoded_values = intrinsics.federated_map(
divide_fn, [summed_decoded_values, summed_weights])
return updated_state, decoded_values
return computation_utils.StatefulAggregateFn(
initialize_fn=initial_state_fn, next_fn=encoded_mean_fn)
def wrap_aggregate_fn(dp_aggregate_fn, sample_value):
tff_types = type_utils.type_from_tensors(sample_value)
@computations.federated_computation
def run_initialize():
return intrinsics.federated_value(dp_aggregate_fn.initialize(),
placement_literals.SERVER)
@computations.federated_computation(run_initialize.type_signature.result,
computation_types.FederatedType(
tff_types,
placement_literals.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, tensor_encoding.core.SimpleEncoder
), encoders, values)
value_type = type_utils.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 = intrinsics.federated_map(
encode_fn, (state, value))
client_encoded_value = intrinsics.federated_broadcast(encoded_value)
client_value = intrinsics.federated_map(decode_fn,
client_encoded_value)
return new_state, client_value
return computation_utils.StatefulBroadcastFn(
initialize_fn=initial_state_fn, next_fn=encoded_broadcast_fn)