Chore: Update protos to match latest TF proto's (#6010)

* Motivation for features / changes
This PR updates our proto files to match the latest Tensorflow protos.

For posterity these protos were synced to Tensorflow at commit
20d085686bb99daefa930205e8ec69a99dae306d. Which was commited on November
1st 2022.

* Technical description of changes
These updates were done by running the
tensorboard/compat/proto/update.sh script. The update to that file was
to fix a little bug in the changes made in #5997

Author: JamesHollyer

tensorboard/compat/proto/BUILD

@@ -68,6 +68,7 @@ tb_proto_library(
         ":debug",
         ":graph",
         ":rewriter_config",
+        ":rpc_options",
         ":step_stats",
     ],
 )
@@ -134,6 +135,11 @@ tb_proto_library(
     ],
 )
 
+tb_proto_library(
+    name = "histogram",
+    srcs = ["histogram.proto"],
+)
+
 tb_proto_library(
     name = "meta_graph",
     srcs = ["meta_graph.proto"],
@@ -186,6 +192,11 @@ tb_proto_library(
     ],
 )
 
+tb_proto_library(
+    name = "rpc_options",
+    srcs = ["rpc_options.proto"],
+)
+
 tb_proto_library(
     name = "saved_object_graph",
     srcs = ["saved_object_graph.proto"],
@@ -227,6 +238,7 @@ tb_proto_library(
     name = "summary",
     srcs = ["summary.proto"],
     deps = [
+        ":histogram",
         ":tensor",
     ],
 )
@@ -310,11 +322,13 @@ tb_proto_library(
         ":full_type",
         ":function",
         ":graph",
+        ":histogram",
         ":meta_graph",
         ":node_def",
         ":op_def",
         ":resource_handle",
         ":rewriter_config",
+        ":rpc_options",
         ":saved_object_graph",
         ":saver",
         ":step_stats",

tensorboard/compat/proto/config.proto

@@ -6,9 +6,10 @@ import "tensorboard/compat/proto/cost_graph.proto";
 import "tensorboard/compat/proto/graph.proto";
 import "tensorboard/compat/proto/step_stats.proto";
 import "tensorboard/compat/proto/cluster.proto";
-import "tensorboard/compat/proto/coordination_config.proto";
 import "tensorboard/compat/proto/debug.proto";
 import "tensorboard/compat/proto/rewriter_config.proto";
+import "tensorboard/compat/proto/rpc_options.proto";
+import "tensorboard/compat/proto/coordination_config.proto";
 
 option cc_enable_arenas = true;
 option java_outer_classname = "ConfigProtos";
@@ -123,6 +124,12 @@ message GPUOptions {
       // created with the default. If this field has values set, then the size
       // of this must match with the above memory_limit_mb.
       repeated int32 priority = 2;
+
+      // Virtual Device ordinal number determines the device ID of the device.
+      // A Virtual device with a lower ordinal number always receives the a
+      // smaller device id. The phyiscal device id and location in the
+      // virtual device list is used to break ties.
+      repeated int32 device_ordinal = 3;
     }
 
     // The multi virtual device settings. If empty (not set), it will create
@@ -132,15 +139,30 @@ message GPUOptions {
     // "visible_device_list" filtering if it is set), and the string represented
     // device names (e.g. /device:GPU:<id>) will refer to the virtual
     // devices and have the <id> field assigned sequentially starting from 0,
-    // according to the order they appear in this list and the "memory_limit"
-    // list inside each element. For example,
+    // according to the order of the virtual devices determined by
+    // device_ordinal and the location in the virtual device list.
+    //
+    // For example,
     //   visible_device_list = "1,0"
     //   virtual_devices { memory_limit: 1GB memory_limit: 2GB }
-    //   virtual_devices {}
-    // will create three virtual devices as:
+    //   virtual_devices { memory_limit: 3GB memory_limit: 4GB }
+    // will create 4 virtual devices as:
     //   /device:GPU:0 -> visible GPU 1 with 1GB memory
     //   /device:GPU:1 -> visible GPU 1 with 2GB memory
-    //   /device:GPU:2 -> visible GPU 0 with all available memory
+    //   /device:GPU:2 -> visible GPU 0 with 3GB memory
+    //   /device:GPU:3 -> visible GPU 0 with 4GB memory
+    //
+    // but
+    //   visible_device_list = "1,0"
+    //   virtual_devices { memory_limit: 1GB memory_limit: 2GB
+    //                     device_ordinal: 10 device_ordinal: 20}
+    //   virtual_devices { memory_limit: 3GB memory_limit: 4GB
+    //                     device_ordinal: 10 device_ordinal: 20}
+    // will create 4 virtual devices as:
+    //   /device:GPU:0 -> visible GPU 1 with 1GB memory  (ordinal 10)
+    //   /device:GPU:1 -> visible GPU 0 with 3GB memory  (ordinal 10)
+    //   /device:GPU:2 -> visible GPU 1 with 2GB memory  (ordinal 20)
+    //   /device:GPU:3 -> visible GPU 0 with 4GB memory  (ordinal 20)
     //
     // NOTE:
     // 1. It's invalid to set both this and "per_process_gpu_memory_fraction"
@@ -355,41 +377,6 @@ message ThreadPoolOptionProto {
   string global_name = 2;
 }
 
-message RPCOptions {
-  // If true, always use RPC to contact the session target.
-  //
-  // If false (the default option), TensorFlow may use an optimized
-  // transport for client-master communication that avoids the RPC
-  // stack. This option is primarily for used testing the RPC stack.
-  bool use_rpc_for_inprocess_master = 1;
-
-  // The compression algorithm to be used. One of "deflate", "gzip".
-  string compression_algorithm = 2;
-
-  // If compression_algorithm is set, the compression level to be used.
-  // From 0 (no compression), up to 3.
-  int32 compression_level = 3;
-
-  // Setting cache_rpc_response to true will enable sender side caching of
-  // response for RecvTensorAsync and RecvBufAsync to allow receiver to retry
-  // requests . This is only necessary when the network fabric is experiencing a
-  // significant error rate.  Without it we'll fail a step on an network error,
-  // while with it we'll be able to complete long steps (like complex
-  // initializations) in the face of some network errors during RecvTensor.
-  bool cache_rpc_response = 4;
-
-  // Disables TCP connection sharing when opening a new RPC channel.
-  bool disable_session_connection_sharing = 5;
-
-  // Setting num_channels_per_target > 0 allows uses of multiple channels to
-  // communicate to the same target. This can be used to improve the aggregate
-  // throughput on high speed links (e.g 100G) where single connection is not
-  // sufficient to maximize link utilization. Note that a single RPC only goes
-  // on a single channel, this only helps in situations where there are multiple
-  // transfers to the same target overlapping in time.
-  int32 num_channels_per_target = 6;
-}
-
 // Metadata about the session.
 //
 // This can be used by the runtime and the Ops for debugging, monitoring, etc.
@@ -606,13 +593,10 @@ message ConfigProto {
     // Session::Extend() may not be supported.
     bool optimize_for_static_graph = 12;
 
-    // This field will eventually be deprecated and replaced by
-    // mlir_bridge_rollout (b/166038521).
-    //
-    // Whether to enable the MLIR-based TF->XLA bridge.
+    // Whether to enable the MLIR-based TF->XLA bridge. This is only used if set
+    // to true. Default value or false is ignored. Use mlir_bridge_rollout for
+    // finer control.
     //
-    // This is a replacement to the existing bridge, and not ready for
-    // production usage yet.
     // If this option is set to true when a session is created, MLIR is used to
     // perform the set of graph transformations to put the graph in a form that
     // can be executed with delegation of some computations to an accelerator.
@@ -644,9 +628,6 @@ message ConfigProto {
       // to process the TF graph.
       MLIR_BRIDGE_ROLLOUT_SAFE_MODE_FALLBACK_ENABLED = 4;
     }
-    // This field is underdevelopment, for now use enable_mlir_bridge
-    // (b/166038521).
-    //
     // Whether to enable the MLIR-based TF->XLA bridge.
     MlirBridgeRollout mlir_bridge_rollout = 17;
 
@@ -798,6 +779,9 @@ message RunMetadata {
   // level idea of what the built graph looks like (since the various graph
   // optimization passes might change the structure of the graph significantly).
   repeated FunctionGraphs function_graphs = 4;
+
+  // Metadata about the session.
+  SessionMetadata session_metadata = 5;
 }
 
 // Defines a connection between two tensors in a `GraphDef`.

tensorboard/compat/proto/coordination_config.proto

@@ -4,6 +4,13 @@ package tensorboard;
 
 option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
 
+// Represents a job type and the number of tasks under this job.
+// For example, ("worker", 20) implies that there will be 20 worker tasks.
+message CoordinatedJob {
+  string name = 1;
+  int32 num_tasks = 2;
+}
+
 // Coordination service configuration parameters.
 // The system picks appropriate values for fields that are not set.
 message CoordinationServiceConfig {
@@ -29,9 +36,9 @@ message CoordinationServiceConfig {
   // recording the state change and the agent stopping heartbeats.
   int64 heartbeat_timeout_in_ms = 5;
 
-  // The list of jobs that partipate in the coordination service. If empty, all
-  // jobs will be included in the coordination service by default.
-  repeated string coordinated_jobs = 6;
+  // The list of `CoordinatedJob`s that will register in coordination service.
+  reserved 6;
+  repeated CoordinatedJob coordinated_job_list = 10;
 
   // Denotes how long to wait for all coordination agents to reach the barriers
   // (after the first shutdown request) before disconnecting together. If
@@ -43,4 +50,10 @@ message CoordinationServiceConfig {
   // find out about the disconnecte agent via stale heartbeats. Used for
   // testing.
   bool agent_destruction_without_shutdown = 8;
+
+  // The list of jobs which are recoverable. If a task in this list fails,
+  // it will not propagate error to other tasks.
+  // If empty, no jobs will be recoverable and every task failure will cause
+  // error propagation to other tasks.
+  repeated string recoverable_jobs = 9;
 }

tensorboard/compat/proto/full_type.proto

@@ -174,6 +174,18 @@ enum FullTypeId {
   //   TFT_LITERAL[TFT_INT32]{1} is the compile-time constant 1.
   TFT_LITERAL = 1003;
 
+  // Encoding types describe a value of a certain type, encoded as a different
+  // type.
+  //
+  // Parametrization:
+  //   TFT_ENCODED[<encoded type>, <encoding type>]
+  //   * <encoded type> may be any type
+  //   * <encoding type> may be any type
+  //
+  // Examples:
+  //   TFT_ENCODING[TFT_INT32, TFT_STRING] is an integer encoded as string.
+  TFT_ENCODED = 1004;
+
   // Type attributes. These always appear in the parametrization of a type,
   // never alone. For example, there is no such thing as a "bool" TensorFlow
   // object (for now).
@@ -217,7 +229,7 @@ enum FullTypeId {
   // is expressed using TFT_PRODUCT.
   //
   //
-  // Parametrization: TFT_ARRAY[<element type>].
+  // Parametrization: TFT_DATASET[<element type>].
   //   * <element type> may be a concrete type or a type symbol. It represents
   //     the data type of the elements produced by the dataset.
   //
@@ -235,6 +247,15 @@ enum FullTypeId {
   // Parametrization: TFT_RAGGED[<element_type>].
   TFT_RAGGED = 10103;
 
+  // Iterators created by tf.data ops and APIs. Very similar to Datasets, except
+  // they are mutable.
+  //
+  //
+  // Parametrization: TFT_ITERATOR[<element type>].
+  //   * <element type> may be a concrete type or a type symbol. It represents
+  //     the data type of the elements produced by the dataset.
+  TFT_ITERATOR = 10104;
+
   // A mutex lock tensor, produced by tf.raw_ops.MutexLock.
   // Unlike strict execution models, where ownership of a lock is denoted by
   // "running after the lock has been acquired", in non-strict mode, lock
@@ -265,7 +286,7 @@ message FullTypeDef {
 
   repeated FullTypeDef args = 2;
 
-  // Literal values of this type object, if the the type admits one.
+  // Literal values of this type object, if the type admits one.
   // For example, a type variable admits a string attribute - its name.
   // Shape-related types may admit int attributes - their static shape values.
   // Fields for more data types to be added as needed.

tensorboard/compat/proto/histogram.proto

@@ -0,0 +1,26 @@
+syntax = "proto3";
+
+package tensorboard;
+
+option cc_enable_arenas = true;
+option java_multiple_files = true;
+option java_package = "org.tensorflow.framework";
+option go_package = "github.com/google/tsl/tsl/go/core/protobuf/summary_go_proto";
+
+// Serialization format for histogram module in
+// tsl/lib/histogram/histogram.h
+message HistogramProto {
+  double min = 1;
+  double max = 2;
+  double num = 3;
+  double sum = 4;
+  double sum_squares = 5;
+
+  // Parallel arrays encoding the bucket boundaries and the bucket values.
+  // bucket(i) is the count for the bucket i.  The range for
+  // a bucket is:
+  //   i == 0:  -DBL_MAX .. bucket_limit(0)
+  //   i != 0:  bucket_limit(i-1) .. bucket_limit(i)
+  repeated double bucket_limit = 6 [packed = true];
+  repeated double bucket = 7 [packed = true];
+}

tensorboard/compat/proto/meta_graph.proto

@@ -17,9 +17,6 @@ option java_multiple_files = true;
 option java_package = "org.tensorflow.framework";
 option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
 
-// NOTE: This protocol buffer is evolving, and will go through revisions in the
-// coming months.
-//
 // Protocol buffer containing the following which are necessary to restart
 // training, run inference. It can be used to serialize/de-serialize memory
 // objects necessary for running computation in a graph when crossing the

tensorboard/compat/proto/rewriter_config.proto

@@ -107,10 +107,17 @@ message RewriterConfig {
   // 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;
-  // Optimize data types for MKL (default is OFF).
+  // Optimize data types for oneDNN (default is OFF).
   // This will try to use bfloat16 on CPUs, which is faster.
   // Note that this can change the numerical stability of the graph.
+  // Note: this is deprecated.
+  // It is replaced by auto_mixed_precision_onednn_bfloat16
   Toggle auto_mixed_precision_mkl = 25;
+  // Optimize data types for oneDNN (default is OFF).
+  // This will try to use bfloat16 on CPUs, which is faster.
+  // Note that this can change the numerical stability of the graph.
+  // Note: this is equivalent to the deprecated option auto_mixed_precision_mkl
+  Toggle auto_mixed_precision_onednn_bfloat16 = 31;
   // Emulate a model using data type float16 on CPU (default is OFF).
   // This will try to emulate the float16 inputs and outputs of an operator
   // on CPU to have better correlation with float16 on GPU; however the
@@ -121,6 +128,8 @@ message RewriterConfig {
   bool disable_meta_optimizer = 19;
   // Optimizers registered by plugin (default is ON)
   Toggle use_plugin_optimizers = 28;
+  // Conditional code motion (default is ON).
+  Toggle experimental_conditional_code_motion = 30;
 
   // Controls how many times we run the optimizers in meta optimizer (default
   // is once).

tensorboard/compat/proto/rpc_options.proto

@@ -0,0 +1,41 @@
+syntax = "proto3";
+
+package tensorboard;
+
+option go_package = "github.com/google/tsl/tsl/go/protobuf/for_core_protos_go_proto";
+
+// RPC options for distributed runtime.
+message RPCOptions {
+  // If true, always use RPC to contact the session target.
+  //
+  // If false (the default option), TensorFlow may use an optimized
+  // transport for client-master communication that avoids the RPC
+  // stack. This option is primarily for used testing the RPC stack.
+  bool use_rpc_for_inprocess_master = 1;
+
+  // The compression algorithm to be used. One of "deflate", "gzip".
+  string compression_algorithm = 2;
+
+  // If compression_algorithm is set, the compression level to be used.
+  // From 0 (no compression), up to 3.
+  int32 compression_level = 3;
+
+  // Setting cache_rpc_response to true will enable sender side caching of
+  // response for RecvTensorAsync and RecvBufAsync to allow receiver to retry
+  // requests . This is only necessary when the network fabric is experiencing a
+  // significant error rate.  Without it we'll fail a step on an network error,
+  // while with it we'll be able to complete long steps (like complex
+  // initializations) in the face of some network errors during RecvTensor.
+  bool cache_rpc_response = 4;
+
+  // Disables TCP connection sharing when opening a new RPC channel.
+  bool disable_session_connection_sharing = 5;
+
+  // Setting num_channels_per_target > 0 allows uses of multiple channels to
+  // communicate to the same target. This can be used to improve the aggregate
+  // throughput on high speed links (e.g 100G) where single connection is not
+  // sufficient to maximize link utilization. Note that a single RPC only goes
+  // on a single channel, this only helps in situations where there are multiple
+  // transfers to the same target overlapping in time.
+  int32 num_channels_per_target = 6;
+}