def test_raises_type_error_with_nonfederated_arg(self): ref = computation_building_blocks.Reference('x', tf.int32) fn = computation_building_blocks.Lambda(ref.name, ref.type_signature, ref) arg = computation_building_blocks.Data('y', tf.int32) with self.assertRaises(TypeError): computation_constructing_utils.create_federated_map(fn, arg)
def test_raises_type_error_with_nonfunctional_fn(self): fn = computation_building_blocks.Reference('x', tf.int32) arg_type = computation_types.FederatedType(tf.bool, placements.CLIENTS, False) arg = computation_building_blocks.Data('y', arg_type) with self.assertRaises(TypeError): computation_constructing_utils.create_federated_map(fn, arg)
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 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. # 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) 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) fn = value_impl.ValueImpl.get_comp(fn) arg = value_impl.ValueImpl.get_comp(arg) comp = computation_constructing_utils.create_federated_map(fn, arg) return value_impl.ValueImpl(comp, self._context_stack)
def test_returns_federated_map(self): ref = computation_building_blocks.Reference('x', tf.int32) fn = computation_building_blocks.Lambda(ref.name, ref.type_signature, ref) arg_type = computation_types.FederatedType(tf.int32, placements.CLIENTS, False) arg = computation_building_blocks.Data('y', arg_type) comp = computation_constructing_utils.create_federated_map(fn, arg) self.assertEqual(comp.tff_repr, 'federated_map(<(x -> x),y>)') self.assertEqual(str(comp.type_signature), '{int32}@CLIENTS')
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 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) type_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(fn, arg) return value_impl.ValueImpl(comp, self._context_stack)
def create_dummy_called_federated_map(parameter_name, parameter_type=tf.int32): r"""Returns a dummy called federated map. Call / \ federated_map 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) arg = computation_building_blocks.Data('data', arg_type) return computation_constructing_utils.create_federated_map(fn, arg)
def test_raises_type_error_with_none_fn(self): arg = computation_building_blocks.Data('y', tf.int32) with self.assertRaises(TypeError): computation_constructing_utils.create_federated_map(None, arg)