def zero_fn(): values = tf.nest.map_structure( lambda s: tf.zeros(s.shape, s.dtype), type_utils.type_to_tf_tensor_specs(part_decoded_x_type)) state_update_tensors = tf.nest.map_structure( lambda s: tf.zeros(s.shape, s.dtype), type_utils.type_to_tf_tensor_specs(state_update_tensors_type)) return _accumulator_value(values, state_update_tensors)
def next_fn(global_state, value, weight=None): """Defines next_fn for StatefulAggregateFn.""" # Weighted aggregation is not supported. # TODO(b/140236959): Add an assertion that weight is None here, so the # contract of this method is better established. Will likely cause some # downstream breaks. del weight ####################################### # Define local tf_computations # TODO(b/129567727): Make most of these tf_computations polymorphic # so type manipulation isn't needed. global_state_type = initialize_fn.type_signature.result @computations.tf_computation(global_state_type) def derive_sample_params(global_state): return query.derive_sample_params(global_state) @computations.tf_computation(derive_sample_params.type_signature.result, value.type_signature.member) def preprocess_record(params, record): # TODO(b/123092620): Once TFF passes the expected container type (instead # of AnonymousTuple), we shouldn't need this. record = from_tff_result_fn(record) return query.preprocess_record(params, record) # TODO(b/123092620): We should have the expected container type here. value_type = value_type_fn(value) tensor_specs = type_utils.type_to_tf_tensor_specs(value_type) @computations.tf_computation def zero(): return query.initial_sample_state(tensor_specs) sample_state_type = zero.type_signature.result @computations.tf_computation(sample_state_type, preprocess_record.type_signature.result) def accumulate(sample_state, preprocessed_record): return query.accumulate_preprocessed_record(sample_state, preprocessed_record) @computations.tf_computation(sample_state_type, sample_state_type) def merge(sample_state_1, sample_state_2): return query.merge_sample_states(sample_state_1, sample_state_2) @computations.tf_computation(merge.type_signature.result) def report(sample_state): return sample_state @computations.tf_computation(sample_state_type, global_state_type) def post_process(sample_state, global_state): result, new_global_state = query.get_noised_result( sample_state, global_state) return new_global_state, result ####################################### # Orchestration logic sample_params = intrinsics.federated_map(derive_sample_params, global_state) client_sample_params = intrinsics.federated_broadcast(sample_params) preprocessed_record = intrinsics.federated_map( preprocess_record, (client_sample_params, value)) agg_result = intrinsics.federated_aggregate(preprocessed_record, zero(), accumulate, merge, report) return intrinsics.federated_map(post_process, (agg_result, global_state))