def test_getattr_comp_construction(self, placement):
federated_value = value_impl.to_value(
computation_building_blocks.Reference(
'test',
computation_types.FederatedType([('a', tf.int32),
('b', tf.bool)], placement,
True)), None,
context_stack_impl.context_stack)
get_a_comp = computation_constructing_utils.construct_federated_getattr_comp(
federated_value, 'a')
self.assertEqual(str(get_a_comp), '(x -> x.a)')
get_b_comp = computation_constructing_utils.construct_federated_getattr_comp(
federated_value, 'b')
self.assertEqual(str(get_b_comp), '(x -> x.b)')
non_federated_arg = value_impl.to_value(
computation_building_blocks.Reference(
'test',
computation_types.NamedTupleType([('a', tf.int32),
('b', tf.bool)])), None,
context_stack_impl.context_stack)
with self.assertRaises(TypeError):
_ = computation_constructing_utils.construct_federated_getattr_comp(
non_federated_arg, 'a')
with self.assertRaisesRegexp(ValueError, 'has no element of name c'):
_ = computation_constructing_utils.construct_federated_getattr_comp(
federated_value, 'c')
def test_federated_getattr_comp_fails_value(self):
x = computation_building_blocks.Reference(
'x', computation_types.to_type([('x', tf.int32)]))
with self.assertRaises(TypeError):
computation_constructing_utils.construct_federated_getattr_comp(
value_impl.to_value(x), 'x')