def test_passes_with_federated_map(self):
intrinsic = computation_building_blocks.Intrinsic(
intrinsic_defs.FEDERATED_MAP.uri,
computation_types.FunctionType([
computation_types.FunctionType(tf.int32, tf.float32),
computation_types.FederatedType(tf.int32, placements.CLIENTS)
], computation_types.FederatedType(tf.float32, placements.CLIENTS)))
tree_analysis.check_intrinsics_whitelisted_for_reduction(intrinsic)
def test_raises_with_federated_mean(self):
intrinsic = computation_building_blocks.Intrinsic(
intrinsic_defs.FEDERATED_MEAN.uri,
computation_types.FunctionType(
computation_types.FederatedType(tf.int32, placements.CLIENTS),
computation_types.FederatedType(tf.int32, placements.SERVER)))
with self.assertRaisesRegex(ValueError,
intrinsic.compact_representation()):
tree_analysis.check_intrinsics_whitelisted_for_reduction(intrinsic)
def test_generic_divide_reduces(self):
uri = intrinsic_defs.GENERIC_DIVIDE.uri
context_stack = context_stack_impl.context_stack
comp = computation_building_blocks.Intrinsic(
uri,
computation_types.FunctionType([tf.float32, tf.float32],
tf.float32))
count_before_reduction = _count_intrinsics(comp, uri)
reduced, modified = value_transformations.replace_all_intrinsics_with_bodies(
comp, context_stack)
count_after_reduction = _count_intrinsics(reduced, uri)
self.assertGreater(count_before_reduction, 0)
self.assertEqual(count_after_reduction, 0)
tree_analysis.check_intrinsics_whitelisted_for_reduction(reduced)
self.assertTrue(modified)
def test_raises_on_none(self):
with self.assertRaises(TypeError):
tree_analysis.check_intrinsics_whitelisted_for_reduction(None)
