def test_returns_federated_apply(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.SERVER, True) arg = computation_building_blocks.Data('y', arg_type) comp = computation_constructing_utils.create_federated_apply(fn, arg) self.assertEqual(comp.tff_repr, 'federated_apply(<(x -> x),y>)') self.assertEqual(str(comp.type_signature), 'int32@SERVER')
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 = computation_building_blocks.Data('data', arg_type) return computation_constructing_utils.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) type_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 = computation_constructing_utils.create_federated_apply(fn, arg) return value_impl.ValueImpl(comp, self._context_stack)
def test_raises_type_error_with_none_arg(self): ref = computation_building_blocks.Reference('x', tf.int32) fn = computation_building_blocks.Lambda(ref.name, ref.type_signature, ref) with self.assertRaises(TypeError): computation_constructing_utils.create_federated_apply(fn, None)
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_apply(None, arg)