Support shapes for remote eager tensor handles.

Since we respond with the shape, all RPCs will happen sync (note
that we may still hide the python overhead, since the op is still scheduled for
execution via the eager executor).

PiperOrigin-RevId: 202207324

1 files changed, 7 insertions, 0 deletions

diff --git a/tensorflow/core/protobuf/eager_service.proto b/tensorflow/core/protobuf/eager_service.proto
index 50294b8a42..5b05a1b3ee 100644
--- a/tensorflow/core/protobuf/eager_service.proto
+++ b/tensorflow/core/protobuf/eager_service.proto
@@ -7,6 +7,7 @@ import "tensorflow/core/framework/device_attributes.proto";
 import "tensorflow/core/framework/function.proto";
 import "tensorflow/core/framework/versions.proto";
 import "tensorflow/core/protobuf/tensorflow_server.proto";
+import "tensorflow/core/framework/tensor_shape.proto";
 
 message RemoteTensorHandle {
   // The ID of the operation that produced this tensor.
@@ -45,6 +46,10 @@ message QueueItem {
   }
 }
 
+message QueueResponse {
+  repeated TensorShapeProto shape = 1;
+}
+
 message CreateContextRequest {
   // Identifies the full cluster, and this particular worker's position within.
   ServerDef server_def = 1;
@@ -84,6 +89,8 @@ message EnqueueRequest {
 }
 
 message EnqueueResponse {
+  // A single operation response for every item in the request.
+  repeated QueueResponse queue_response = 1;
 }
 
 message WaitQueueDoneRequest {