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)