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 construct_tensorflow_calling_lambda_on_concrete_arg( parameter: building_blocks.Reference, body: building_blocks.ComputationBuildingBlock, concrete_arg: building_blocks.ComputationBuildingBlock): """Generates TensorFlow for lambda invocation with given arg, body and param. That is, generates TensorFlow block encapsulating the logic represented by invoking a function with parameter `parameter` and body `body`, with argument `concrete_arg`. Via the guarantee made in `compiled_computation_transforms.TupleCalledGraphs`, this function makes the claim that the computations which define `concrete_arg` will be executed exactly once in the generated TenosorFlow. Args: parameter: Instance of `building_blocks.Reference` defining the parameter of the function to be generated and invoked, as described above. After calling this transformation, every instance of parameter` in `body` will represent a reference to `concrete_arg`. body: `building_blocks.ComputationBuildingBlock` representing the body of the function for which we are generating TensorFlow. concrete_arg: `building_blocks.ComputationBuildingBlock` representing the argument to be passed to the resulting function. `concrete_arg` will then be referred to by every occurrence of `parameter` in `body`. Therefore `concrete_arg` must have an equivalent type signature to that of `parameter`. Returns: A called `building_blocks.CompiledComputation`, as specified above. Raises: TypeError: If the arguments are of the wrong types, or the type signature of `concrete_arg` does not match that of `parameter`. """ py_typecheck.check_type(parameter, building_blocks.Reference) py_typecheck.check_type(body, building_blocks.ComputationBuildingBlock) py_typecheck.check_type(concrete_arg, building_blocks.ComputationBuildingBlock) type_utils.check_equivalent_types(parameter.type_signature, concrete_arg.type_signature) def _generate_simple_tensorflow(comp): tf_parser_callable = tree_to_cc_transformations.TFParser() comp, _ = tree_transformations.insert_called_tf_identity_at_leaves( comp) comp, _ = transformation_utils.transform_postorder( comp, tf_parser_callable) return comp encapsulating_lambda = _generate_simple_tensorflow( building_blocks.Lambda(parameter.name, parameter.type_signature, body)) comp_called = _generate_simple_tensorflow( building_blocks.Call(encapsulating_lambda, concrete_arg)) return comp_called
def federated_secure_sum(self, value, bitwidth): """Implements `federated_secure_sum` as defined in `api/intrinsics.py`.""" value = value_impl.to_value(value, None, self._context_stack) value = value_utils.ensure_federated_value(value, placements.CLIENTS, 'value to be summed') type_utils.check_is_structure_of_integers(value.type_signature) bitwidth = value_impl.to_value(bitwidth, None, self._context_stack) type_utils.check_equivalent_types(value.type_signature.member, bitwidth.type_signature) value = value_impl.ValueImpl.get_comp(value) bitwidth = value_impl.ValueImpl.get_comp(bitwidth) comp = building_block_factory.create_federated_secure_sum(value, bitwidth) return value_impl.ValueImpl(comp, self._context_stack)
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])))
async def embed_tf_scalar_constant(executor, type_spec, val): """Embeds a constant `val` of TFF type `type_spec` in `executor`. Args: executor: An instance of `tff.framework.Executor`. type_spec: An instance of `tff.Type`. val: A scalar value. Returns: An instance of `tff.framework.ExecutorValue` containing an embedded value. """ py_typecheck.check_type(executor, executor_base.Executor) fn_building_block = (building_block_factory.create_tensorflow_constant( type_spec, val)) embedded_val = await executor.create_call(await executor.create_value( fn_building_block.function.proto, fn_building_block.function.type_signature)) type_utils.check_equivalent_types(embedded_val.type_signature, type_spec) return embedded_val
async def _compute_intrinsic_federated_reduce(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) != 3: raise ValueError( 'Expected 3 elements in the `federated_reduce()` 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] op_type = arg.type_signature[2] type_utils.check_equivalent_types( op_type, type_factory.reduction_op(zero_type, item_type)) val = arg.internal_representation[0] py_typecheck.check_type(val, list) child = self._target_executors[placement_literals.SERVER][0] async def _move(v): return await child.create_value(await v.compute(), item_type) items = await asyncio.gather(*[_move(v) for v in val]) zero = await child.create_value( await (await self.create_selection(arg, index=1)).compute(), zero_type) op = await child.create_value(arg.internal_representation[2], op_type) result = zero for item in items: result = await child.create_call( op, await child.create_tuple( anonymous_tuple.AnonymousTuple([(None, result), (None, item)]))) return FederatedExecutorValue([result], computation_types.FederatedType( result.type_signature, placement_literals.SERVER, all_equal=True))
async def _embed_tf_scalar_constant(executor, type_spec, val): """Embeds a constant `val` of TFF type `type_spec` in `executor`. Args: executor: An instance of `tff.framework.Executor`. type_spec: An instance of `tff.Type`. val: A scalar value. Returns: An instance of `tff.framework.ExecutorValue` containing an embedded value. """ # TODO(b/134543154): Perhaps graduate this and the function below it into a # separate library, so that it can be used in other places. py_typecheck.check_type(executor, executor_base.Executor) fn_building_block = (building_block_factory.create_tensorflow_constant( type_spec, val)) embedded_val = await executor.create_call(await executor.create_value( fn_building_block.function.proto, fn_building_block.function.type_signature)) type_utils.check_equivalent_types(embedded_val.type_signature, type_spec) return embedded_val
def parse_federated_aggregate_argument_types(type_spec): """Verifies and parses `type_spec` into constituents. Args: type_spec: An instance of `computation_types.NamedTupleType`. Returns: A tuple of (value_type, zero_type, accumulate_type, merge_type, report_type) for the 5 type constituents. """ py_typecheck.check_type(type_spec, computation_types.NamedTupleType) py_typecheck.check_len(type_spec, 5) value_type = type_spec[0] py_typecheck.check_type(value_type, computation_types.FederatedType) item_type = value_type.member zero_type = type_spec[1] accumulate_type = type_spec[2] type_utils.check_equivalent_types( accumulate_type, type_factory.reduction_op(zero_type, item_type)) merge_type = type_spec[3] type_utils.check_equivalent_types(merge_type, type_factory.binary_op(zero_type)) report_type = type_spec[4] py_typecheck.check_type(report_type, computation_types.FunctionType) type_utils.check_equivalent_types(report_type.parameter, zero_type) return value_type, zero_type, accumulate_type, merge_type, report_type
async def _compute_intrinsic_federated_aggregate(self, arg): val_type, zero_type, accumulate_type, _, report_type = ( executor_utils.parse_federated_aggregate_argument_types( arg.type_signature)) py_typecheck.check_type(arg.internal_representation, anonymous_tuple.AnonymousTuple) py_typecheck.check_len(arg.internal_representation, 5) # 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( FederatingExecutorValue( 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( FederatingExecutorValue( anonymous_tuple.AnonymousTuple([ (None, report), (None, pre_report.internal_representation) ]), computation_types.NamedTupleType( (report_type, pre_report.type_signature))))
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 = ( building_block_factory.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_factory.binary_op(type_spec)) return embedded_val
def test_check_equivalent_types(self): type_utils.check_equivalent_types(tf.int32, tf.int32) with self.assertRaises(TypeError): type_utils.check_equivalent_types(tf.int32, tf.bool)