def federated_map_all_equal(self, fn, arg): """Implements `federated_map` as defined in `api/intrinsic.py`. Implements `federated_map` as defined in `api/intrinsic.py` with an argument with the `all_equal` bit set. Args: fn: As in `api/intrinsics.py`. arg: As in `api/intrinsics.py`, with the `all_equal` bit set. 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 clients 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. arg = self.federated_zip(arg) value_utils.check_federated_value_placement(arg, placements.CLIENTS, 'value to be mapped') # 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)) fn = value_impl.ValueImpl.get_comp(fn) arg = value_impl.ValueImpl.get_comp(arg) comp = computation_constructing_utils.create_federated_map_all_equal( fn, arg) return value_impl.ValueImpl(comp, self._context_stack)
def create_dummy_called_federated_map_all_equal(parameter_name, parameter_type=tf.int32): r"""Returns a dummy called federated map. Call / \ federated_map_all_equal 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.CLIENTS, all_equal=True) arg = computation_building_blocks.Data('data', arg_type) return computation_constructing_utils.create_federated_map_all_equal(fn, arg)
def test_converts_federated_map_all_equal_to_federated_map(self): fed_type_all_equal = computation_types.FederatedType( tf.int32, placements.CLIENTS, all_equal=True) normalized_fed_type = computation_types.FederatedType( tf.int32, placements.CLIENTS) int_ref = building_blocks.Reference('x', tf.int32) int_identity = building_blocks.Lambda('x', tf.int32, int_ref) federated_int_ref = building_blocks.Reference('y', fed_type_all_equal) called_federated_map_all_equal = computation_constructing_utils.create_federated_map_all_equal( int_identity, federated_int_ref) normalized_federated_map = mapreduce_transformations.normalize_all_equal_bit( called_federated_map_all_equal) self.assertEqual(called_federated_map_all_equal.function.uri, intrinsic_defs.FEDERATED_MAP_ALL_EQUAL.uri) self.assertIsInstance(normalized_federated_map, building_blocks.Call) self.assertIsInstance(normalized_federated_map.function, building_blocks.Intrinsic) self.assertEqual(normalized_federated_map.function.uri, intrinsic_defs.FEDERATED_MAP.uri) self.assertEqual(normalized_federated_map.type_signature, normalized_fed_type)