class ExecutorsTest(parameterized.TestCase):
@executor_test_utils.executors
def test_without_arguments(self):
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
result = add_one(5)
self.assertEqual(result, 6)
@executor_test_utils.executors()
def test_with_no_arguments(self):
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
result = add_one(5)
self.assertEqual(result, 6)
@executor_test_utils.executors(
('reference', reference_executor.ReferenceExecutor()),)
def test_with_one_argument(self):
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
result = add_one(5)
self.assertEqual(result, 6)
@executor_test_utils.executors(
('reference', reference_executor.ReferenceExecutor()),
('local', executor_stacks.local_executor_factory()),
)
def test_with_two_argument(self):
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
result = add_one(5)
self.assertEqual(result, 6)
def test_setattr_federated_named_tuple_type_with_unnamed_element(self):
@computations.federated_computation(
computation_types.FederatedType([('a', tf.int32),
(None, tf.float32),
('b', tf.bool)],
placement_literals.CLIENTS))
def foo(x):
x.b = False
return x
with context_stack_impl.context_stack.install(
reference_executor.ReferenceExecutor()):
self.assertEqual(
foo([[5, 1.0, True], [0, 2.0, True], [-5, 3.0, False]]), [
anonymous_tuple.AnonymousTuple([
('a', 5),
(None, 1.0),
('b', False),
]),
anonymous_tuple.AnonymousTuple([
('a', 0),
(None, 2.0),
('b', False),
]),
anonymous_tuple.AnonymousTuple([
('a', -5),
(None, 3.0),
('b', False),
])
])
def test_with_reference_executor(self):
context_stack = context_stack_impl.context_stack
executor = reference_executor.ReferenceExecutor()
self.assertIsNot(context_stack.current, executor)
set_default_executor.set_default_executor(executor)
self.assertIs(context_stack.current, executor)
def test_setattr_named_tuple_type_bool(self):
@computations.federated_computation([('a', tf.int32), ('b', tf.bool)])
def foo(x):
x.b = False
return x
with context_stack_impl.context_stack.install(
reference_executor.ReferenceExecutor()):
self.assertEqual(foo([5, True]).b, False)
self.assertEqual(foo([5, False]).b, False)
self.assertEqual(foo([0, True]).a, 0)
def decorator(fn, named_executors=None):
"""Construct a custom `parameterized.named_parameter` decorator for `fn`."""
if not named_executors:
named_executors = [
('reference', reference_executor.ReferenceExecutor(compiler=None)),
('local', executor_stacks.local_executor_factory()),
]
named_parameters_decorator = parameterized.named_parameters(named_executors)
fn = executor_decorator(fn)
fn = named_parameters_decorator(fn)
return fn
def test_setattr_named_tuple_with_unnamed_element(self):
@computations.federated_computation([('a', tf.int32),
(None, tf.float32),
('b', tf.bool)])
def foo(x):
x.b = False
return x
with context_stack_impl.context_stack.install(
reference_executor.ReferenceExecutor()):
self.assertEqual(foo([5, 1.0, True]).b, False)
self.assertEqual(foo([5, 1.0, False]).b, False)
self.assertEqual(foo([5, 1.0, True])[1], 1.0)
def test_setattr_federated_named_tuple_int_on_server(self):
@computations.federated_computation(
computation_types.FederatedType([('a', tf.int32), ('b', tf.bool)],
placement_literals.SERVER,
all_equal=True))
def foo(x):
x.a = 10
return x
with context_stack_impl.context_stack.install(
reference_executor.ReferenceExecutor()):
self.assertEqual(foo([5, True]),
structure.Struct([('a', 10), ('b', True)]))
def test_reference_executor(self):
set_default_executor.set_default_executor(
reference_executor.ReferenceExecutor())
self.assertIsInstance(context_stack_impl.context_stack.current,
reference_executor.ReferenceExecutor)
@computations.tf_computation(computation_types.SequenceType(tf.int32))
def comp(ds):
return ds.take(5).reduce(np.int32(0), lambda x, y: x + y)
ds = tf.data.Dataset.range(1).map(lambda x: tf.constant(5)).repeat(10)
v = comp(ds)
self.assertEqual(v, 25)
def decorator(fn, *named_executors):
"""Construct a custom `parameterized.named_parameter` decorator for `fn`."""
wraps_decorator = functools.wraps(fn)
if not named_executors:
named_executors = [
('reference', reference_executor.ReferenceExecutor()),
('local', executor_stacks.local_executor_factory()),
]
named_parameters_decorator = parameterized.named_parameters(
*named_executors)
fn = executor_decorator(fn)
fn = named_parameters_decorator(fn)
fn = wraps_decorator(fn)
return fn
def test_setattr_federated_named_tuple_int(self):
@computations.federated_computation(
computation_types.FederatedType([('a', tf.int32), ('b', tf.bool)],
placement_literals.CLIENTS))
def foo(x):
x.a = 10
return x
with context_stack_impl.context_stack.install(
reference_executor.ReferenceExecutor()):
self.assertEqual(foo([[5, True], [0, False], [-5, True]]), [
anonymous_tuple.AnonymousTuple([('a', 10), ('b', True)]),
anonymous_tuple.AnonymousTuple([('a', 10), ('b', False)]),
anonymous_tuple.AnonymousTuple([('a', 10), ('b', True)])
])
def test_setattr_federated_named_tuple_type_bool(self):
@computations.federated_computation(
computation_types.FederatedType([('a', tf.int32), ('b', tf.bool)],
placement_literals.CLIENTS))
def foo(x):
x.b = False
return x
with context_stack_impl.context_stack.install(
reference_executor.ReferenceExecutor()):
self.assertEqual(foo([[5, True], [0, False], [-5, True]]), [
structure.Struct([('a', 5), ('b', False)]),
structure.Struct([('a', 0), ('b', False)]),
structure.Struct([('a', -5), ('b', False)])
])
('sum_example_with_no_federated_secure_sum',
get_iterative_process_for_sum_example_with_no_federated_secure_sum()),
('sum_example_with_no_update',
get_iterative_process_for_sum_example_with_no_update()),
('sum_example_with_no_server_state',
get_iterative_process_for_sum_example_with_no_server_state()),
('minimal_sum_example',
get_iterative_process_for_minimal_sum_example()),
('example_with_unused_lambda_arg',
test_utils.get_iterative_process_for_example_with_unused_lambda_arg()),
('example_with_unused_tf_computation_arg',
test_utils.get_iterative_process_for_example_with_unused_tf_computation_arg()),
)
# pyformat: enable
def test_returns_canonical_form(self, ip):
cf = canonical_form_utils.get_canonical_form_for_iterative_process(ip)
self.assertIsInstance(cf, canonical_form.CanonicalForm)
def test_raises_value_error_for_sum_example_with_no_aggregation(self):
ip = get_iterative_process_for_sum_example_with_no_aggregation()
with self.assertRaises(ValueError):
canonical_form_utils.get_canonical_form_for_iterative_process(ip)
if __name__ == '__main__':
reference_executor = reference_executor.ReferenceExecutor()
default_executor.set_default_executor(reference_executor)
test.main()
int32_sequence = computation_types.SequenceType(tf.int32)
@computations.tf_computation(tf.int32, tf.int32)
def square_error(x, y):
return tf.pow(x - y, 2)
@computations.federated_computation(tf.int32, int32_sequence)
def sum_of_square_errors(x, y):
@computations.federated_computation(tf.int32)
def mapping_func(v):
return square_error(x, v)
return intrinsics.sequence_sum(intrinsics.sequence_map(mapping_func, y))
self.assertEqual(sum_of_square_errors(10, [11, 8, 13]), 14)
if __name__ == '__main__':
# We need to be able to individually test all components of the executor, and
# the compiler pipeline will potentially interfere with this process by
# performing reductions. Since this executor is intended to be the standards
# to compare against, it is the compiler pipeline that should get tested
# against this implementation, not the other way round.
executor_without_compiler = reference_executor.ReferenceExecutor()
# We must enable eager behavior as we cannot currently process Datasets in
# graph mode.
tf.compat.v1.enable_v2_behavior()
with context_stack_impl.context_stack.install(executor_without_compiler):
test.main()
def _make_default_context(stack):
return reference_executor.ReferenceExecutor(
compiler_pipeline.CompilerPipeline(stack))
def __init__(self):
super(ContextStackImpl, self).__init__()
self._stack = [
reference_executor.ReferenceExecutor(
compiler_pipeline.CompilerPipeline(self))
]
