Update TF compat protos

tensorboard/compat/proto/config.proto

@@ -265,6 +265,43 @@ message GPUOptions {
     // system memory size for better resource estimation of multi-tenancy(one
     // gpu with multiple model) use case.
     int32 gpu_system_memory_size_in_mb = 16;
+
+    // If true, save information needed for created a PjRt GPU client for
+    // creating a client with remote devices.
+    bool populate_pjrt_gpu_client_creation_info = 17;
+
+    // node_id for use when creating a PjRt GPU client with remote devices,
+    // which enumerates jobs*tasks from a ServerDef.
+    int32 node_id = 18;
+
+    // Whether to merge data transfer streams into the compute stream in the
+    // same stream group. Stream merging helps reduce the overhead caused by
+    // stream synchronization, especially when data transfers are frequent. For
+    // example, setting "merge_host_to_device_stream = true" will make the
+    // compute stream responsible for both computation and host to device memory
+    // copy.
+    message StreamMergeOptions {
+      // If true, the compute stream will be used for host_to_device copy as
+      // well. It's no longer necessary to record an event before the copy to
+      // let the copy stream wait for the compute stream to finish. There is
+      // also no need to wait for the copy to complete before executing the
+      // callback function.
+      bool merge_host_to_device_stream = 1;
+
+      // If true, the compute stream will be used for device_to_host copy as
+      // well. It's no longer necessary to record an event before the copy to
+      // let the copy stream wait for the compute stream to finish.
+      bool merge_device_to_host_stream = 2;
+
+      // If true, the compute stream will be used for device_to_device copy as
+      // well. It's no longer necessary to record an event before the copy to
+      // let the copy stream wait for the compute stream of the sending device
+      // to finish. There is also no need to wait for the compute stream of the
+      // receiving device to finish if the copy is within the same device.
+      bool merge_device_to_device_stream = 3;
+    }
+
+    StreamMergeOptions stream_merge_options = 19;
   }
 
   // Everything inside experimental is subject to change and is not subject
@@ -499,6 +536,9 @@ message ConfigProto {
   // Options that apply to all GPUs.
   GPUOptions gpu_options = 6;
 
+  // Options that apply to pluggable devices.
+  GPUOptions pluggable_device_options = 18;
+
   // Whether soft placement is allowed. If allow_soft_placement is true,
   // an op will be placed on CPU if
   //   1. there's no GPU implementation for the OP
@@ -676,6 +716,10 @@ message ConfigProto {
     // If true, use Pathways with TFRT API for multi host support.
     bool enable_multi_host = 27;
 
+    // If true, use ifrt as the backend for TFRT. This is only used when
+    // `use_tfrt` is true.
+    bool tfrt_use_ifrt = 32;
+
     // Port for the Pathways server. Ignored if enable_multi_host=false.
     int32 backend_server_port = 28;
 
@@ -730,12 +774,12 @@ message ConfigProto {
 
     reserved 25;
 
-    // Next: 32
+    // Next: 33
   }
 
   Experimental experimental = 16;
 
-  // Next: 18
+  // Next: 19
 }
 
 // Options for a single Run() call.

tensorboard/compat/proto/coordination_config.proto

@@ -67,4 +67,8 @@ message CoordinationServiceConfig {
   // not specify any config. This field allows users to explicitly disable
   // coordination service under all situations.
   bool force_disable = 12;
+
+  // Use long polling to get error from coordination service as the error
+  // propagation mechanism.
+  bool poll_for_error_from_service_at_startup = 13;
 }

tensorboard/compat/proto/meta_graph.proto

@@ -257,74 +257,18 @@ message TensorInfo {
 
 // SignatureDef defines the signature of a computation supported by a TensorFlow
 // graph.
-//
-// For example, a model with two loss computations, sharing a single input,
-// might have the following signature_def map, in a MetaGraphDef message.
-//
-// Note that across the two SignatureDefs "loss_A" and "loss_B", the input key,
-// output key, and method_name are identical, and will be used by system(s) that
-// implement or rely upon this particular loss method. The output tensor names
-// differ, demonstrating how different outputs can exist for the same method.
-//
-// signature_def {
-//   key: "loss_A"
-//   value {
-//     inputs {
-//       key: "input"
-//       value {
-//         name: "input:0"
-//         dtype: DT_STRING
-//         tensor_shape: ...
-//       }
-//     }
-//     outputs {
-//       key: "loss_output"
-//       value {
-//         name: "loss_output_A:0"
-//         dtype: DT_FLOAT
-//         tensor_shape: ...
-//       }
-//     }
-//     method_name: "some/package/compute_loss"
-//   }
-//   ...
-// }
-// signature_def {
-//   key: "loss_B"
-//   value {
-//     inputs {
-//       key: "input"
-//       value {
-//         name: "input:0"
-//         dtype: DT_STRING
-//         tensor_shape: ...
-//       }
-//     }
-//     outputs {
-//       key: "loss_output"
-//       value {
-//         name: "loss_output_B:0"
-//         dtype: DT_FLOAT
-//         tensor_shape: ...
-//       }
-//     }
-//     method_name: "some/package/compute_loss"
-//   }
-//   ...
-// }
 message SignatureDef {
   // Named input parameters.
   map<string, TensorInfo> inputs = 1;
   // Named output parameters.
   map<string, TensorInfo> outputs = 2;
-  // Extensible method_name information enabling third-party users to mark a
-  // SignatureDef as supporting a particular method. This enables producers and
-  // consumers of SignatureDefs, e.g. a model definition library and a serving
-  // library to have a clear hand-off regarding the semantics of a computation.
+  // Deprecated: TensorFlow 2 always sets this to a fixed value;
+  // open-source TF Serving stopped checking by default since release 2.4.
   //
-  // Note that multiple SignatureDefs in a single MetaGraphDef may have the same
-  // method_name. This is commonly used to support multi-headed computation,
-  // where a single graph computation may return multiple results.
+  // In TensorFlow 1, the method_name enabled users to mark a SignatureDef as
+  // supporting a particular method. Multiple SignatureDefs in a single
+  // MetaGraphDef could have the same method_name (e.g., to support multi-headed
+  // computation).
   string method_name = 3;
   // Named input to corresponding default values if any.
   map<string, TensorProto> defaults = 4;

tensorboard/compat/proto/resource_handle.proto

@@ -34,7 +34,9 @@ message ResourceHandleProto {
 
   // Protocol buffer representing a pair of (data type, tensor shape).
   message DtypeAndShape {
+    // Data type of the tensor.
     DataType dtype = 1;
+    // Shape of the tensor.
     TensorShapeProto shape = 2;
   }
 

tensorboard/compat/proto/rewriter_config.proto

@@ -102,8 +102,8 @@ message RewriterConfig {
   // Enable the swap of kernel implementations based on the device placement
   // (default is ON).
   Toggle implementation_selector = 22;
-  // Optimize data types for CUDA (default is OFF).
-  // This will try to use float16 on GPU which is faster.
+  // Optimize data types for CUDA/oneDNN (default is OFF).
+  // This will try to use float16 on GPU/CPU which is faster.
   // Note that this can change the numerical stability of the graph and may
   // require the use of loss scaling to maintain model convergence.
   Toggle auto_mixed_precision = 23;

tensorboard/compat/proto/tensor.proto

@@ -14,6 +14,7 @@ option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/framewo
 
 // Protocol buffer representing a tensor.
 message TensorProto {
+  // Data type of the tensor.
   DataType dtype = 1;
 
   // Shape of the tensor.  TODO(touts): sort out the 0-rank issues.