Update TF compat protos.

Author: arcra

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

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.