async def _compute_intrinsic_federated_weighted_mean(self, arg): type_utils.check_valid_federated_weighted_mean_argument_tuple_type( arg.type_signature) zipped_arg = await self._compute_intrinsic_federated_zip_at_clients(arg) # TODO(b/134543154): Replace with something that produces a section of # plain TensorFlow code instead of constructing a lambda (so that this # can be executed directly on top of a plain TensorFlow-based executor). multiply_blk = intrinsic_utils.construct_binary_operator_with_upcast( zipped_arg.type_signature.member, tf.multiply) sum_of_products = await self._compute_intrinsic_federated_sum( await self._compute_intrinsic_federated_map( FederatedExecutorValue( anonymous_tuple.AnonymousTuple([ (None, multiply_blk.proto), (None, zipped_arg.internal_representation) ]), computation_types.NamedTupleType( [multiply_blk.type_signature, zipped_arg.type_signature])))) total_weight = await self._compute_intrinsic_federated_sum( FederatedExecutorValue(arg.internal_representation[1], arg.type_signature[1])) divide_arg = await self._compute_intrinsic_federated_zip_at_server( await self.create_tuple( anonymous_tuple.AnonymousTuple([(None, sum_of_products), (None, total_weight)]))) divide_blk = intrinsic_utils.construct_binary_operator_with_upcast( divide_arg.type_signature.member, tf.divide) return await self._compute_intrinsic_federated_apply( FederatedExecutorValue( anonymous_tuple.AnonymousTuple([ (None, divide_blk.proto), (None, divide_arg.internal_representation) ]), computation_types.NamedTupleType( [divide_blk.type_signature, divide_arg.type_signature])))
def test_construct_generic_raises_federated_type(self): bad_type = computation_types.FederatedType( [[tf.int32, tf.int32], computation_types.TensorType(tf.float32, [2])], placement_literals.CLIENTS) with self.assertRaisesRegex(TypeError, 'argument that is not a two-tuple'): intrinsic_utils.construct_binary_operator_with_upcast( bad_type, tf.multiply)
def test_raises_non_callable_op(self): bad_type_ref = computation_building_blocks.Reference( 'x', [tf.float32, tf.float32]) with self.assertRaisesRegex(TypeError, 'non-callable'): intrinsic_utils.apply_binary_operator_with_upcast( bad_type_ref, tf.constant(0)) with self.assertRaisesRegex(TypeError, 'non-callable'): intrinsic_utils.construct_binary_operator_with_upcast( bad_type_ref, tf.constant(0))
def test_raises_tuple_scalar_multiplied_by_nonscalar(self): bad_type_ref = computation_building_blocks.Reference( 'x', [tf.int32, computation_types.TensorType(tf.float32, [2])]) with self.assertRaisesRegex(TypeError, 'incompatible with upcasted'): intrinsic_utils.apply_binary_operator_with_upcast( bad_type_ref, tf.multiply) with self.assertRaisesRegex(TypeError, 'incompatible with upcasted'): intrinsic_utils.construct_binary_operator_with_upcast( bad_type_ref.type_signature, tf.multiply)
def test_raises_incompatible_tuple_and_tensor(self): bad_type_ref = computation_building_blocks.Reference( 'x', computation_types.FederatedType([[tf.int32, tf.int32], tf.float32], placement_literals.CLIENTS)) with self.assertRaisesRegex(TypeError, 'incompatible with upcasted'): intrinsic_utils.apply_binary_operator_with_upcast( bad_type_ref, tf.multiply) with self.assertRaisesRegex(TypeError, 'incompatible with upcasted'): intrinsic_utils.construct_binary_operator_with_upcast( bad_type_ref.type_signature.member, tf.multiply)
def test_construct_integer_type_signature(self): ref = computation_building_blocks.Reference('x', [tf.int32, tf.int32]) multiplier = intrinsic_utils.construct_binary_operator_with_upcast( ref.type_signature, tf.multiply) self.assertEqual( multiplier.type_signature, type_constructors.binary_op(computation_types.to_type(tf.int32)))
def federated_sum(x): zero = intrinsic_utils.zero_for(x.type_signature.member, context_stack) plus_op = value_impl.ValueImpl( intrinsic_utils.construct_binary_operator_with_upcast( computation_types.NamedTupleType( [x.type_signature.member, x.type_signature.member]), tf.add), context_stack) return federated_reduce([x, zero, plus_op])
def test_construct_divide_op_named_tuple_with_scalar_type_signature(self): type_spec = computation_types.to_type([[('a', tf.float32), ('b', tf.float32)], tf.float32]) multiplier = intrinsic_utils.construct_binary_operator_with_upcast( type_spec, tf.divide) expected_function_type = computation_types.FunctionType( type_spec, type_spec[0]) self.assertEqual(multiplier.type_signature, expected_function_type)
def test_construct_op_raises_on_none_operator(self): with self.assertRaisesRegex(TypeError, 'found non-callable'): intrinsic_utils.construct_binary_operator_with_upcast( tf.int32, None)