async def _compute_intrinsic_federated_sum(self, arg): py_typecheck.check_type(arg.type_signature, computation_types.FederatedType) zero_building_block = ( computation_constructing_utils.construct_tensorflow_constant( arg.type_signature.member, 0)) zero = await self.create_call(await self.create_value( zero_building_block.function.proto, zero_building_block.function.type_signature)) type_utils.check_equivalent_types(arg.type_signature.member, zero.type_signature) # TODO(b/134543154): There is an opportunity here to import something more # in line with the usage (no building block wrapping, etc.) plus_building_block = (computation_constructing_utils. construct_tensorflow_binary_operator( zero.type_signature, tf.add)) plus = await self.create_value(plus_building_block.proto, plus_building_block.type_signature) type_utils.check_equivalent_types( plus.type_signature, type_constructors.binary_op(zero.type_signature)) return await self._compute_intrinsic_federated_reduce( FederatedExecutorValue( anonymous_tuple.AnonymousTuple([ (None, arg.internal_representation), (None, zero.internal_representation), (None, plus.internal_representation) ]), computation_types.NamedTupleType([ arg.type_signature, zero.type_signature, plus.type_signature ])))
def test_construct_integer_type_signature(self): ref = computation_building_blocks.Reference('x', [tf.int32, tf.int32]) multiplier = intrinsic_utils.construct_binary_operator_with_upcast( ref.type_signature, tf.multiply) self.assertEqual( multiplier.type_signature, type_constructors.binary_op(computation_types.to_type(tf.int32)))
async def _compute_intrinsic_federated_aggregate(self, arg): py_typecheck.check_type(arg.type_signature, computation_types.NamedTupleType) py_typecheck.check_type(arg.internal_representation, anonymous_tuple.AnonymousTuple) if len(arg.internal_representation) != 5: raise ValueError( 'Expected 5 elements in the `federated_aggregate()` argument tuple, ' 'found {}.'.format(len(arg.internal_representation))) val_type = arg.type_signature[0] py_typecheck.check_type(val_type, computation_types.FederatedType) item_type = val_type.member zero_type = arg.type_signature[1] accumulate_type = arg.type_signature[2] type_utils.check_equivalent_types( accumulate_type, type_constructors.reduction_op(zero_type, item_type)) merge_type = arg.type_signature[3] type_utils.check_equivalent_types( merge_type, type_constructors.binary_op(zero_type)) report_type = arg.type_signature[4] py_typecheck.check_type(report_type, computation_types.FunctionType) type_utils.check_equivalent_types(report_type.parameter, zero_type) # NOTE: This is a simple initial implementation that simply forwards this # to `federated_reduce()`. The more complete implementation would be able # to take advantage of the parallelism afforded by `merge` to reduce the # cost from liner (with respect to the number of clients) to sub-linear. # TODO(b/134543154): Expand this implementation to take advantage of the # parallelism afforded by `merge`. val = arg.internal_representation[0] zero = arg.internal_representation[1] accumulate = arg.internal_representation[2] pre_report = await self._compute_intrinsic_federated_reduce( FederatedExecutorValue( anonymous_tuple.AnonymousTuple([(None, val), (None, zero), (None, accumulate)]), computation_types.NamedTupleType( [val_type, zero_type, accumulate_type]))) py_typecheck.check_type(pre_report.type_signature, computation_types.FederatedType) type_utils.check_equivalent_types(pre_report.type_signature.member, report_type.parameter) report = arg.internal_representation[4] return await self._compute_intrinsic_federated_apply( FederatedExecutorValue( anonymous_tuple.AnonymousTuple([ (None, report), (None, pre_report.internal_representation) ]), computation_types.NamedTupleType( [report_type, pre_report.type_signature])))
def plus_for(type_spec, context_stack): """Constructs PLUS intrinsic that operates on values of TFF type `type_spec`. Args: type_spec: An instance of `types.Type` or something convertible to it. intrinsic. context_stack: The context stack to use. Returns: The `PLUS` intrinsic of type `<T,T> -> T`, where `T` represents `type_spec`. """ type_spec = computation_types.to_type(type_spec) py_typecheck.check_type(context_stack, context_stack_base.ContextStack) return value_impl.ValueImpl( computation_building_blocks.Intrinsic( intrinsic_defs.GENERIC_PLUS.uri, type_constructors.binary_op(type_spec)), context_stack)
async def _embed_tf_binary_operator(executor, type_spec, op): """Embeds a binary operator `op` on `type_spec`-typed values in `executor`. Args: executor: An instance of `tff.framework.Executor`. type_spec: An instance of `tff.Type` of the type of values that the binary operator accepts as input and returns as output. op: An operator function (such as `tf.add` or `tf.multiply`) to apply to the tensor-level constituents of the values, pointwise. Returns: An instance of `tff.framework.ExecutorValue` representing the operator in a form embedded into the executor. """ # TODO(b/134543154): There is an opportunity here to import something more # in line with the usage (no building block wrapping, etc.) fn_building_block = ( computation_constructing_utils.create_tensorflow_binary_operator( type_spec, op)) embedded_val = await executor.create_value( fn_building_block.proto, fn_building_block.type_signature) type_utils.check_equivalent_types(embedded_val.type_signature, type_constructors.binary_op(type_spec)) return embedded_val
# a = generic_partial_reduce(x, zero, accumulate, INTERMEDIATE_AGGREGATORS) # b = generic_reduce(a, zero, merge, SERVER) # c = generic_map(report, b) # return c # # Actual implementations might vary. # # Type signature: <{T}@CLIENTS,U,(<U,T>->U),(<U,U>->U),(U->R)> -> R@SERVER FEDERATED_AGGREGATE = IntrinsicDef( 'FEDERATED_AGGREGATE', 'federated_aggregate', computation_types.FunctionType(parameter=[ type_constructors.at_clients(computation_types.AbstractType('T')), computation_types.AbstractType('U'), type_constructors.reduction_op(computation_types.AbstractType('U'), computation_types.AbstractType('T')), type_constructors.binary_op(computation_types.AbstractType('U')), computation_types.FunctionType(computation_types.AbstractType('U'), computation_types.AbstractType('R')) ], result=type_constructors.at_server( computation_types.AbstractType('R')))) # Applies a given function to a value on the server. # # Type signature: <(T->U),T@SERVER> -> U@SERVER FEDERATED_APPLY = IntrinsicDef( 'FEDERATED_APPLY', 'federated_apply', computation_types.FunctionType(parameter=[ computation_types.FunctionType(computation_types.AbstractType('T'), computation_types.AbstractType('U')), type_constructors.at_server(computation_types.AbstractType('T')),
def test_binary_op(self): self.assertEqual(str(type_constructors.binary_op(tf.bool)), '(<bool,bool> -> bool)')