def test_different_input_output(self):
ex = sizing_executor.SizingExecutor(
eager_tf_executor.EagerTFExecutor())
tensor_type = computation_types.TensorType(tf.int32, 10)
@computations.tf_computation(tensor_type)
def return_constant(x):
del x
return tf.constant(0, tf.int32)
async def _make():
v1 = await ex.create_value(return_constant)
v2 = await ex.create_value([0, 1, 2, 3, 4, 5, 6, 7, 8, 9],
tensor_type)
v3 = await ex.create_call(v1, v2)
return await v3.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history, [[10, tf.int32]])
self.assertCountEqual(ex.aggregate_history, [[1, tf.int32]])
def test_executor_create_value_two_arg_computation(self):
ex = eager_tf_executor.EagerTFExecutor()
@computations.tf_computation(tf.int32, tf.int32)
def comp(a, b):
return a + b
comp_proto = computation_impl.ComputationImpl.get_proto(comp)
val = asyncio.get_event_loop().run_until_complete(
ex.create_value(
comp_proto,
computation_types.FunctionType([tf.int32, tf.int32],
tf.int32)))
self.assertIsInstance(val, eager_tf_executor.EagerValue)
self.assertEqual(str(val.type_signature), '(<int32,int32> -> int32)')
self.assertTrue(callable(val.internal_representation))
arg = anonymous_tuple.AnonymousTuple([('a', tf.constant(10)),
('b', tf.constant(10))])
result = val.internal_representation(arg)
self.assertIsInstance(result, tf.Tensor)
self.assertEqual(result.numpy(), 20)
def test_executor_service_value_unavailable_after_dispose(self):
env = TestEnv(eager_tf_executor.EagerTFExecutor())
value_proto, _ = executor_service_utils.serialize_value(
tf.constant(10.0).numpy(), tf.float32)
# Create the value
response = env.stub.CreateValue(
executor_pb2.CreateValueRequest(value=value_proto))
self.assertIsInstance(response, executor_pb2.CreateValueResponse)
value_id = str(response.value_ref.id)
# Check that the value appears in the _values map
env.get_value_future_directly(value_id)
# Dispose of the value
dispose_request = executor_pb2.DisposeRequest()
dispose_request.value_ref.append(response.value_ref)
response = env.stub.Dispose(dispose_request)
self.assertIsInstance(response, executor_pb2.DisposeResponse)
# Check that the value is gone from the _values map
# get_value_future_directly is used here so that we can catch the
# exception rather than having it occur on the GRPC thread.
with self.assertRaises(KeyError):
env.get_value_future_directly(value_id)
def create_executor(
self,
*,
cardinalities: Optional[executor_factory.CardinalitiesType] = None,
placement: Optional[placement_literals.PlacementLiteral] = None
) -> executor_base.Executor:
if cardinalities:
raise ValueError(
'Unplaced executors cannot accept nonempty cardinalities as '
'arguments. Received cardinalities: {}.'.format(cardinalities))
if placement == placement_literals.CLIENTS:
device = self._get_next_client_device()
elif placement == placement_literals.SERVER:
device = self._server_device
else:
device = None
eager_ex = eager_tf_executor.EagerTFExecutor(device=device)
return _wrap_executor_in_threading_stack(
eager_ex,
use_caching=self._use_caching,
can_resolve_references=self._can_resolve_references)
def test_simple(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
tensor_type = computation_types.TensorType(tf.int32, 10)
@computations.tf_computation(tensor_type)
def add_one(x):
return tf.add(x, 1)
async def _make():
v1 = await ex.create_value(add_one)
v2 = await ex.create_value(
tf.constant([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], tf.int32), tensor_type)
v3 = await ex.create_call(v1, v2)
v4 = await ex.create_struct(structure.Struct([('foo', v3)]))
v5 = await ex.create_selection(v4, 0)
return await v5.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history, [[10, tf.int32]])
self.assertCountEqual(ex.aggregate_history, [[10, tf.int32]])
def test_executor_service_create_and_select_from_tuple(self):
ex_factory = executor_stacks.ResourceManagingExecutorFactory(
lambda _: eager_tf_executor.EagerTFExecutor())
env = TestEnv(ex_factory)
value_proto, _ = executor_serialization.serialize_value(10, tf.int32)
response = env.stub.CreateValue(
executor_pb2.CreateValueRequest(value=value_proto))
self.assertIsInstance(response, executor_pb2.CreateValueResponse)
ten_ref = response.value_ref
self.assertEqual(env.get_value(ten_ref.id), 10)
value_proto, _ = executor_serialization.serialize_value(20, tf.int32)
response = env.stub.CreateValue(
executor_pb2.CreateValueRequest(value=value_proto))
self.assertIsInstance(response, executor_pb2.CreateValueResponse)
twenty_ref = response.value_ref
self.assertEqual(env.get_value(twenty_ref.id), 20)
response = env.stub.CreateStruct(
executor_pb2.CreateStructRequest(element=[
executor_pb2.CreateStructRequest.Element(
name='a', value_ref=ten_ref),
executor_pb2.CreateStructRequest.Element(
name='b', value_ref=twenty_ref)
]))
self.assertIsInstance(response, executor_pb2.CreateStructResponse)
tuple_ref = response.value_ref
self.assertEqual(str(env.get_value(tuple_ref.id)), '<a=10,b=20>')
for arg_name, arg_val, result_val in [('name', 'a', 10), ('name', 'b', 20),
('index', 0, 10), ('index', 1, 20)]:
response = env.stub.CreateSelection(
executor_pb2.CreateSelectionRequest(
source_ref=tuple_ref, **{arg_name: arg_val}))
self.assertIsInstance(response, executor_pb2.CreateSelectionResponse)
selection_ref = response.value_ref
self.assertEqual(env.get_value(selection_ref.id), result_val)
del env
def test_with_federated_map(self):
eager_ex = eager_tf_executor.EagerTFExecutor()
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory(
{placement_literals.SERVER: eager_ex})
federated_ex = federating_executor.FederatingExecutor(factory, eager_ex)
ex = reference_resolving_executor.ReferenceResolvingExecutor(federated_ex)
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
@computations.federated_computation(type_factory.at_server(tf.int32))
def comp(x):
return intrinsics.federated_map(add_one, x)
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(
ex.create_value(10, type_factory.at_server(tf.int32)))
v3 = loop.run_until_complete(ex.create_call(v1, v2))
result = loop.run_until_complete(v3.compute())
self.assertEqual(result.numpy(), 11)
def test_with_one_arg_tf_comp_in_two_arg_fed_comp(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
@tensorflow_computation.tf_computation(tf.int32, tf.int32)
def add_numbers(x, y):
return x + y
@federated_computation.federated_computation(tf.int32, tf.int32)
def comp(x, y):
return add_numbers(x, x), add_numbers(x, y), add_numbers(y, y)
v1 = asyncio.run(ex.create_value(comp))
v2 = asyncio.run(ex.create_value(10, tf.int32))
v3 = asyncio.run(ex.create_value(20, tf.int32))
v4 = asyncio.run(
ex.create_struct(structure.Struct([('x', v2), ('y', v3)])))
v5 = asyncio.run(ex.create_call(v1, v4))
result = asyncio.run(v5.compute())
self.assertEqual(
str(structure.map_structure(lambda x: x.numpy(), result)),
'<20,30,40>')
def test_with_nested_lambdas(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32, tf.int32)
def add_numbers(x, y):
return x + y
@computations.federated_computation(tf.int32)
def comp(x):
@computations.federated_computation(tf.int32)
def nested_comp(y):
return add_numbers(x, y)
return nested_comp(1)
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(ex.create_value(10, tf.int32))
v3 = loop.run_until_complete(ex.create_call(v1, v2))
result = loop.run_until_complete(v3.compute())
self.assertEqual(result.numpy(), 11)
def test_with_functional_parameter(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32)
def add_one(x):
return x + 1
@computations.federated_computation(
computation_types.FunctionType(tf.int32, tf.int32), tf.int32)
def comp(f, x):
return f(f(x))
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(ex.create_value(add_one))
v3 = loop.run_until_complete(ex.create_value(10, tf.int32))
v4 = loop.run_until_complete(
ex.create_struct(structure.Struct([('f', v2), ('x', v3)])))
v5 = loop.run_until_complete(ex.create_call(v1, v4))
result = loop.run_until_complete(v5.compute())
self.assertEqual(result.numpy(), 12)
def test_executor_service_create_no_arg_computation_value_and_call(self):
ex_factory = executor_test_utils.BasicTestExFactory(
eager_tf_executor.EagerTFExecutor())
env = TestEnv(ex_factory)
@tensorflow_computation.tf_computation
def comp():
return tf.constant(10)
value_proto, _ = value_serialization.serialize_value(comp)
response = env.stub.CreateValue(
executor_pb2.CreateValueRequest(executor=env.executor_pb,
value=value_proto))
self.assertIsInstance(response, executor_pb2.CreateValueResponse)
response = env.stub.CreateCall(
executor_pb2.CreateCallRequest(executor=env.executor_pb,
function_ref=response.value_ref))
self.assertIsInstance(response, executor_pb2.CreateCallResponse)
value_id = str(response.value_ref.id)
value = env.get_value(value_id)
self.assertEqual(value, 10)
del env
def test_with_one_arg_tf_comp_in_two_arg_fed_comp(self):
ex = reference_resolving_executor.ReferenceResolvingExecutor(
eager_tf_executor.EagerTFExecutor())
loop = asyncio.get_event_loop()
@computations.tf_computation(tf.int32, tf.int32)
def add_numbers(x, y):
return x + y
@computations.federated_computation(tf.int32, tf.int32)
def comp(x, y):
return add_numbers(x, x), add_numbers(x, y), add_numbers(y, y)
v1 = loop.run_until_complete(ex.create_value(comp))
v2 = loop.run_until_complete(ex.create_value(10, tf.int32))
v3 = loop.run_until_complete(ex.create_value(20, tf.int32))
v4 = loop.run_until_complete(
ex.create_struct(structure.Struct([(None, v2), (None, v3)])))
v5 = loop.run_until_complete(ex.create_call(v1, v4))
result = loop.run_until_complete(v5.compute())
self.assertEqual(
str(structure.map_structure(lambda x: x.numpy(), result)),
'<20,30,40>')
def test_executor_create_call_take_two_from_stateful_dataset(self):
vocab = ['a', 'b', 'c', 'd', 'e', 'f']
@tensorflow_computation.tf_computation(
computation_types.SequenceType(tf.string))
def comp(ds):
table = tf.lookup.StaticVocabularyTable(
tf.lookup.KeyValueTensorInitializer(
vocab, tf.range(len(vocab), dtype=tf.int64)),
num_oov_buckets=1)
ds = ds.map(table.lookup)
return ds.take(2)
ds = tf.data.Dataset.from_tensor_slices(vocab)
ex = eager_tf_executor.EagerTFExecutor()
comp = asyncio.run(ex.create_value(comp))
arg = asyncio.run(ex.create_value(ds, comp.type_signature.parameter))
result = asyncio.run(ex.create_call(comp, arg))
self.assertIsInstance(result, eager_tf_executor.EagerValue)
self.assertEqual(str(result.type_signature), 'int64*')
self.assertIn('Dataset', type(result.internal_representation).__name__)
self.assertCountEqual(
[x.numpy() for x in result.internal_representation], [0, 1])
def test_multiple_inputs(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
int_type = computation_types.TensorType(tf.int32, 10)
float_type = computation_types.TensorType(tf.float64, 10)
@computations.tf_computation(float_type, int_type)
def add(x, y):
x = tf.cast(x, tf.int64)
y = tf.cast(y, tf.int64)
return x + y
async def _make():
v1 = await ex.create_value(add)
v2 = await ex.create_value([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], float_type)
v3 = await ex.create_value([0, 1, 2, 3, 4, 5, 6, 7, 8, 9], int_type)
v4 = await ex.create_struct(structure.Struct([(None, v2), (None, v3)]))
v5 = await ex.create_call(v1, v4)
return await v5.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history,
[[10, tf.int32], [10, tf.float64]])
self.assertCountEqual(ex.aggregate_history, [[10, tf.int64]])
def test_nested_tuple(self):
ex = sizing_executor.SizingExecutor(eager_tf_executor.EagerTFExecutor())
a = computation_types.TensorType(tf.int32, [4])
b = computation_types.TensorType(tf.bool, [2])
c = computation_types.TensorType(tf.int64, [2, 3])
inner_type = computation_types.StructType([('a', a), ('b', b), ('c', c)])
outer_type = computation_types.StructType([('a', inner_type),
('b', inner_type)])
inner_type_val = collections.OrderedDict()
inner_type_val['a'] = [0, 1, 2, 3]
inner_type_val['b'] = [True, False]
inner_type_val['c'] = [[1, 2, 3], [4, 5, 6]]
outer_type_val = collections.OrderedDict()
outer_type_val['a'] = inner_type_val
outer_type_val['b'] = inner_type_val
async def _make():
v1 = await ex.create_value(outer_type_val, outer_type)
return await v1.compute()
asyncio.get_event_loop().run_until_complete(_make())
self.assertCountEqual(ex.broadcast_history,
[[4, tf.int32], [2, tf.bool], [6, tf.int64],
[4, tf.int32], [2, tf.bool], [6, tf.int64]])
def test_executor_create_call_with_dataset_of_tuples(self):
element = collections.namedtuple('_', 'a b')
@tensorflow_computation.tf_computation(
computation_types.SequenceType(element(tf.int32, tf.int32)))
def comp(ds):
return ds.reduce(element(np.int32(0), np.int32(0)),
lambda p, q: element(p.a + q.a, p.b + q.b))
ds = tf.data.Dataset.from_tensor_slices(
element([10, 20, 30], [4, 5, 6]))
ex = eager_tf_executor.EagerTFExecutor()
comp = asyncio.run(ex.create_value(comp))
arg = asyncio.run(ex.create_value(ds, comp.type_signature.parameter))
result = asyncio.run(ex.create_call(comp, arg))
self.assertIsInstance(result, eager_tf_executor.EagerValue)
self.assertEqual(str(result.type_signature), '<a=int32,b=int32>')
self.assertIsInstance(result.internal_representation, structure.Struct)
self.assertCountEqual(dir(result.internal_representation), ['a', 'b'])
self.assertIsInstance(result.internal_representation.a, tf.Tensor)
self.assertIsInstance(result.internal_representation.b, tf.Tensor)
self.assertEqual(result.internal_representation.a, 60)
self.assertEqual(result.internal_representation.b, 15)
def test_with_federated_map(self):
eager_ex = eager_tf_executor.EagerTFExecutor()
factory = federated_resolving_strategy.FederatedResolvingStrategy.factory(
{placements.SERVER: eager_ex})
federated_ex = federating_executor.FederatingExecutor(
factory, eager_ex)
ex = reference_resolving_executor.ReferenceResolvingExecutor(
federated_ex)
@tensorflow_computation.tf_computation(tf.int32)
def add_one(x):
return x + 1
@federated_computation.federated_computation(
computation_types.at_server(tf.int32))
def comp(x):
return intrinsics.federated_map(add_one, x)
v1 = asyncio.run(ex.create_value(comp))
v2 = asyncio.run(
ex.create_value(10, computation_types.at_server(tf.int32)))
v3 = asyncio.run(ex.create_call(v1, v2))
result = asyncio.run(v3.compute())
self.assertEqual(result.numpy(), 11)
def test_executor_create_value_two_arg_computation(self):
ex = eager_tf_executor.EagerTFExecutor()
@tensorflow_computation.tf_computation(tf.int32, tf.int32)
def comp(a, b):
return a + b
comp_proto = computation_impl.ConcreteComputation.get_proto(comp)
val = asyncio.run(
ex.create_value(
comp_proto,
computation_types.FunctionType(
computation_types.StructType([('a', tf.int32),
('b', tf.int32)]),
tf.int32)))
self.assertIsInstance(val, eager_tf_executor.EagerValue)
self.assertEqual(str(val.type_signature),
'(<a=int32,b=int32> -> int32)')
self.assertTrue(callable(val.internal_representation))
arg = structure.Struct([('a', tf.constant(10)),
('b', tf.constant(10))])
result = val.internal_representation(arg)
self.assertIsInstance(result, tf.Tensor)
self.assertEqual(result, 20)
def create_test_executor_factory():
executor = eager_tf_executor.EagerTFExecutor()
executor = reference_resolving_executor.ReferenceResolvingExecutor(
executor)
return executor_test_utils.BasicTestExFactory(executor)
def _create_bottom_stack(device=None):
return _complete_stack(eager_tf_executor.EagerTFExecutor(device=device))
def _stack_fn(x): del x # Unused return eager_tf_executor.EagerTFExecutor()
def create_test_executor_factory():
executor = eager_tf_executor.EagerTFExecutor()
executor = caching_executor.CachingExecutor(executor)
executor = reference_resolving_executor.ReferenceResolvingExecutor(
executor)
return executor_factory.ExecutorFactoryImpl(lambda _: executor)
def _create_bottom_stack():
return reference_resolving_executor.ReferenceResolvingExecutor(
caching_executor.CachingExecutor(
thread_delegating_executor.ThreadDelegatingExecutor(
eager_tf_executor.EagerTFExecutor())))
def ex_fn(device):
return data_executor.DataExecutor(
eager_tf_executor.EagerTFExecutor(device),
TestDataBackend(data_constant))
def test_clear_executor_calls_cleanup(self, mock_cleanup):
ex_factory = executor_stacks.ResourceManagingExecutorFactory(
lambda _: eager_tf_executor.EagerTFExecutor())
env = TestEnv(ex_factory)
env.stub.ClearExecutor(executor_pb2.ClearExecutorRequest())
mock_cleanup.assert_called_once()
def test_executor_construction_with_no_device_name(self):
eager_tf_executor.EagerTFExecutor()
def create_bottom_stack():
executor = eager_tf_executor.EagerTFExecutor()
return reference_resolving_executor.ReferenceResolvingExecutor(
executor)
def test_executor_constructor_fails_if_not_in_eager_mode(self):
with tf.Graph().as_default():
with self.assertRaises(RuntimeError):
eager_tf_executor.EagerTFExecutor()
def _make_executor_and_tracer_for_test():
tracer = executor_test_utils.TracingExecutor(
eager_tf_executor.EagerTFExecutor())
ex = caching_executor.CachingExecutor(tracer)
return ex, tracer
def test_executor_construction_with_correct_device_name(self):
eager_tf_executor.EagerTFExecutor(
tf.config.list_logical_devices('CPU')[0])
