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])