def test_returns_string_for_federated_map(self):
comp = test_utils.create_dummy_called_federated_map(parameter_name='a')
compact_string = comp.compact_representation()
self.assertEqual(compact_string, 'federated_map(<(a -> a),data>)')
formatted_string = comp.formatted_representation()
# pyformat: disable
self.assertEqual(
formatted_string,
'federated_map(<\n'
' (a -> a),\n'
' data\n'
'>)'
)
# pyformat: enable
structural_string = comp.structural_representation()
# pyformat: disable
self.assertEqual(
structural_string,
' Call\n'
' / \\\n'
'federated_map Tuple\n'
' |\n'
' [Lambda(a), data]\n'
' |\n'
' Ref(a)'
)
def test_noops_in_case_of_distinct_maps(self):
called_intrinsic1 = test_utils.create_dummy_called_federated_map(
parameter_name='a', parameter_type=tf.int32)
called_intrinsic2 = test_utils.create_dummy_called_federated_map(
parameter_name='a', parameter_type=tf.float32)
calls = building_blocks.Tuple((called_intrinsic1, called_intrinsic2))
comp = calls
deduped_and_merged_comp, deduped_modified = compiler_transformations.dedupe_and_merge_tuple_intrinsics(
comp, intrinsic_defs.FEDERATED_MAP.uri)
directly_merged_comp, directly_modified = transformations.merge_tuple_intrinsics(
comp, intrinsic_defs.FEDERATED_MAP.uri)
self.assertTrue(deduped_modified)
self.assertTrue(directly_modified)
self.assertEqual(deduped_and_merged_comp.compact_representation(),
directly_merged_comp.compact_representation())