def federated_map(self, fn, arg): """Implements `federated_map` as defined in `api/intrinsics.py`.""" # TODO(b/113112108): Possibly lift the restriction that the mapped value # must be placed at the server or clients. Would occur after adding support # for placement labels in the federated types, and expanding the type # specification of the intrinsic this is based on to work with federated # values of arbitrary placement. arg = value_impl.to_value(arg, None, self._context_stack) arg = value_utils.ensure_federated_value(arg, label='value to be mapped') fn = value_impl.to_value(fn, None, self._context_stack, parameter_type_hint=arg.type_signature.member) py_typecheck.check_type(fn, value_base.Value) py_typecheck.check_type(fn.type_signature, computation_types.FunctionType) if not type_utils.is_assignable_from(fn.type_signature.parameter, arg.type_signature.member): raise TypeError( 'The mapping function expects a parameter of type {}, but member ' 'constituents of the mapped value are of incompatible type {}.' .format(fn.type_signature.parameter, arg.type_signature.member)) # TODO(b/144384398): Change structure to one that maps the placement type # to the building_block function that fits it, in a way that allows the # appropriate type checks. if arg.type_signature.placement is placements.SERVER: if not arg.type_signature.all_equal: raise TypeError( 'Arguments placed at {} should be equal at all locations.'. format(placements.SERVER)) fn = value_impl.ValueImpl.get_comp(fn) arg = value_impl.ValueImpl.get_comp(arg) comp = building_block_factory.create_federated_apply(fn, arg) elif arg.type_signature.placement is placements.CLIENTS: fn = value_impl.ValueImpl.get_comp(fn) arg = value_impl.ValueImpl.get_comp(arg) comp = building_block_factory.create_federated_map(fn, arg) else: raise TypeError( 'The argument should be placed at {} or {}, placed at {} instead.' .format(placements.SERVER, placements.CLIENTS, arg.type_signature.placement)) return value_impl.ValueImpl(comp, self._context_stack)
def create_dummy_called_federated_apply(parameter_name, parameter_type=tf.int32): r"""Returns a dummy called federated apply. Call / \ federated_apply Tuple | [Lambda(x), data] | Ref(x) Args: parameter_name: The name of the parameter. parameter_type: The type of the parameter. """ fn = create_identity_function(parameter_name, parameter_type) arg_type = computation_types.FederatedType(parameter_type, placements.SERVER) arg = building_blocks.Data('data', arg_type) return building_block_factory.create_federated_apply(fn, arg)
def federated_apply(self, fn, arg): """Implements `federated_apply` as defined in `api/intrinsics.py`. Args: fn: As in `api/intrinsics.py`. arg: As in `api/intrinsics.py`. Returns: As in `api/intrinsics.py`. Raises: TypeError: As in `api/intrinsics.py`. """ fn = value_impl.to_value(fn, None, self._context_stack) py_typecheck.check_type(fn, value_base.Value) py_typecheck.check_type(fn.type_signature, computation_types.FunctionType) arg = value_impl.to_value(arg, None, self._context_stack) if isinstance(arg.type_signature, computation_types.NamedTupleType): if len(anonymous_tuple.to_elements(arg.type_signature)) >= 2: # We've been passed a value which the user expects to be zipped. arg = self.federated_zip(arg) value_utils.check_federated_value_placement(arg, placements.SERVER, 'the argument') if not arg.type_signature.all_equal: raise TypeError('The argument should be equal at all locations.') if not type_utils.is_assignable_from(fn.type_signature.parameter, arg.type_signature.member): raise TypeError( 'The function to apply expects a parameter of type {}, but member ' 'constituents of the argument are of an incompatible type {}.'. format(fn.type_signature.parameter, arg.type_signature.member)) fn = value_impl.ValueImpl.get_comp(fn) arg = value_impl.ValueImpl.get_comp(arg) comp = building_block_factory.create_federated_apply(fn, arg) return value_impl.ValueImpl(comp, self._context_stack)
def test_compiles_lambda_under_federated_comp_to_tf(self): ref_to_x = building_blocks.Reference( 'x', computation_types.StructType([tf.int32, tf.float32])) identity_lambda = building_blocks.Lambda(ref_to_x.name, ref_to_x.type_signature, ref_to_x) federated_data = building_blocks.Data( 'a', computation_types.FederatedType( computation_types.StructType([tf.int32, tf.float32]), placements.SERVER)) applied = building_block_factory.create_federated_apply( identity_lambda, federated_data) transformed = compiler.compile_local_subcomputations_to_tensorflow( applied) self.assertIsInstance(transformed, building_blocks.Call) self.assertIsInstance(transformed.function, building_blocks.Intrinsic) self.assertIsInstance(transformed.argument[0], building_blocks.CompiledComputation) self.assertEqual(transformed.argument[1], federated_data) self.assertEqual(transformed.argument[0].type_signature, identity_lambda.type_signature)
def federated_map(fn, arg): """Maps a federated value pointwise using a mapping function. The function `fn` is applied separately across the group of devices represented by the placement type of `arg`. For example, if `value` has placement type `tff.CLIENTS`, then `fn` is applied to each client individually. In particular, this operation does not alter the placement of the federated value. Args: fn: A mapping function to apply pointwise to member constituents of `arg`. The parameter of this function must be of the same type as the member constituents of `arg`. arg: A value of a TFF federated type (or a value that can be implicitly converted into a TFF federated type, e.g., by zipping) placed at `tff.CLIENTS` or `tff.SERVER`. Returns: A federated value with the same placement as `arg` that represents the result of `fn` on the member constituent of `arg`. Raises: TypeError: If the arguments are not of the appropriate types. """ # TODO(b/113112108): Possibly lift the restriction that the mapped value # must be placed at the server or clients. Would occur after adding support # for placement labels in the federated types, and expanding the type # specification of the intrinsic this is based on to work with federated # values of arbitrary placement. arg = value_impl.to_value(arg, None) arg = value_utils.ensure_federated_value(arg, label='value to be mapped') fn = value_impl.to_value(fn, None, parameter_type_hint=arg.type_signature.member) py_typecheck.check_type(fn, value_impl.Value) py_typecheck.check_type(fn.type_signature, computation_types.FunctionType) if not fn.type_signature.parameter.is_assignable_from( arg.type_signature.member): raise TypeError( 'The mapping function expects a parameter of type {}, but member ' 'constituents of the mapped value are of incompatible type {}.'. format(fn.type_signature.parameter, arg.type_signature.member)) # TODO(b/144384398): Change structure to one that maps the placement type # to the building_block function that fits it, in a way that allows the # appropriate type checks. if arg.type_signature.placement is placements.SERVER: if not arg.type_signature.all_equal: raise TypeError( 'Arguments placed at {} should be equal at all locations.'. format(placements.SERVER)) comp = building_block_factory.create_federated_apply(fn.comp, arg.comp) elif arg.type_signature.placement is placements.CLIENTS: comp = building_block_factory.create_federated_map(fn.comp, arg.comp) else: raise TypeError( 'Expected `arg` to have a type with a supported placement, ' 'found {}.'.format(arg.type_signature.placement)) comp = _bind_comp_as_reference(comp) return value_impl.Value(comp)
def federated_map(self, fn, arg): """Implements `federated_map` as defined in `api/intrinsics.py`. Args: fn: As in `api/intrinsics.py`. arg: As in `api/intrinsics.py`. Returns: As in `api/intrinsics.py`. Raises: TypeError: As in `api/intrinsics.py`. """ # TODO(b/113112108): Possibly lift the restriction that the mapped value # must be placed at the server or clients. Would occur after adding support # for placement labels in the federated types, and expanding the type # specification of the intrinsic this is based on to work with federated # values of arbitrary placement. arg = value_impl.to_value(arg, None, self._context_stack) if isinstance(arg.type_signature, computation_types.NamedTupleType): if len(anonymous_tuple.to_elements(arg.type_signature)) >= 2: # We've been passed a value which the user expects to be zipped. named_type_signatures = anonymous_tuple.to_elements(arg.type_signature) _, first_type_signature = named_type_signatures[0] for _, type_signature in named_type_signatures: py_typecheck.check_type(type_signature, computation_types.FederatedType) if type_signature.placement is not first_type_signature.placement: raise TypeError( 'You cannot apply federated_map on nested values with mixed ' 'placements (was given a nested value of type {}).'.format( arg.type_signature)) arg = self.federated_zip(arg) py_typecheck.check_type(arg.type_signature, computation_types.FederatedType) # TODO(b/113112108): Add support for polymorphic templates auto-instantiated # here based on the actual type of the argument. fn = value_impl.to_value(fn, None, self._context_stack) py_typecheck.check_type(fn, value_base.Value) py_typecheck.check_type(fn.type_signature, computation_types.FunctionType) if not type_utils.is_assignable_from(fn.type_signature.parameter, arg.type_signature.member): raise TypeError( 'The mapping function expects a parameter of type {}, but member ' 'constituents of the mapped value are of incompatible type {}.' .format(fn.type_signature.parameter, arg.type_signature.member)) # TODO(b/144384398): Change structure to one that maps the placement type # to the building_block function that fits it, in a way that allows the # appropriate type checks. if arg.type_signature.placement is placements.SERVER: if not arg.type_signature.all_equal: raise TypeError( 'Arguments placed at {} should be equal at all locations.'.format( placements.SERVER)) fn = value_impl.ValueImpl.get_comp(fn) arg = value_impl.ValueImpl.get_comp(arg) comp = building_block_factory.create_federated_apply(fn, arg) elif arg.type_signature.placement is placements.CLIENTS: fn = value_impl.ValueImpl.get_comp(fn) arg = value_impl.ValueImpl.get_comp(arg) comp = building_block_factory.create_federated_map(fn, arg) else: raise TypeError( 'The argument should be placed at {} or {}, placed at {} instead.' .format(placements.SERVER, placements.CLIENTS, arg.type_signature.placement)) return value_impl.ValueImpl(comp, self._context_stack)