wehub/tensorflow--tensorflow
1
0
Fork 0
Code Issues Pull Requests Actions 3 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
cffconvert / validate (push) Has been skipped
Details
License Check / license-check (push) Failing after 2s
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 12:14:16 +08:00
commit 8a852e4b4e
36502 changed files with 9277225 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tensorflow/core/protobuf/BUILD
+234
View File
@@ -0,0 +1,234 @@
# copybara:uncomment_begin(oss-unused)
# load("//net/grpc:cc_grpc_library.bzl", "cc_grpc_library")
# load("//net/grpc/go/build_defs:go_grpc_library.bzl", "go_grpc_library")
# copybara:uncomment_end
# For platform specific build config
load(
"//tensorflow:tensorflow.bzl",
"if_google",
)
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_additional_all_protos",
"tf_proto_library",
"tf_pyclif_proto_library",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
default_visibility = [
"//tensorflow:internal",
"//tensorflow/core:__subpackages__",
"//tensorflow_models:__subpackages__",
],
features = if_google(["-parse_headers"]),
licenses = ["notice"],
)
COMMON_PROTO_SRCS = [
"config.proto",
"cluster.proto",
"debug.proto",
"device_filters.proto",
"device_properties.proto",
"queue_runner.proto",
"rewriter_config.proto",
"tensor_bundle.proto",
"saver.proto",
"status.proto",
"verifier_config.proto",
]
[
[
tf_pyclif_proto_library(
name = "%s_pyclif" % proto_name,
proto_lib = ":for_core_protos",
proto_srcfile = "%s.proto" % proto_name,
visibility = ["//visibility:public"],
),
]
for proto_name in [
"config",
"device_properties",
"meta_graph",
"saved_model",
"tensorflow_server",
]
]
tf_proto_library(
name = "autotuning_proto",
srcs = ["autotuning.proto"],
make_default_target_header_only = True,
protodeps = ["@xla//xla:autotuning_proto"],
exports = ["@xla//xla:autotuning_proto"],
)
tf_proto_library(
name = "conv_autotuning_proto",
srcs = ["conv_autotuning.proto"],
make_default_target_header_only = True,
protodeps = [
"@xla//xla/tsl/protobuf:dnn_proto",
],
)
tf_proto_library(
name = "worker_proto",
srcs = ["worker.proto"],
protodeps = tf_additional_all_protos(),
visibility = if_google(
["//tensorflow:internal"],
["//visibility:public"],
),
)
tf_proto_library(
name = "worker_service_proto",
srcs = ["worker_service.proto"],
has_services = 1,
create_service = True,
protodeps = [":worker_proto"],
)
tf_proto_library(
name = "master_proto",
srcs = ["master.proto"],
protodeps = tf_additional_all_protos(),
visibility = ["//tensorflow:internal"],
)
tf_proto_library(
name = "master_service_proto",
srcs = ["master_service.proto"],
has_services = 1,
create_service = True,
protodeps = [":master_proto"],
)
tf_proto_library(
name = "eager_service_proto",
srcs = ["eager_service.proto"],
has_services = 1,
create_grpc_library = True,
create_service = True,
protodeps = tf_additional_all_protos(),
)
tf_proto_library(
name = "replay_log_proto",
srcs = ["replay_log.proto"],
protodeps = [
":master_proto",
] + tf_additional_all_protos(),
)
tf_proto_library(
name = "error_codes_proto_impl",
srcs = ["error_codes.proto"],
make_default_target_header_only = True,
protodeps = ["@xla//xla/tsl/protobuf:error_codes_proto_impl"],
visibility = ["//visibility:public"],
exports = ["@xla//xla/tsl/protobuf:error_codes_proto_impl"],
)
exports_files(
srcs = ["error_codes.proto"] + COMMON_PROTO_SRCS + [
# Protos which are not needed on mobile builds, but should be included
# in protos_all.
#
# Note that some protos are in neither core_proto_srcs nor this
# filegroup; e.g. ones with individual proto_library targets.
"control_flow.proto",
# TODO(ebrevdo): Re-enable once CriticalSection is in core.
# "critical_section.proto",
"snapshot.proto",
"data_service.proto",
"service_config.proto",
"debug_event.proto",
"composite_tensor_variant.proto",
"meta_graph.proto",
"named_tensor.proto",
"remote_tensor_handle.proto",
"saved_model.proto",
"saved_object_graph.proto",
"struct.proto",
"tensorflow_server.proto",
"trackable_object_graph.proto",
"transport_options.proto",
"core_platform_payloads.proto",
"fingerprint.proto",
],
)
tf_proto_library(
name = "for_core_protos",
srcs = COMMON_PROTO_SRCS + [
# Protos which are not needed on mobile builds, but should be included
# in protos_all.
#
# Note that some protos are in neither core_proto_srcs nor this
# filegroup; e.g. ones with individual proto_library targets.
"bfc_memory_map.proto",
"control_flow.proto",
# TODO(ebrevdo): Re-enable once CriticalSection is in core.
# "critical_section.proto",
# TODO: Move snapshot.proto and service_config.proto to a separate package.
# NOTE: Creating an alias and adding the files does not work in OSS.
# NOTE: tf_proto_library requires files to be in the same package.
"snapshot.proto",
"data_service.proto",
"service_config.proto",
"debug_event.proto",
"composite_tensor_variant.proto",
"meta_graph.proto",
"named_tensor.proto",
"remote_tensor_handle.proto",
"rpc_options.proto",
"saved_model.proto",
"saved_object_graph.proto",
"struct.proto",
"tensorflow_server.proto",
"trackable_object_graph.proto",
"transport_options.proto",
"core_platform_payloads.proto",
"fingerprint.proto",
],
make_default_target_header_only = True,
protodeps = [
":error_codes_proto_impl",
"//tensorflow/core/framework:protos_all",
"@xla//xla/tsl/protobuf:bfc_memory_map_proto",
"@xla//xla/tsl/protobuf:coordination_config_proto",
"@xla//xla/tsl/protobuf:rpc_options_proto",
"@xla//xla/tsl/protobuf:status_proto",
],
tags = ["alt_dep=//third_party/tensorflow/core:protos_all"],
visibility = ["//visibility:public"],
exports = [
"@xla//xla/tsl/protobuf:bfc_memory_map_proto",
"@xla//xla/tsl/protobuf:rpc_options_proto",
"@xla//xla/tsl/protobuf:status_proto",
],
)
# copybara:uncomment_begin(google-only)
# cc_grpc_library(
# name = "worker_service_cc_grpc_proto",
# srcs = [":worker_service_proto"],
# generate_mocks = True,
# service_namespace = "grpc_gen",
# deps = [":worker_service_proto_cc"],
# )
#
# cc_grpc_library(
# name = "master_service_cc_grpc_proto",
# srcs = [":master_service_proto"],
# compatible_with = ["//buildenv/target:non_prod"],
# generate_mocks = True,
# service_namespace = "grpc_gen",
# deps = [":master_service_proto_cc"],
# )
# copybara:uncomment_end
tensorflow/core/protobuf/autotuning.proto
+22
View File
@@ -0,0 +1,22 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.dummy;
import public "xla/autotuning.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
tensorflow/core/protobuf/bfc_memory_map.proto
+22
View File
@@ -0,0 +1,22 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.dummy;
import public "xla/tsl/protobuf/bfc_memory_map.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
tensorflow/core/protobuf/cluster.proto
+87
View File
@@ -0,0 +1,87 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
option java_outer_classname = "ClusterProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.distruntime";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// This file contains protos to be used when defining a TensorFlow
// cluster.
//
// EXAMPLES
// --------
//
// 1. A single-process cluster, containing "/job:local/task:0".
//
// Cluster:
// job { name: 'local' tasks { key: 0 value: 'localhost:2222' } }
//
// Server:
// cluster { $CLUSTER } job_name: 'local' task_index: 0
//
// 2. A two-process cluster, containing "/job:local/task:{0,1}".
//
// Cluster:
// job { name: 'local' tasks { key: 0 value: 'localhost:2222' }
// tasks { key: 1 value: 'localhost:2223' } }
//
// Servers:
// cluster { $CLUSTER } job_name: 'local' task_index: 0
// cluster { $CLUSTER } job_name: 'local' task_index: 1
//
// 3. A two-job cluster, containing "/job:worker/task:{0,1,2}" and
// "/job:ps/task:{0,1}".
//
// Cluster:
// job { name: 'worker' tasks { key: 0 value: 'worker1:2222' }
// tasks { key: 1 value: 'worker2:2222' }
// tasks { key: 2 value: 'worker3:2222' } }
// job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
// tasks { key: 1 value: 'ps1:2222' } }
//
// Servers:
// cluster { $CLUSTER } job_name: 'worker' task_index: 0
// cluster { $CLUSTER } job_name: 'worker' task_index: 1
// cluster { $CLUSTER } job_name: 'worker' task_index: 2
// cluster { $CLUSTER } job_name: 'ps' task_index: 0
// cluster { $CLUSTER } job_name: 'ps' task_index: 1
// Defines a single job in a TensorFlow cluster.
message JobDef {
// The name of this job.
string name = 1;
// Mapping from task ID to "hostname:port" string.
//
// If the `name` field contains "worker", and the `tasks` map contains a
// mapping from 7 to "example.org:2222", then the device prefix
// "/job:worker/task:7" will be assigned to "example.org:2222".
//
// If a job has multiple replicas, host-ports will be comma-delimited, with
// one entry for each replica.
map<int32, string> tasks = 2;
}
// Defines a TensorFlow cluster as a set of jobs.
message ClusterDef {
// The jobs that comprise the cluster.
repeated JobDef job = 1;
}
tensorflow/core/protobuf/composite_tensor_variant.proto
+31
View File
@@ -0,0 +1,31 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/protobuf/struct.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Metadata for CompositeTensorVariant, used when serializing as Variant.
//
// We define a new message here (rather than directly using TypeSpecProto for
// the metadata string) to retain flexibility to change the metadata encoding
// to support additional features.
message CompositeTensorVariantMetadata {
TypeSpecProto type_spec_proto = 1;
}
tensorflow/core/protobuf/config.proto
+1080
View File
File diff suppressed because it is too large Load Diff
tensorflow/core/protobuf/control_flow.proto
+106
View File
@@ -0,0 +1,106 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
option java_outer_classname = "ControlFlowProtos";
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";
// Control flow context related protocol buffers.
// Protocol buffer representing the values in ControlFlowContext.
message ValuesDef {
// Value names that have been seen in this context.
repeated string values = 1;
// Value names referenced by but external to this context.
map<string, string> external_values = 2;
}
// Container for any kind of control flow context. Any other control flow
// contexts that are added below should also be added here.
message ControlFlowContextDef {
oneof ctxt {
CondContextDef cond_ctxt = 1;
WhileContextDef while_ctxt = 2;
}
}
// Protocol buffer representing a CondContext object.
message CondContextDef {
// Name of the context.
string context_name = 1;
// Name of the pred tensor.
string pred_name = 2;
// Name of the pivot tensor.
string pivot_name = 3;
// Branch prediction. 0 or 1.
int32 branch = 4;
// Values and external values in control flow context.
ValuesDef values_def = 5;
// Contexts contained inside this context (e.g. nested conds).
repeated ControlFlowContextDef nested_contexts = 6;
}
// Protocol buffer representing a WhileContext object.
message WhileContextDef {
// Name of the context.
string context_name = 1;
// The number of iterations allowed to run in parallel.
int32 parallel_iterations = 2;
// Whether backprop is enabled for this while loop.
bool back_prop = 3;
// Whether GPU-CPU memory swap is enabled for this loop.
bool swap_memory = 4;
// Name of the pivot tensor.
string pivot_name = 5;
// Name of the pivot_for_pred tensor.
string pivot_for_pred_name = 6;
// Name of the pivot_for_body tensor.
string pivot_for_body_name = 7;
// List of names for exit tensors.
repeated string loop_exit_names = 8;
// List of names for enter tensors.
repeated string loop_enter_names = 10;
// Values and external values in control flow context.
ValuesDef values_def = 9;
// Optional name of the maximum_iterations tensor.
string maximum_iterations_name = 11;
// Contexts contained inside this context (e.g. nested whiles).
repeated ControlFlowContextDef nested_contexts = 12;
// Next available id: 13.
}
tensorflow/core/protobuf/conv_autotuning.proto
+69
View File
@@ -0,0 +1,69 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
// This is used for convolution logging. Also see
// tensorflow/core/protobuf/autotuing.h
syntax = "proto3";
package tensorflow;
import "xla/tsl/protobuf/dnn.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// A convolution. Currently it's only used for logging. In the future, we may
// want to use it in the API as well.
message ConvolutionProto {
stream_executor.dnn.ConvolutionKind kind = 1;
stream_executor.dnn.TensorDescriptorProto input = 2;
stream_executor.dnn.TensorDescriptorProto filter = 3;
stream_executor.dnn.TensorDescriptorProto output = 4;
stream_executor.dnn.ConvolutionDescriptorProto conv_desc = 5;
// result = conv_scale * conv(...) + side_value_scale * side_value.
// side_value is an arbitrary buffer if activation is not none. Otherwise, it
// has to be the result buffer (using its old values).
double conv_scale = 6;
double side_value_scale = 7;
stream_executor.dnn.ActivationMode activation = 8;
int64 input_address = 9;
int64 filter_address = 10;
int64 output_address = 11;
int64 bias_address = 12;
int64 side_input_address = 13;
}
message MatmulProto {
stream_executor.dnn.DataType ab_dtype = 1;
stream_executor.dnn.DataType c_dtype = 2;
bool trans_a = 3;
bool trans_b = 4;
uint64 m = 5;
uint64 n = 6;
uint64 k = 7;
int64 lda = 8;
int64 ldb = 9;
int64 ldc = 10;
stream_executor.dnn.ActivationMode activation = 11;
int64 a_address = 12;
int64 b_address = 13;
int64 c_address = 14;
int64 bias_address = 15;
}
tensorflow/core/protobuf/core_platform_payloads.proto
+40
View File
@@ -0,0 +1,40 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.core.platform;
option cc_enable_arenas = true;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// If included as a payload, this message contains the error source information
// where the error was raised.
// URI: "type.googleapis.com/tensorflow.core.platform.ErrorSourceProto"
message ErrorSourceProto {
enum ErrorSource {
UNKNOWN = 0;
TPU_COMPILE_OP = 1;
// Old bridge.
TF_XLA_BRIDGE = 2;
// TPUBridge.
MLIR_BRIDGE_PHASE_1 = 3;
// LegalizeToHlo.
MLIR_BRIDGE_PHASE_2 = 4;
// eager::RemoteMgr.
EAGER_REMOTE_MGR = 5;
}
ErrorSource error_source = 1;
}
tensorflow/core/protobuf/critical_section.proto
+39
View File
@@ -0,0 +1,39 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
option java_outer_classname = "CriticalSectionProtos";
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";
// Protocol buffer representing a CriticalSection.
message CriticalSectionDef {
// Name of the critical section handle.
string critical_section_name = 1;
}
// Protocol buffer representing a CriticalSection execution.
message CriticalSectionExecutionDef {
// Name of the critical section handle.
string execute_in_critical_section_name = 1;
// Whether this operation requires exclusive access to its resources,
// (i.e., no other CriticalSections may request the same resources).
bool exclusive_resource_access = 2;
}
tensorflow/core/protobuf/data_service.proto
+110
View File
@@ -0,0 +1,110 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.data;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Next tag: 2
message ProcessingModeDef {
// Specifies how data is sharded among tf.data service workers.
enum ShardingPolicy {
// No sharding will be performed. Each worker produces the entire dataset
// without any sharding. With this mode, the best practice is to shuffle the
// dataset nondeterministically so that workers process the dataset in
// different orders.
OFF = 0;
// The input dataset is dynamically split among workers at runtime. Each
// worker gets the next split when it reads data from the dispatcher. There
// is no fixed sharding with this mode.
DYNAMIC = 1;
// The following are static sharding policies. The semantics are similar to
// `tf.data.experimental.AutoShardPolicy`. These policies require:
// * The tf.data service cluster has a fixed size, and you need to specify
// the workers in DispatcherConfig.
// * Each client only reads from the local tf.data service worker.
//
// Shards by input files (each worker will get a set of files to process).
// When this option is selected, make sure that there is at least as many
// files as workers. If there are fewer input files than workers, a runtime
// error will be raised.
FILE = 2;
// Shards by elements produced by the dataset. Each worker will process the
// whole dataset and discard the portion that is not for itself. Note that
// for this mode to correctly partitions the dataset elements, the dataset
// needs to produce elements in a deterministic order.
DATA = 3;
// Attempts FILE-based sharding, falling back to DATA-based sharding on
// failures.
FILE_OR_DATA = 4;
// Looks for the presence of `shard(SHARD_HINT, ...)` which is treated as a
// placeholder to replace with `shard(num_workers, worker_index)`.
HINT = 5;
}
ShardingPolicy sharding_policy = 1;
}
// tf.data service deployment mode.
enum DeploymentMode {
DEPLOYMENT_MODE_UNSPECIFIED = 0;
// tf.data service workers colocate with TF workers.
DEPLOYMENT_MODE_COLOCATED = 1;
// tf.data service workers run in dedicated tf.data hosts.
DEPLOYMENT_MODE_REMOTE = 2;
// tf.data service workers run in colocated TF hosts and dedicated tf.data
// hosts.
DEPLOYMENT_MODE_HYBRID = 3;
}
// Metadata related to tf.data service datasets.
// Next tag: 4
message DataServiceMetadata {
oneof optional_element_spec {
// Serialized element spec.
bytes element_spec = 1;
}
enum Compression {
COMPRESSION_UNSPECIFIED = 0;
// No compression.
COMPRESSION_OFF = 1;
// AUTO compression, either none or snappy compression as defined in
// tensorflow/core/platform/snappy.h.
COMPRESSION_SNAPPY = 2;
// Forced a snappy compression as in tensorflow/core/platform/snappy.h.
COMPRESSION_FORCED_SNAPPY = 3;
}
Compression compression = 2;
// Cardinality of the dataset.
int64 cardinality = 3;
}
message CrossTrainerCacheOptions {
string trainer_id = 1;
}
// Data service config available to the client through GetDataServiceConfig RPC.
// Next tag: 2
message DataServiceConfig {
DeploymentMode deployment_mode = 1;
}
tensorflow/core/protobuf/debug.proto
+109
View File
@@ -0,0 +1,109 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
option java_outer_classname = "DebugProtos";
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";
// Option for watching a node in TensorFlow Debugger (tfdbg).
message DebugTensorWatch {
// Name of the node to watch.
// Use "*" for wildcard. But note: currently, regex is not supported in
// general.
string node_name = 1;
// Output slot to watch.
// The semantics of output_slot == -1 is that all outputs of the node
// will be watched (i.e., a wildcard).
// Other negative values of output_slot are invalid and will lead to
// errors currently.
int32 output_slot = 2;
// Name(s) of the debugging op(s).
// One or more than one probes on a tensor.
// e.g., {"DebugIdentity", "DebugNanCount"}
repeated string debug_ops = 3;
// URL(s) for debug targets(s).
//
// Supported URL formats are:
// - file:///foo/tfdbg_dump: Writes out Event content to file
// /foo/tfdbg_dump. Assumes all directories can be created if they don't
// already exist.
// - grpc://localhost:11011: Sends an RPC request to an EventListener
// service running at localhost:11011 with the event.
// - memcbk:///event_key: Routes tensors to clients using the
// callback registered with the DebugCallbackRegistry for event_key.
//
// Each debug op listed in debug_ops will publish its output tensor (debug
// signal) to all URLs in debug_urls.
//
// N.B. Session::Run() supports concurrent invocations of the same inputs
// (feed keys), outputs and target nodes. If such concurrent invocations
// are to be debugged, the callers of Session::Run() must use distinct
// debug_urls to make sure that the streamed or dumped events do not overlap
// among the invocations.
// TODO(cais): More visible documentation of this in g3docs.
repeated string debug_urls = 4;
// Do not error out if debug op creation fails (e.g., due to dtype
// incompatibility). Instead, just log the failure.
bool tolerate_debug_op_creation_failures = 5;
}
// Options for initializing DebuggerState in TensorFlow Debugger (tfdbg).
message DebugOptions {
// Debugging options
repeated DebugTensorWatch debug_tensor_watch_opts = 4;
// Caller-specified global step count.
// Note that this is distinct from the session run count and the executor
// step count.
int64 global_step = 10;
// Whether the total disk usage of tfdbg is to be reset to zero
// in this Session.run call. This is used by wrappers and hooks
// such as the local CLI ones to indicate that the dumped tensors
// are cleaned up from the disk after each Session.run.
bool reset_disk_byte_usage = 11;
}
message DebuggedSourceFile {
// The host name on which a source code file is located.
string host = 1;
// Path to the source code file.
string file_path = 2;
// The timestamp at which the source code file is last modified.
int64 last_modified = 3;
// Byte size of the file.
int64 bytes = 4;
// Line-by-line content of the source code file.
repeated string lines = 5;
}
message DebuggedSourceFiles {
// A collection of source code files.
repeated DebuggedSourceFile source_files = 1;
}
tensorflow/core/protobuf/debug_event.proto
+315
View File
@@ -0,0 +1,315 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/framework/graph_debug_info.proto";
import "tensorflow/core/framework/tensor.proto";
option cc_enable_arenas = true;
option java_outer_classname = "DebugEventProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.util";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Available modes for extracting debugging information from a Tensor.
// TODO(cais): Document the detailed column names and semantics in a separate
// markdown file once the implementation settles.
enum TensorDebugMode {
UNSPECIFIED = 0;
// Only records what tensors are computed, eagerly or in graphs.
// No information regarding the value of the tensor is available.
NO_TENSOR = 1;
// A minimalist health summary for float-type tensors.
// Contains information only about the presence/absence of pathological
// values including Infinity and NaN.
// Applicable only to float dtypes.
CURT_HEALTH = 2;
// A concise health summary for float-type tensors.
// Contains more information that CURT_HEALTH.
// Infinity and NaN are treated differently.
// Applicable only to float and integer dtypes.
CONCISE_HEALTH = 3;
// A detailed health summary.
// Contains further detailed information than `CONCISE_HEALTH`.
// Information about device, dtype and shape are included.
// Counts for various types of values (Infinity, NaN, negative, zero,
// positive) are included.
// Applicable to float, integer and boolean dtypes.
FULL_HEALTH = 4;
// Provides full runtime shape information, up to a maximum rank, beyond
// which the dimension sizes are truncated.
SHAPE = 5;
// Full numeric summary.
// Including device, dtype, shape, counts of various types of values
// (Infinity, NaN, negative, zero, positive), and summary statistics
// (minimum, maximum, mean and variance).
// Applicable to float, integer and boolean dtypes.
FULL_NUMERICS = 6;
// Full tensor value.
FULL_TENSOR = 7;
// Reduce the elements of a tensor to a rank-1 tensor of shape [3], in which
// - the 1st element is -inf if any element of the tensor is -inf,
// or zero otherwise.
// - the 2nd element is +inf if any element of the tensor is +inf,
// or zero otherwise.
// - the 3rd element is nan if any element of the tensor is nan, or zero
// otherwise.
REDUCE_INF_NAN_THREE_SLOTS = 8;
}
// An Event related to the debugging of a TensorFlow program.
message DebugEvent {
// Timestamp in seconds (with microsecond precision).
double wall_time = 1;
// Step of training (if available).
int64 step = 2;
oneof what {
// Metadata related to this debugging data.
DebugMetadata debug_metadata = 3;
// The content of a source file.
SourceFile source_file = 4;
// A stack frame (filename, line number and column number, function name and
// code string) with ID.
StackFrameWithId stack_frame_with_id = 6;
// The creation of an op within a graph (e.g., a FuncGraph compiled from
// a Python function).
GraphOpCreation graph_op_creation = 7;
// Information about a debugged graph.
DebuggedGraph debugged_graph = 8;
// Execution of an op or a Graph (e.g., a tf.function).
Execution execution = 9;
// A graph execution trace: Contains information about the intermediate
// tensors computed during the graph execution.
GraphExecutionTrace graph_execution_trace = 10;
// The ID of the graph (i.e., FuncGraph) executed here: applicable only
// to the execution of a FuncGraph.
string graph_id = 11;
// A device on which debugger-instrumented ops and/or tensors reside.
DebuggedDevice debugged_device = 12;
}
}
// Metadata about the debugger and the debugged TensorFlow program.
message DebugMetadata {
// Version of TensorFlow.
string tensorflow_version = 1;
// Version of the DebugEvent file format.
// Has a format of "debug.Event:<number>", e.g., "debug.Event:1".
string file_version = 2;
// A unique ID for the current run of tfdbg.
// A run of tfdbg is defined as a TensorFlow job instrumented by tfdbg.
// Multiple hosts in a distributed TensorFlow job instrumented by tfdbg
// have the same ID.
string tfdbg_run_id = 3;
}
// Content of a source file involved in the execution of the debugged TensorFlow
// program.
message SourceFile {
// Path to the file.
string file_path = 1;
// Name of the host on which the file is located.
string host_name = 2;
// Line-by-line content of the file.
repeated string lines = 3;
}
// A stack frame with ID.
message StackFrameWithId {
// A unique ID for the stack frame: A UUID-like string.
string id = 1;
// Stack frame, i.e., a frame of a stack trace, containing information
// regarding the file name, line number, function name, code content
// of the line, and column number (if available).
GraphDebugInfo.FileLineCol file_line_col = 2;
}
// Code location information: A stack trace with host-name information.
// Instead of encoding the detailed stack trace, this proto refers to IDs of
// stack frames stored as `StackFrameWithId` protos.
message CodeLocation {
// Host name on which the source files are located.
string host_name = 1;
// ID to a stack frame, each of which is pointed to
// by a unique ID. The ordering of the frames is consistent with Python's
// `traceback.extract_tb()`.
repeated string stack_frame_ids = 2;
}
// The creation of an op in a TensorFlow Graph (e.g., FuncGraph in TF2).
message GraphOpCreation {
// Type of the op (e.g., "MatMul").
string op_type = 1;
// Name of the op (e.g., "Dense/MatMul_1").
string op_name = 2;
// Name of the graph that the op is a part of (if available).
string graph_name = 3;
// Unique ID of the graph (generated by debugger).
// This is the ID of the immediately-enclosing graph.
string graph_id = 4;
// Name of the device that the op is assigned to (if available).
string device_name = 5;
// Names of the input tensors to the op.
repeated string input_names = 6;
// Number of output tensors emitted by the op.
int32 num_outputs = 7;
// The unique ID for code location (stack trace) of the op's creation.
CodeLocation code_location = 8;
// Unique IDs for the output tensors of this op.
repeated int32 output_tensor_ids = 9;
}
// A debugger-instrumented graph.
message DebuggedGraph {
// An ID for the graph.
// This can be used up to look up graph names. Generated by the debugger.
string graph_id = 1;
// Name of the graph (if available).
string graph_name = 2;
// Names of the instrumented ops. This can be used to look up op name
// based on the numeric-summary tensors (2nd column).
repeated string instrumented_ops = 3;
// Original (uninstrumented) GraphDef (if available).
bytes original_graph_def = 4;
// An encoded version of a GraphDef.
// This graph may include the debugger-inserted ops.
bytes instrumented_graph_def = 5;
// IDs of the immediate enclosing context (graph), if any.
string outer_context_id = 6;
}
// A device on which ops and/or tensors are instrumented by the debugger.
message DebuggedDevice {
// Name of the device.
string device_name = 1;
// A debugger-generated ID for the device. Guaranteed to be unique within
// the scope of the debugged TensorFlow program, including single-host and
// multi-host settings.
// TODO(cais): Test the uniqueness guarantee in multi-host settings.
int32 device_id = 2;
}
// Data relating to the eager execution of an op or a Graph.
// For a op that generates N output tensors (N >= 0), only one
// Execution proto will be used to describe the execution event.
message Execution {
// Op type (e.g., "MatMul").
// In the case of a Graph, this is the name of the Graph.
string op_type = 1;
// Number of output tensors.
int32 num_outputs = 2;
// The graph that's executed: applicable only to the eager
// execution of a FuncGraph.
string graph_id = 3;
// IDs of the input tensors (if available).
repeated int64 input_tensor_ids = 4;
// IDs of the output tensors (if availbable).
// If specified, must have the same length as tensor_protos.
repeated int64 output_tensor_ids = 5;
// Type of the tensor value encapsulated in this proto.
TensorDebugMode tensor_debug_mode = 6;
// Output Tensor values in the type described by `tensor_value_type`.
// The length of this should match `num_outputs`.
repeated TensorProto tensor_protos = 7;
// Stack trace of the eager execution.
CodeLocation code_location = 8;
// Debugged-generated IDs of the devices on which the output tensors reside.
// To look up details about the device (e.g., name), cross-reference this
// field with the DebuggedDevice messages.
repeated int32 output_tensor_device_ids = 9;
// TODO(cais): When backporting to V1 Session.run() support, add more fields
// such as fetches and feeds.
}
// Data relating to an execution of a Graph (e.g., an eager execution of a
// FuncGraph).
// The values of the intermediate tensors computed in the graph are recorded
// in this proto. A graph execution may correspond to one or more pieces of
// `GraphExecutionTrace`, depending on whether the instrumented tensor values
// are summarized in an aggregated or separate fashion.
message GraphExecutionTrace {
// Unique ID of the context that the executed op(s) belong to (e.g., a
// compiled concrete tf.function).
string tfdbg_context_id = 1;
// Name of the op (applicable only in the case of the `FULL_TENSOR` trace
// level).
string op_name = 2;
// Output slot of the tensor (applicable only in the case of the `FULL_TENSOR`
// trace level).
int32 output_slot = 3;
// Type of the tensor value encapsulated in this proto.
TensorDebugMode tensor_debug_mode = 4;
// Tensor value in the type described by `tensor_value_type`.
// This tensor may summarize the value of a single intermediate op of the
// graph, or those of multiple intermediate tensors.
TensorProto tensor_proto = 5;
// Name of the device that the op belongs to.
string device_name = 6;
}
tensorflow/core/protobuf/device_filters.proto
+73
View File
@@ -0,0 +1,73 @@
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
option java_outer_classname = "DeviceFiltersProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.distruntime";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// This file contains protos to be used when defining a TensorFlow
// cluster.
//
// Configure device filters for remote tasks in the cluster. When associated
// with a ClusterDef in setting up the cluster, a remote task will ignore all
// devices which do not match any of its filters. Device filters must be
// configured at the cluster startup, and cannot be updated once the cluster is
// up and running.
//
// EXAMPLES
// --------
//
// A two-job cluster with the following ClusterDef:
//
// Cluster:
// job { name: 'worker' tasks { key: 0 value: 'worker1:2222' }
// tasks { key: 1 value: 'worker2:2222' } }
// job { name: 'ps' tasks { key: 0 value: 'ps0:2222' }
// tasks { key: 1 value: 'ps1:2222' } }
//
// Set device filters to isolate worker tasks:
//
// ClusterDeviceFilters:
// job { name: 'worker' tasks { key: 0
// value: device_filter '/job:ps'
// device_filter '/job:worker/task:0' }
// tasks { key: 1
// value: device_filter '/job:ps'
// device_filter '/job:worker/task:1' } }
// Defines the device filters for a remote task.
message TaskDeviceFilters {
repeated string device_filters = 1;
}
// Defines the device filters for tasks in a job.
message JobDeviceFilters {
// The name of this job.
string name = 1;
// Mapping from task ID to task device filters.
map<int32, TaskDeviceFilters> tasks = 2;
}
// Defines the device filters for jobs in a cluster.
message ClusterDeviceFilters {
repeated JobDeviceFilters jobs = 1;
}
tensorflow/core/protobuf/device_properties.proto
+58
View File
@@ -0,0 +1,58 @@
/* Copyright 2017 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
option java_outer_classname = "DevicePropertiesProtos";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
message DeviceProperties {
// Device type (CPU, GPU, ...)
string type = 1;
// Vendor (Intel, nvidia, ...)
string vendor = 2;
// Model (Haswell, K40, ...)
string model = 3;
// Core Frequency in Mhz
int64 frequency = 4;
// Number of cores
int64 num_cores = 5;
// Version of the tools and libraries used with this device (e.g. gcc 4.9,
// cudnn 5.1)
map<string, string> environment = 6;
// Number of registers per core.
int64 num_registers = 7;
// L1 cache size in bytes
int64 l1_cache_size = 8;
// L2 cache size in bytes
int64 l2_cache_size = 9;
// L3 cache size in bytes
int64 l3_cache_size = 10;
// Shared memory size per multiprocessor in bytes. This field is
// applicable to GPUs only.
int64 shared_memory_size_per_multiprocessor = 11;
// Memory size in bytes
int64 memory_size = 12;
// Memory bandwidth in KB/s
int64 bandwidth = 13;
}
message NamedDevice {
string name = 1;
DeviceProperties properties = 2;
}
tensorflow/core/protobuf/eager_service.proto
+385
View File
@@ -0,0 +1,385 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.eager;
import "tensorflow/core/framework/attr_value.proto";
import "tensorflow/core/framework/device_attributes.proto";
import "tensorflow/core/framework/function.proto";
import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/versions.proto";
import "tensorflow/core/protobuf/remote_tensor_handle.proto";
import "tensorflow/core/protobuf/tensorflow_server.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// A proto representation of an eager operation.
message Operation {
// A unique identifier for the operation. Set by the client so that the client
// can uniquely identify the outputs of the scheduled operation.
//
// In the initial implementation, sending duplicate IDs has undefined
// behaviour, but additional constraints may be placed upon this in the
// future.
int64 id = 1;
string name = 2;
message Input {
oneof item {
RemoteTensorHandle remote_handle = 1;
TensorProto tensor = 2;
}
}
repeated Input op_inputs = 10;
// Control Operation IDs that will be respected when ops are re-ordered by
// async execution. If async execution (+ op re-ordering) is not enabled, this
// should have no effect.
repeated int64 control_op_ids = 4;
map<string, AttrValue> attrs = 5;
string device = 6;
// Indicates whether the op is a component of a multi-device function.
bool is_component_function = 7;
// Set when is_component_function is true. It's initially generated
// when we create an FunctionLibraryRuntime::Options (negative value) and used
// to create Rendezvous for function execution. All components of a
// multi-device function should use the same step id to make sure that they
// can communicate through Send/Recv ops.
int64 func_step_id = 8;
// Indicates whether the op is a function.
bool is_function = 9;
reserved 3;
}
message QueueItem {
// The remote executor should be able to handle either executing ops directly,
// or releasing any unused tensor handles, since the tensor lifetime is
// maintained by the client.
oneof item {
RemoteTensorHandle handle_to_decref = 1;
Operation operation = 2;
SendTensorOp send_tensor = 3;
// Takes a FunctionDef and makes it enqueable on the remote worker.
RegisterFunctionOp register_function = 4;
CleanupFunctionOp cleanup_function = 5;
// A remote executor is created to execute ops/functions asynchronously
// enqueued in streaming call. Request with this item type waits for pending
// nodes to finish on the remote executor and report status.
SyncRemoteExecutorForStream sync_remote_executor_for_stream = 6;
SendPackedHandleOp send_packed_handle = 7;
// Takes a FunctionDef signature name and remove it from the remote worker.
RemoveFunctionOp remove_function = 8;
}
}
message QueueResponse {
// `shape` and `tensor` cannot be set in the same response.
// Shapes of output tensors for creating remote TensorHandles.
repeated TensorShapeProto shape = 1;
// Optional. If set, represents the output devices of a function.
repeated string device = 3;
// Output tensors of a remote function. Set when Operation.id is invalid.
repeated TensorProto tensor = 2;
}
message CreateContextRequest {
// Identifies the full cluster, and this particular worker's position within.
ServerDef server_def = 1;
// Whether the ops on the worker should be executed synchronously or
// asynchronously. By default, ops are executed synchronously.
bool async = 2;
// Number of seconds to keep the context alive. If more than keep_alive_secs
// has passed since a particular context has been communicated with, it will
// be garbage collected.
int64 keep_alive_secs = 3;
// This is the version for all the ops that will be enqueued by the client.
VersionDef version_def = 4;
// Device attributes in the cluster
repeated DeviceAttributes cluster_device_attributes = 6;
// The ID of the created context. This is usually a randomly generated number,
// that will be used to identify the context in future requests to the
// service. Contexts are not persisted through server restarts.
// This ID will be used for all future communications as well. It is essential
// that both ends use this ID for selecting a rendezvous to get everything to
// match.
fixed64 context_id = 7;
// The view ID of the context.
fixed64 context_view_id = 8;
// For a multi device function, if false, eagerly copy all remote inputs to
// the default function device; if true, lazily copy remote inputs to their
// target devices after function instantiation to avoid redundant copies.
bool lazy_copy_remote_function_inputs = 9;
// If true, clears resource managers created in the worker environment.
bool clear_existing_contexts = 10;
reserved 5;
}
message CreateContextResponse {
// List of devices that are locally accessible to the worker.
repeated DeviceAttributes device_attributes = 2;
reserved 1;
}
message UpdateContextRequest {
// Identifies the full cluster, and this particular worker's position within.
ServerDef server_def = 1;
// Device attributes in the cluster.
// If this field is empty, it indicates that this is a simple update request
// that only increments the cluster view ID and does not require changes to
// the workers it connects to.
repeated DeviceAttributes cluster_device_attributes = 2;
// The ID of the context to be updated. A context with the specified ID must
// already exist on the recepient server of this request.
fixed64 context_id = 3;
// The view ID of the context, which should be contiguously incremented when
// updating the same context.
fixed64 context_view_id = 4;
}
message UpdateContextResponse {
// List of devices that are locally accessible to the worker.
repeated DeviceAttributes device_attributes = 1;
}
message EnqueueRequest {
fixed64 context_id = 1;
repeated QueueItem queue = 3;
}
message EnqueueResponse {
// A single operation response for every item in the request.
repeated QueueResponse queue_response = 1;
}
message WaitQueueDoneRequest {
fixed64 context_id = 1;
// Ids to wait on. If empty, wait on everything currently pending.
repeated int64 op_id = 2;
}
message WaitQueueDoneResponse {
// TODO(nareshmodi): Consider adding NodeExecStats here to be able to
// propagate some stats.
}
message RunComponentFunctionRequest {
fixed64 context_id = 1;
Operation operation = 2;
// The output indices of its parent function.
repeated int32 output_num = 3;
}
message RunComponentFunctionResponse {
repeated TensorShapeProto shape = 1;
repeated TensorProto tensor = 2;
}
message KeepAliveRequest {
fixed64 context_id = 1;
}
message KeepAliveResponse {
// If the requested context_id is on the remote host, set the context view ID.
fixed64 context_view_id = 1;
}
message CloseContextRequest {
fixed64 context_id = 1;
fixed64 context_view_id = 2;
}
message CloseContextResponse {}
message RegisterFunctionOp {
FunctionDef function_def = 1;
// If true, it means that function_def is produced by graph partition during
// multi-device function instantiation.
bool is_component_function = 2;
// All necessary FunctionDefs and GradientDefs to expand `function_def`.
// When is_component_function is true, `function_def` could be a nested
// function, since some nodes in its parent's function body could be
// replaced with a new function by the graph optimization passes. No need to
// add FunctionDefs here to the function cache in EagerContext since they
// won't be executed as KernelAndDevices.
FunctionDefLibrary library = 3;
}
message RemoveFunctionOp {
// The signature name of function_def to be removed from remote workers.
string function_name = 1;
}
// Cleanup the step state of a multi-device function (e.g. tensors buffered by
// a `Send` op but not picked up by its corresponding `Recv` op).
message CleanupFunctionOp {
int64 step_id = 1;
}
message SyncRemoteExecutorForStream {}
message SendTensorOp {
// All remote tensors are identified by <Op ID, Output num>. To mimic this
// situation when directly sending tensors, we include an "artificial" op ID
// (which would have corresponded to the _Recv op when not using SendTensor).
int64 op_id = 1;
// The index within the repeated field is the output number that will help
// uniquely identify (along with the above op_id) the particular tensor.
repeated TensorProto tensors = 2;
// The device on which the tensors should be resident.
string device_name = 3;
}
// Send a packed TensorHandle to a remote worker.
message SendPackedHandleOp {
// Op id of the remote packed TensorHandle.
int64 op_id = 1;
message LocalTensorHandle {
TensorProto tensor = 1;
// Device where the tensor is produced.
string device = 2;
}
message Handle {
oneof item {
LocalTensorHandle local_handle = 1;
RemoteTensorHandle remote_handle = 2;
}
}
repeated Handle handles = 2;
string device_name = 3;
}
////////////////////////////////////////////////////////////////////////////////
//
// Eager Service defines a TensorFlow service that executes operations eagerly
// on a set of local devices, on behalf of a remote Eager executor.
//
// The service impl will keep track of the various clients and devices it has
// access to and allows the client to enqueue ops on any devices that it is able
// to access and schedule data transfers from/to any of the peers.
//
// A client can generate multiple contexts to be able to independently execute
// operations, but cannot share data between the two contexts.
//
// NOTE: Even though contexts generated by clients should be independent, the
// lower level tensorflow execution engine is not, so they might share some data
// (e.g. a Device's ResourceMgr).
//
////////////////////////////////////////////////////////////////////////////////
service EagerService {
// This initializes the worker, informing it about the other workers in the
// cluster and exchanging authentication tokens which will be used in all
// other RPCs to detect whether the worker has restarted.
rpc CreateContext(CreateContextRequest) returns (CreateContextResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// This updates the eager context on an existing worker when updating the set
// of servers in a distributed eager cluster.
rpc UpdateContext(UpdateContextRequest) returns (UpdateContextResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// This takes a list of Execute and DeleteTensorHandle operations and enqueues
// (in async mode) or executes (in sync mode) them on the remote server.
// All outputs of ops which were not explicitly deleted with
// DeleteTensorHandle entries will be assumed to be alive and are usable by
// future calls to Enqueue.
rpc Enqueue(EnqueueRequest) returns (EnqueueResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// A streaming version of Enqueue.
// Current server implementation sends one response per received request.
// The benefit for using a streaming version is that subsequent requests
// can be sent without waiting for a response to the previous request. This
// synchronization is required in the regular Enqueue call because gRPC does
// not guarantee to preserve request order.
rpc StreamingEnqueue(stream EnqueueRequest) returns (stream EnqueueResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// Takes a set of op IDs and waits until those ops are done. Returns any error
// in the stream so far.
rpc WaitQueueDone(WaitQueueDoneRequest) returns (WaitQueueDoneResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// This takes an Eager operation and executes it in async mode on the remote
// server. Different from EnqueueRequest, ops/functions sent through this
// type of requests are allowed to execute in parallel and no ordering is
// preserved by RPC stream or executor.
// This request type should only be used for executing component functions.
// Ordering of component functions should be enforced by their corresponding
// main functions. The runtime ensures the following invarients for component
// functions (CFs) and their main functions (MFs):
// (1) MF1 -> MF2 ==> CF1 -> CF2 ("->" indicates order of execution);
// (2) MF1 || MF2 ==> CF1 || CF2 ("||" indicates possible parallel execution);
// (3) For CF1 and CF2 that come from the same MF, CF1 || CF2
// For executing ops/main functions, use Enqueue or StreamingEnqueue instead
// for correct ordering.
rpc RunComponentFunction(RunComponentFunctionRequest)
returns (RunComponentFunctionResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// Contexts are always created with a deadline and no RPCs within a deadline
// will trigger a context garbage collection. KeepAlive calls can be used to
// delay this. It can also be used to validate the existence of a context ID
// on remote eager worker. If the context is on remote worker, return the same
// ID and the current context view ID. This is useful for checking if the
// remote worker (potentially with the same task name and hostname / port) is
// replaced with a new process.
rpc KeepAlive(KeepAliveRequest) returns (KeepAliveResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// Closes the context. No calls to other methods using the existing context ID
// are valid after this.
rpc CloseContext(CloseContextRequest) returns (CloseContextResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
}
tensorflow/core/protobuf/error_codes.proto
+26
View File
@@ -0,0 +1,26 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
// Add a dummy package name. Having no package, like
// core/lib/core/error_codes.proto, or having tensorflow.error, like
// tsl/protobuf/error_codes.proto, results in name collision errors in generated
// code for some users that use JS through J2CL.
package tensorflow.error.dummy;
import public "xla/tsl/protobuf/error_codes.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
tensorflow/core/protobuf/fingerprint.proto
+48
View File
@@ -0,0 +1,48 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/framework/versions.proto";
option cc_enable_arenas = true;
option java_outer_classname = "FingerprintProtos";
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";
// Protocol buffer representing a SavedModel Fingerprint.
//
// If there are multiple MetaGraphDefs in the SavedModel, the FingerprintDef
// corresponds to the first one.
message FingerprintDef {
// Hash of the saved_model.pb, referred to as a "checksum".
uint64 saved_model_checksum = 1;
// Hash of regularized graph_def.
uint64 graph_def_program_hash = 2;
// Hash of the regularized (sorted) SignatureDefs.
uint64 signature_def_hash = 3;
// Hash of the regularized SavedObjectGraph.
uint64 saved_object_graph_hash = 4;
// Hash of the checkpoint.
uint64 checkpoint_hash = 5;
// An UUID for the model, chosen at random, not related to the hashes.
string uuid = 7;
// Version specification of the fingerprint.
VersionDef version = 6;
// TODO(b/290068219): add USM version when GA
}
tensorflow/core/protobuf/master.proto
+353
View File
@@ -0,0 +1,353 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
syntax = "proto3";
package tensorflow;
import "tensorflow/core/framework/device_attributes.proto";
import "tensorflow/core/framework/graph.proto";
import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/protobuf/config.proto";
import "tensorflow/core/protobuf/error_codes.proto";
import "tensorflow/core/protobuf/named_tensor.proto";
option cc_enable_arenas = true;
option java_outer_classname = "DistributedRuntimeProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.distruntime";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
////////////////////////////////////////////////////////////////////////////////
//
// CreateSession method request/response protos.
//
////////////////////////////////////////////////////////////////////////////////
message CreateSessionRequest {
// The initial graph definition.
GraphDef graph_def = 1;
// Configuration options.
ConfigProto config = 2;
// The target string used from the client's perspective.
string target = 3;
}
message CreateSessionResponse {
// The session handle to be used in subsequent calls for the created session.
//
// The client must arrange to call CloseSession with this returned
// session handle to close the session.
string session_handle = 1;
// The initial version number for the graph, to be used in the next call
// to ExtendSession.
int64 graph_version = 2;
}
////////////////////////////////////////////////////////////////////////////////
//
// ExtendSession method request/response protos.
//
// The "graph_def" specifies a set of nodes to be added to the session's graph.
//
// A typical "graph_def" will contain:
//
// * Zero or more new nodes with names that do not exist in the server-side
// graph. These will be added to the graph.
//
// PRECONDITION: The server-side current version is req.current_version.
// None of the names in req.graph_def appeared in previous successful calls to
// CreateSession or ExtendSession with the same session_handle.
// POSTCONDITION: The server-side current version is resp.new_version.
//
////////////////////////////////////////////////////////////////////////////////
message ExtendSessionRequest {
// REQUIRED: session_handle must be returned by a CreateSession call
// to the same master service.
string session_handle = 1;
// REQUIRED: The nodes to be added to the session's graph. If any node has
// the same name as an existing node, the operation will fail with
// ILLEGAL_ARGUMENT.
GraphDef graph_def = 2;
// REQUIRED: The version number of the graph to be extended. This will be
// tested against the current server-side version number, and the operation
// will fail with FAILED_PRECONDITION if they do not match.
int64 current_graph_version = 3;
}
message ExtendSessionResponse {
// TODO(mrry): Return something about the operation?
// The new version number for the extended graph, to be used in the next call
// to ExtendSession.
int64 new_graph_version = 4;
}
////////////////////////////////////////////////////////////////////////////////
//
// RunStep method request/response protos.
//
// The caller should provide the feeds needed by the graph and specify
// what nodes should be fetched.
//
////////////////////////////////////////////////////////////////////////////////
message RunStepRequest {
// REQUIRED: session_handle must be returned by a CreateSession call
// to the same master service.
string session_handle = 1;
// Tensors to be fed in the step. Each feed is a named tensor.
repeated NamedTensorProto feed = 2;
// Fetches. A list of tensor names. The caller expects a tensor to
// be returned for each fetch[i] (see RunStepResponse.tensor). The
// order of specified fetches does not change the execution order.
repeated string fetch = 3;
// Target Nodes. A list of node names. The named nodes will be run
// to but their outputs will not be fetched.
repeated string target = 4;
// Options for the run call.
RunOptions options = 5;
// Partial run handle (optional). If specified, this will be a partial run
// execution, run up to the specified fetches.
string partial_run_handle = 6;
// If true then some errors, e.g., execution errors that have long
// error messages, may return an OK RunStepResponse with the actual
// error saved in the status_code/status_error_message fields of the
// response body. This is a workaround since the RPC subsystem may
// truncate long metadata messages.
bool store_errors_in_response_body = 7;
// Unique identifier for this request. Every RunStepRequest must
// have a unique request_id, and retried RunStepRequest must have
// the same request_id. If request_id is zero, retry detection is disabled.
int64 request_id = 8;
}
message RunStepResponse {
// NOTE: The order of the returned tensors may or may not match
// the fetch order specified in RunStepRequest.
repeated NamedTensorProto tensor = 1;
// Returned metadata if requested in the options.
RunMetadata metadata = 2;
// If store_errors_in_response_body is true in the request, then
// optionally the server may return an OK status for the RPC and
// fill the true status into the fields below, to allow for messages
// that are too long to fit in metadata.
error.Code status_code = 3;
string status_error_message = 4;
}
////////////////////////////////////////////////////////////////////////////////
//
// PartialRunSetup method request/response protos.
//
// The caller should provide the future partial run feeds, fetches, and targets.
// Then the caller can use RunStepRequest with is_partial set to make partial
// run calls.
//
////////////////////////////////////////////////////////////////////////////////
message PartialRunSetupRequest {
// REQUIRED: session_handle must be returned by a CreateSession call
// to the same master service.
string session_handle = 1;
// Tensors to be fed in future steps.
repeated string feed = 2;
// Fetches. A list of tensor names. The caller expects a tensor to be returned
// for each fetch[i] (see RunStepResponse.tensor), for corresponding partial
// RunStepRequests. The order of specified fetches does not change the
// execution order.
repeated string fetch = 3;
// Target Nodes. A list of node names. The named nodes will be run in future
// steps, but their outputs will not be fetched.
repeated string target = 4;
// Unique identifier for this request. Every PartialRunSetupRequest must
// have a unique request_id, and retried PartialRunSetupRequest must have
// the same request_id. If request_id is zero, retry detection is disabled.
int64 request_id = 5;
}
message PartialRunSetupResponse {
// The unique handle corresponding to the ongoing partial run call setup by
// the invocation to PartialRunSetup. This handle may be passed to
// RunStepRequest to send and receive tensors for this partial run.
string partial_run_handle = 1;
}
////////////////////////////////////////////////////////////////////////////////
//
// CloseSession method request/response protos.
//
////////////////////////////////////////////////////////////////////////////////
message CloseSessionRequest {
// REQUIRED: session_handle must be returned by a CreateSession call
// to the same master service.
string session_handle = 1;
}
message CloseSessionResponse {}
// Reset() allows misbehaving or slow sessions to be aborted and closed, and
// causes their resources eventually to be released. Reset() does not wait
// for the computations in old sessions to cease; it merely starts the
// process of tearing them down. However, if a new session is started after
// a Reset(), the new session is isolated from changes that old sessions
// (started prior to the Reset()) may continue to make to resources, provided
// all those resources are in containers listed in "containers".
//
// Old sessions may continue to have side-effects on resources not in
// containers listed in "containers", and thus may affect future
// sessions' results in ways that are hard to predict. Thus, if well-defined
// behavior is desired, is it recommended that all containers be listed in
// "containers". Similarly, if a device_filter is specified, results may be
// hard to predict.
message ResetRequest {
// A list of container names, which may be empty.
//
// If 'container' is not empty, releases resources in the given
// containers in all devices.
//
// If 'container' is empty, releases resources in the default
// container in all devices.
repeated string container = 1;
// When any filters are present, only devices that match the filters
// will be reset. Each filter can be partially specified,
// e.g. "/job:ps" "/job:worker/replica:3", etc.
repeated string device_filters = 2;
}
message ResetResponse {}
////////////////////////////////////////////////////////////////////////////////
//
// ListDevices method request/response protos.
//
// Returns information about the TensorFlow devices that are available
// to this master.
//
////////////////////////////////////////////////////////////////////////////////
message ListDevicesRequest {
// Optional: session_handle must be returned by a CreateSession call to the
// same master service.
//
// When session_handle is empty, the ClusterSpec provided when the master was
// started is used to compute the available devices. If the session_handle is
// provided but not recognized, an error is returned. Finally, if a valid
// session_handle is provided, the cluster configuration for that session is
// used when computing the response.
string session_handle = 1;
}
message ListDevicesResponse {
repeated DeviceAttributes local_device = 1;
repeated DeviceAttributes remote_device = 2;
}
////////////////////////////////////////////////////////////////////////////////
//
// MakeCallable method request/response protos.
//
////////////////////////////////////////////////////////////////////////////////
message MakeCallableRequest {
// REQUIRED: session_handle must be returned by a CreateSession call
// to the same master service.
string session_handle = 1;
// Options that define the behavior of the created callable.
CallableOptions options = 2;
// Unique identifier for this request. Every MakeCallableRequest must
// have a unique request_id, and retried MakeCallableRequest must have
// the same request_id. If request_id is zero, retry detection is disabled.
int64 request_id = 3;
}
message MakeCallableResponse {
// A handle to the created callable.
int64 handle = 1;
}
////////////////////////////////////////////////////////////////////////////////
//
// RunCallable method request/response protos.
//
////////////////////////////////////////////////////////////////////////////////
message RunCallableRequest {
// REQUIRED: session_handle must be returned by a CreateSession call
// to the same master service.
string session_handle = 1;
// REQUIRED: handle must be returned by a MakeCallable call to the same
// master service.
int64 handle = 2;
// Values of the tensors passed as arguments to the callable, in the order
// defined in the CallableOptions.feed field passed to MakeCallable.
repeated TensorProto feed = 3;
// Unique identifier for this request. Every RunCallableRequest must
// have a unique request_id, and retried RunCallableRequest must have
// the same request_id. If request_id is zero, retry detection is disabled.
int64 request_id = 4;
}
message RunCallableResponse {
// Values of the tensors returned by the callable, in the order defined in the
// CallableOptions.fetch field passed to MakeCallable.
repeated TensorProto fetch = 1;
// Returned metadata if requested in the options.
RunMetadata metadata = 2;
}
////////////////////////////////////////////////////////////////////////////////
//
// ReleaseCallable method request/response protos.
//
////////////////////////////////////////////////////////////////////////////////
message ReleaseCallableRequest {
// REQUIRED: session_handle must be returned by a CreateSession call
// to the same master service.
string session_handle = 1;
// REQUIRED: handle must be returned by a MakeCallable call to the same
// master service.
int64 handle = 2;
}
message ReleaseCallableResponse {}
tensorflow/core/protobuf/master_service.proto
+121
View File
@@ -0,0 +1,121 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
syntax = "proto3";
package tensorflow.grpc;
import "tensorflow/core/protobuf/master.proto";
option java_outer_classname = "MasterServiceProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.distruntime";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
////////////////////////////////////////////////////////////////////////////////
//
// MasterService defines a TensorFlow service with which a client can
// interact to execute a distributed TensorFlow computation.
//
// A master service keeps track of multiple "master sessions". Each
// session encapsulates a computation graph and its associated state,
// and typically corresponds to a single "client session" (e.g. a
// `tensorflow::Session` instance).
//
// A session is responsible for the following:
// * assigning each node to a device (locally or remotely) using a
// placement algorithm. This may make decisions based on collected
// statistics from the workers in the system (e.g., memory usage,
// bandwidth consumption, etc.)
//
// * inserting intermediate nodes and edges to support cross-device
// and cross-process data flows and resource management.
//
// * issuing commands to workers to execute the subgraphs associated
// with those workers.
//
// Typically, a client carries out an iterative computation
// (e.g. training) by invoking RPCs against the master in a
// client-side loop. The client first creates a client session that
// connects to a particular master (using gRPC for example). The
// master creates a corresponding master session that is hosted on
// the master and caches state between the client's invocations.
//
// After the session is established, the master returns an opaque
// handle to the client that can be used to associate the client and
// master sessions.
//
// The client may send an initial graph to the master in the
// CreateSession call, and add nodes to the graph using ExtendSession.
//
// The most frequent operation a master is "RunStep", which implements
// the `Session::Run()` API. It supports feeding in arguments,
// executing a dataflow computation, and fetching arguments.
//
// Finally, when the client no longer needs the session, it should
// close the session by invoking CloseSession, which allows the master
// to reclaim resources associated with the session. The master may
// implement a garbage collection scheme that closes sessions that
// have been inactive for some time.
//
// For example, the following pseudo-code illustrates how a client
// interacts with a master:
//
// stub = NewStub("/job:mnist/replica:0/task:0")
// {handle} = stub->CreateSession({graph_def})
// do {
// stub->RunStep({handle, {feeds}, {fetches}})
// // The client can evaluate a predicate locally, based on the
// // result of `fetches`, to determine whether to terminate. For
// // example, it might fetch the loss and evaluate whether it is less
// // than some threshold.
// } while (!should_stop({fetches}));
// stub->CloseSession({handle})
//
////////////////////////////////////////////////////////////////////////////////
service MasterService {
// Creates a session.
rpc CreateSession(CreateSessionRequest) returns (CreateSessionResponse);
// Extends a session.
rpc ExtendSession(ExtendSessionRequest) returns (ExtendSessionResponse);
// Prepares future partial run calls.
rpc PartialRunSetup(PartialRunSetupRequest) returns (PartialRunSetupResponse);
// Drives the graph computation.
rpc RunStep(RunStepRequest) returns (RunStepResponse);
// Closes a session.
rpc CloseSession(CloseSessionRequest) returns (CloseSessionResponse);
// List the devices usable by the master.
rpc ListDevices(ListDevicesRequest) returns (ListDevicesResponse);
// Close and abandon all existing sessions. Ongoing computations
// will no longer affect fresh ones via the resources in containers listed in
// the ResetRequest. See ResetRequest for more details.
rpc Reset(ResetRequest) returns (ResetResponse);
// Registers a callable for execution with RunCallable.
rpc MakeCallable(MakeCallableRequest) returns (MakeCallableResponse);
// Executes a callable registered with MakeCallable.
rpc RunCallable(RunCallableRequest) returns (RunCallableResponse);
// Frees resources associated with a callable registered with MakeCallable.
rpc ReleaseCallable(ReleaseCallableRequest) returns (ReleaseCallableResponse);
}
tensorflow/core/protobuf/meta_graph.proto
+301
View File
@@ -0,0 +1,301 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "google/protobuf/any.proto";
import "tensorflow/core/framework/graph.proto";
import "tensorflow/core/framework/op_def.proto";
import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";
import "tensorflow/core/protobuf/saved_object_graph.proto";
import "tensorflow/core/protobuf/saver.proto";
import "tensorflow/core/protobuf/struct.proto";
option cc_enable_arenas = true;
option java_outer_classname = "MetaGraphProtos";
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";
// 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
// process boundary. It can be used for long term storage of graphs,
// cross-language execution of graphs, etc.
// MetaInfoDef
// GraphDef
// SaverDef
// CollectionDef
// TensorInfo
// SignatureDef
message MetaGraphDef {
// Meta information regarding the graph to be exported. To be used by users
// of this protocol buffer to encode information regarding their meta graph.
message MetaInfoDef {
// User specified Version string. Can be the name of the model and revision,
// steps this model has been trained to, etc.
string meta_graph_version = 1;
// A copy of the OpDefs used by the producer of this graph_def.
// Descriptions and Ops not used in graph_def are stripped out.
OpList stripped_op_list = 2;
// A serialized protobuf. Can be the time this meta graph is created, or
// modified, or name of the model.
google.protobuf.Any any_info = 3;
// User supplied tag(s) on the meta_graph and included graph_def.
//
// MetaGraphDefs should be tagged with their capabilities or use-cases.
// Examples: "train", "serve", "gpu", "tpu", etc.
// These tags enable loaders to access the MetaGraph(s) appropriate for a
// specific use-case or runtime environment.
repeated string tags = 4;
// The __version__ string of the tensorflow build used to write this graph.
// This will be populated by the framework, which will overwrite any user
// supplied value.
string tensorflow_version = 5;
// The __git_version__ string of the tensorflow build used to write this
// graph. This will be populated by the framework, which will overwrite any
// user supplied value.
string tensorflow_git_version = 6;
// A flag to denote whether default-valued attrs have been stripped from
// the nodes in this graph_def.
bool stripped_default_attrs = 7;
// FunctionDef name to aliases mapping.
map<string, string> function_aliases = 8;
}
MetaInfoDef meta_info_def = 1;
// GraphDef.
GraphDef graph_def = 2;
// SaverDef.
SaverDef saver_def = 3;
// collection_def: Map from collection name to collections.
// See CollectionDef section for details.
map<string, CollectionDef> collection_def = 4;
// signature_def: Map from user supplied key for a signature to a single
// SignatureDef.
map<string, SignatureDef> signature_def = 5;
// Asset file def to be used with the defined graph.
repeated AssetFileDef asset_file_def = 6;
// Extra information about the structure of functions and stateful objects.
SavedObjectGraph object_graph_def = 7;
}
// CollectionDef should cover most collections.
// To add a user-defined collection, do one of the following:
// 1. For simple data types, such as string, int, float:
// tf.add_to_collection("your_collection_name", your_simple_value)
// strings will be stored as bytes_list.
//
// 2. For Protobuf types, there are three ways to add them:
// 1) tf.add_to_collection("your_collection_name",
// your_proto.SerializeToString())
//
// collection_def {
// key: "user_defined_bytes_collection"
// value {
// bytes_list {
// value: "queue_name: \"test_queue\"\n"
// }
// }
// }
//
// or
//
// 2) tf.add_to_collection("your_collection_name", str(your_proto))
//
// collection_def {
// key: "user_defined_string_collection"
// value {
// bytes_list {
// value: "\n\ntest_queue"
// }
// }
// }
//
// or
//
// 3) any_buf = any_pb2.Any()
// tf.add_to_collection("your_collection_name",
// any_buf.Pack(your_proto))
//
// collection_def {
// key: "user_defined_any_collection"
// value {
// any_list {
// value {
// type_url: "type.googleapis.com/tensorflow.QueueRunnerDef"
// value: "\n\ntest_queue"
// }
// }
// }
// }
//
// 3. For Python objects, implement to_proto() and from_proto(), and register
// them in the following manner:
// ops.register_proto_function("your_collection_name",
// proto_type,
// to_proto=YourPythonObject.to_proto,
// from_proto=YourPythonObject.from_proto)
// These functions will be invoked to serialize and de-serialize the
// collection. For example,
// ops.register_proto_function(ops.GraphKeys.GLOBAL_VARIABLES,
// proto_type=variable_pb2.VariableDef,
// to_proto=Variable.to_proto,
// from_proto=Variable.from_proto)
message CollectionDef {
// NodeList is used for collecting nodes in graph. For example
// collection_def {
// key: "summaries"
// value {
// node_list {
// value: "input_producer/ScalarSummary:0"
// value: "shuffle_batch/ScalarSummary:0"
// value: "ImageSummary:0"
// }
// }
message NodeList {
repeated string value = 1;
}
// BytesList is used for collecting strings and serialized protobufs. For
// example:
// collection_def {
// key: "trainable_variables"
// value {
// bytes_list {
// value: "\n\017conv1/weights:0\022\024conv1/weights/Assign
// \032\024conv1/weights/read:0"
// value: "\n\016conv1/biases:0\022\023conv1/biases/Assign\032
// \023conv1/biases/read:0"
// }
// }
// }
message BytesList {
repeated bytes value = 1;
}
// Int64List is used for collecting int, int64 and long values.
message Int64List {
repeated int64 value = 1 [packed = true];
}
// FloatList is used for collecting float values.
message FloatList {
repeated float value = 1 [packed = true];
}
// AnyList is used for collecting Any protos.
message AnyList {
repeated google.protobuf.Any value = 1;
}
oneof kind {
NodeList node_list = 1;
BytesList bytes_list = 2;
Int64List int64_list = 3;
FloatList float_list = 4;
AnyList any_list = 5;
}
}
// Information about a Tensor necessary for feeding or retrieval.
message TensorInfo {
// For sparse tensors, The COO encoding stores a triple of values, indices,
// and shape.
message CooSparse {
// The shape of the values Tensor is [?]. Its dtype must be the dtype of
// the SparseTensor as a whole, given in the enclosing TensorInfo.
string values_tensor_name = 1;
// The indices Tensor must have dtype int64 and shape [?, ?].
string indices_tensor_name = 2;
// The dynamic logical shape represented by the SparseTensor is recorded in
// the Tensor referenced here. It must have dtype int64 and shape [?].
string dense_shape_tensor_name = 3;
}
// Generic encoding for composite tensors.
message CompositeTensor {
// The serialized TypeSpec for the composite tensor.
TypeSpecProto type_spec = 1;
// A TensorInfo for each flattened component tensor.
repeated TensorInfo components = 2;
}
oneof encoding {
// For dense `Tensor`s, the name of the tensor in the graph.
string name = 1;
// There are many possible encodings of sparse matrices
// (https://en.wikipedia.org/wiki/Sparse_matrix). Currently, TensorFlow
// uses only the COO encoding. This is supported and documented in the
// SparseTensor Python class.
CooSparse coo_sparse = 4;
// Generic encoding for CompositeTensors.
CompositeTensor composite_tensor = 5;
}
DataType dtype = 2;
// The static shape should be recorded here, to the extent that it can
// be known in advance. In the case of a SparseTensor, this field describes
// the logical shape of the represented tensor (aka dense_shape).
TensorShapeProto tensor_shape = 3;
}
// SignatureDef defines the signature of a computation supported by a TensorFlow
// graph.
message SignatureDef {
// Named input parameters.
map<string, TensorInfo> inputs = 1;
// Named output parameters.
map<string, TensorInfo> outputs = 2;
// Deprecated: TensorFlow 2 always sets this to a fixed value;
// open-source TF Serving stopped checking by default since release 2.4.
//
// 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;
}
// An asset file def for a single file or a set of sharded files with the same
// name.
message AssetFileDef {
// The tensor to bind the asset filename to.
TensorInfo tensor_info = 1;
// The filename within an assets directory. Note: does not include the path
// prefix, i.e. directories. For an asset at /tmp/path/vocab.txt, the filename
// would be "vocab.txt".
string filename = 2;
}
tensorflow/core/protobuf/named_tensor.proto
+40
View File
@@ -0,0 +1,40 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/framework/tensor.proto";
option cc_enable_arenas = true;
option java_outer_classname = "NamedTensorProtos";
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";
// A pair of tensor name and tensor values.
message NamedTensorProto {
// Name of the tensor.
string name = 1;
// The client can populate a TensorProto using a tensorflow::Tensor`, or
// directly using the protobuf field accessors.
//
// The client specifies whether the returned tensor values should be
// filled tensor fields (float_val, int_val, etc.) or encoded in a
// compact form in tensor.tensor_content.
TensorProto tensor = 2;
}
tensorflow/core/protobuf/queue_runner.proto
+45
View File
@@ -0,0 +1,45 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/protobuf/error_codes.proto";
option cc_enable_arenas = true;
option java_outer_classname = "QueueRunnerProtos";
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";
// Protocol buffer representing a QueueRunner.
message QueueRunnerDef {
// Queue name.
string queue_name = 1;
// A list of enqueue operations.
repeated string enqueue_op_name = 2;
// The operation to run to close the queue.
string close_op_name = 3;
// The operation to run to cancel the queue.
string cancel_op_name = 4;
// A list of exception types considered to signal a safely closed queue
// if raised during enqueue operations.
repeated error.Code queue_closed_exception_types = 5;
}
tensorflow/core/protobuf/remote_tensor_handle.proto
+49
View File
@@ -0,0 +1,49 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.eager;
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";
option cc_enable_arenas = true;
option java_outer_classname = "RemoteTensorHandleProtos";
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";
message ResourceDtypeAndShape {
DataType dtype = 1;
TensorShapeProto shape = 2;
}
message RemoteTensorHandle {
// The ID of the operation that produced this tensor.
int64 op_id = 1;
// The index into the outputs of the operation that produced this tensor.
int32 output_num = 2;
// Device where the tensor is located. Cannot be empty.
// For multi-device functions, it's the default device passed to placer.
string device = 3;
// Device of the operation producing this tensor. Can be empty if the
// operation producing this tensor is a multi-device function.
string op_device = 4;
// Tensor type.
DataType dtype = 5;
// Optional data types and shapes of a remote resource variable.
repeated ResourceDtypeAndShape resource_dtypes_and_shapes = 6;
}
tensorflow/core/protobuf/replay_log.proto
+62
View File
@@ -0,0 +1,62 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/protobuf/master.proto";
option cc_enable_arenas = true;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Records the creation of a new replay session. We record the device listing
// here to capture the state of the cluster.
message NewReplaySession {
ListDevicesResponse devices = 1;
string session_handle = 2;
}
message ReplayOp {
double start_time_us = 31;
double end_time_us = 32;
oneof op {
CreateSessionRequest create_session = 1;
ExtendSessionRequest extend_session = 2;
PartialRunSetupRequest partial_run_setup = 3;
RunStepRequest run_step = 4;
CloseSessionRequest close_session = 5;
ListDevicesRequest list_devices = 6;
ResetRequest reset_request = 7;
MakeCallableRequest make_callable = 8;
RunCallableRequest run_callable = 9;
ReleaseCallableRequest release_callable = 10;
NewReplaySession new_replay_session = 11;
}
oneof response {
CreateSessionResponse create_session_response = 21;
ExtendSessionResponse extend_session_response = 22;
PartialRunSetupResponse partial_run_setup_response = 23;
RunStepResponse run_step_response = 24;
CloseSessionResponse close_session_response = 25;
ListDevicesResponse list_devices_response = 26;
ResetResponse reset_request_response = 27;
MakeCallableResponse make_callable_response = 28;
RunCallableResponse run_callable_response = 29;
ReleaseCallableResponse release_callable_response = 30;
}
}
tensorflow/core/protobuf/rewriter_config.proto
+256
View File
@@ -0,0 +1,256 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/framework/attr_value.proto";
import "tensorflow/core/protobuf/verifier_config.proto";
option cc_enable_arenas = true;
option java_outer_classname = "RewriterConfigProtos";
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";
message AutoParallelOptions {
bool enable = 1;
int32 num_replicas = 2;
}
message ScopedAllocatorOptions {
// If present, only perform optimization for these ops.
repeated string enable_op = 1;
}
message RewriterConfig {
// Graph rewriting is experimental and subject to change, not covered by any
// API stability guarantees.
// Configuration options for the meta-optimizer. Unless otherwise noted, these
// configuration options do not apply to explicitly triggered optimization
// passes in the optimizers field.
enum Toggle {
DEFAULT = 0;
ON = 1;
OFF = 2;
// Enable some aggressive optimizations that use assumptions that TF graphs
// may break. For example, assume the shape of a placeholder matches its
// actual feed.
AGGRESSIVE = 3;
// Run MLIR pass if there's one implemented in TFG, do nothing otherwise.
// I.e., if there's no corresponding TFG pass, it's an OFF. This is supposed
// to be mapped with `ON` and there's no `AGGRESSIVE` in MLIR pass now.
EXPERIMENTAL_MLIR = 4;
// Run both MLIR and Grappler passes consecutively and MLIR pass will come
// first.
EXPERIMENTAL_BOTH = 5;
}
// Enum for layout conversion between NCHW and NHWC on CPU. Default is OFF.
enum CpuLayout {
NO_CONVERSION_ON_CPU = 0;
NCHW_TO_NHWC = 1;
NHWC_TO_NCHW = 2;
}
// Enum controlling the number of times to run optimizers. The default is to
// run them twice.
enum NumIterationsType {
DEFAULT_NUM_ITERS = 0;
ONE = 1;
TWO = 2;
}
// CPU Conversion settings between NHCW and NCHW.
CpuLayout cpu_layout_conversion = 50;
// Optimize tensor layouts (default is ON)
// e.g. This will try to use NCHW layout on GPU which is faster.
Toggle layout_optimizer = 1;
// Fold constants (default is ON)
// Statically infer the value of tensors when possible, and materialize the
// result using constants.
Toggle constant_folding = 3;
// Shape optimizations (default is ON)
// Simplify computations made on shapes.
Toggle shape_optimization = 13;
// Remapping (default is ON)
// Remap subgraphs onto more efficient implementations.
Toggle remapping = 14;
// Common subgraph elimination (default is ON)
// e.g. Simplify arithmetic ops; merge ops with same value (like constants).
Toggle common_subgraph_elimination = 24;
// Arithmetic optimizations (default is ON)
// e.g. Simplify arithmetic ops; merge ops with same value (like constants).
Toggle arithmetic_optimization = 7;
// Control dependency optimizations (default is ON).
// Remove redundant control dependencies, which may enable other optimization.
Toggle dependency_optimization = 8;
// Loop optimizations (default is ON).
Toggle loop_optimization = 9;
// Function optimizations (default is ON).
Toggle function_optimization = 10;
// Strips debug-related nodes from the graph (off by default).
Toggle debug_stripper = 11;
// If true, don't remove unnecessary ops from the graph
bool disable_model_pruning = 2;
// Try to allocate some independent Op outputs contiguously in order to
// merge or eliminate downstream Ops (off by default).
Toggle scoped_allocator_optimization = 15;
// Force small ops onto the CPU (default is OFF).
Toggle pin_to_host_optimization = 18;
// Enable the swap of kernel implementations based on the device placement
// (default is ON).
Toggle implementation_selector = 22;
// 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;
// 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
// computation in the operator is based on float32.
// Note that this can change the numerical stability of the graph.
Toggle auto_mixed_precision_cpu = 29;
// Disable the entire meta optimizer (off by default).
bool disable_meta_optimizer = 19;
// Disable the TFG optimizer (off by default).
bool disable_tfg_optimizer = 32;
// 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).
NumIterationsType meta_optimizer_iterations = 12;
// The minimum number of nodes in a graph to optimizer. For smaller graphs,
// optimization is skipped.
// 0 means the system picks an appropriate number.
// < 0 means do not skip optimization.
int32 min_graph_nodes = 17;
// Disable optimizations that assume compressed tensors. Note that this flag
// is experimental and may be removed in the future.
bool experimental_disable_compressed_tensor_optimization = 26;
// Disable folding quantization emulation ops such as FakeQuantWithMinMax* and
// QuantizeAndDequantize*. Some compilers (e.g. the TF-to-tflite converter)
// have to extract quantization configs (e.g. min/max range, number of bits,
// and per-channel) from the quantization emulation ops. Note that this flag
// is experimental and may be removed in the future. See b/174138564 for more
// details.
bool experimental_disable_folding_quantization_emulation = 27;
enum MemOptType {
// The default setting (SCHEDULING and SWAPPING HEURISTICS only)
DEFAULT_MEM_OPT = 0;
// Disabled in the meta-optimizer.
NO_MEM_OPT = 1;
// Driven by manual op-level annotations.
MANUAL = 2;
// Driven by heuristics. The behavior of these heuristics is subject to
// change. Currently includes an experimental recomputation and swapping
// heuristics. Manual annotations are respected, but additional nodes are
// selected automatically.
// Swapping heuristic will move a tensor from the GPU to the CPU and move
// it back when needed to reduce peak memory usage.
SWAPPING_HEURISTICS = 4;
// Recomputation heuristics will recompute ops (such as Relu activation)
// during backprop instead of storing them, reducing peak memory usage.
RECOMPUTATION_HEURISTICS = 5;
// Scheduling will split big ops such as AddN and try to enforce a schedule
// of the new computations that decreases peak memory usage.
SCHEDULING_HEURISTICS = 6;
// Use any combination of swapping and recomputation heuristics.
HEURISTICS = 3;
}
// Configures memory optimization passes through the meta-optimizer. Has no
// effect on manually requested memory optimization passes in the optimizers
// field.
MemOptType memory_optimization = 4;
// A node name scope for node names which are valid outputs of recomputations.
// Inputs to nodes that match this scope may be recomputed (subject either to
// manual annotation of those input nodes or to manual annotation and
// heuristics depending on memory_optimization), but the nodes themselves will
// not be recomputed. This matches any sub-scopes as well, meaning the scope
// can appear not just as a top-level scope. For example, if the value is
// "gradients/", the default, it will match node name "gradients/foo",
// "foo/gradients/bar", but not "foo_gradients/"
string memory_optimizer_target_node_name_scope = 6;
// Maximum number of milliseconds to spend optimizing a single graph before
// timing out. If less than or equal to 0 (default value) the optimizer will
// never time out.
int64 meta_optimizer_timeout_ms = 20;
// Configures AutoParallel optimization passes either through the
// meta-optimizer or when manually specified through the optimizers field.
AutoParallelOptions auto_parallel = 5;
// If true, any optimization pass failing will cause the MetaOptimizer to
// stop with an error. By default - or when set to false, failing passes are
// skipped silently.
bool fail_on_optimizer_errors = 21;
ScopedAllocatorOptions scoped_allocator_opts = 16;
// If non-empty, will use this as an alternative way to specify a list of
// optimizations to turn on and the order of the optimizations (replacing the
// meta-optimizer).
//
// Of the RewriterConfig options, only the AutoParallel configuration options
// (the auto_parallel field) apply to manually requested optimization passes
// ("autoparallel"). Memory optimization passes ("memory") invoked here are
// not configurable (in contrast to memory optimization passes through the
// meta-optimizer) and act only on manual op annotations.
//
// Custom optimizers (see custom_optimizers) that are not part of this
// schedule will be run after - in the order that they were specified.
repeated string optimizers = 100;
// Message to describe custom graph optimizer and its parameters
message CustomGraphOptimizer {
string name = 1;
map<string, AttrValue> parameter_map = 2;
}
// list of CustomGraphOptimizers to apply.
repeated CustomGraphOptimizer custom_optimizers = 200;
// VerifierConfig specifying the verifiers to be run after every optimizer.
VerifierConfig inter_optimizer_verifier_config = 300;
// VerifierConfig specifying the verifiers to be run at the end, after all
// optimizers have run.
VerifierConfig post_optimization_verifier_config = 301;
}
tensorflow/core/protobuf/rpc_options.proto
+22
View File
@@ -0,0 +1,22 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.dummy;
import public "xla/tsl/protobuf/rpc_options.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
tensorflow/core/protobuf/saved_model.proto
+38
View File
@@ -0,0 +1,38 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/protobuf/meta_graph.proto";
option cc_enable_arenas = true;
option java_outer_classname = "SavedModelProtos";
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";
// SavedModel is the high level serialization format for TensorFlow Models.
// See [todo: doc links, similar to session_bundle] for more information.
message SavedModel {
// The schema version of the SavedModel instance. Used for versioning when
// making future changes to the specification/implementation. Initial value
// at release will be 1.
int64 saved_model_schema_version = 1;
// One or more MetaGraphs.
repeated MetaGraphDef meta_graphs = 2;
}
tensorflow/core/protobuf/saved_object_graph.proto
+266
View File
@@ -0,0 +1,266 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "google/protobuf/any.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";
import "tensorflow/core/framework/variable.proto";
import "tensorflow/core/framework/versions.proto";
import "tensorflow/core/protobuf/struct.proto";
import "tensorflow/core/protobuf/trackable_object_graph.proto";
option cc_enable_arenas = true;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// A SavedObjectGraph is part of object-based SavedModels in TF 2.0. It
// describes the directed graph of Python objects (or equivalent in other
// languages) that make up a model, with nodes[0] at the root.
// SavedObjectGraph shares some structure with TrackableObjectGraph, but
// SavedObjectGraph belongs to the MetaGraph and contains pointers to functions
// and type information, while TrackableObjectGraph lives in the checkpoint
// and contains pointers only to variable values.
message SavedObjectGraph {
// Flattened list of objects in the object graph.
//
// The position of the object in this list indicates its id.
// Nodes[0] is considered the root node.
repeated SavedObject nodes = 1;
// Information about captures and output structures in concrete functions.
// Referenced from SavedBareConcreteFunction and SavedFunction.
map<string, SavedConcreteFunction> concrete_functions = 2;
}
message SavedObject {
// Objects which this object depends on: named edges in the dependency
// graph.
//
// Note: All kinds of SavedObject may have children, except
// "constant" and "captured_tensor".
repeated TrackableObjectGraph.TrackableObject.ObjectReference children = 1;
// Ordered list of dependencies that must be loaded before this object.
// SavedModel loads with the bottom-up approach, by first creating all objects
// (in the order defined by the dependencies), then connecting the edges.
repeated TrackableObjectGraph.TrackableObject.ObjectReference dependencies =
15;
// Removed when forking SavedObject from TrackableObjectGraph.
reserved "attributes";
reserved 2;
// Slot variables owned by this object. This describes the three-way
// (optimizer, variable, slot variable) relationship; none of the three
// depend on the others directly.
//
// Note: currently only valid if kind == "user_object".
repeated TrackableObjectGraph.TrackableObject.SlotVariableReference
slot_variables = 3;
oneof kind {
SavedUserObject user_object = 4;
SavedAsset asset = 5;
SavedFunction function = 6;
SavedVariable variable = 7;
SavedBareConcreteFunction bare_concrete_function = 8;
SavedConstant constant = 9;
SavedResource resource = 10;
CapturedTensor captured_tensor = 12;
}
// Stores the functions used to save and restore this object. At most one of
// `saveable_objects` or `registered_saver` is defined for each SavedObject.
// See the comment below for the difference between SaveableObject and
// registered savers.
map<string, SaveableObject> saveable_objects = 11;
// The fields below are filled when the user serializes a registered Trackable
// class or an object with a registered saver function.
//
// Registered classes may save additional metadata and supersede the
// default loading process where nodes are recreated from the proto.
// If the registered class cannot be found, then the object will load as one
// one of the default trackable objects: Autotrackable (a class similar to
// tf.Module), tf.function, or tf.Variable.
//
// Unlike SaveableObjects, which store the functions for saving and restoring
// from tensors, registered savers allow Trackables to write checkpoint shards
// directly (e.g. for performance or coordination reasons).
// *All registered savers must be available when loading the SavedModel.*
// The name of the registered class of the form "{package}.{class_name}".
// This field is used to search for the registered class at loading time.
string registered_name = 13;
// The user-generated proto storing metadata for this object, to be passed to
// the registered classes's _deserialize_from_proto method when this object is
// loaded from the SavedModel.
google.protobuf.Any serialized_user_proto = 14;
// String name of the registered saver. At most one of `saveable_objects` or
// `registered_saver` is defined for each SavedObject.
string registered_saver = 16;
}
// A SavedUserObject is an object (in the object-oriented language of the
// TensorFlow program) of some user- or framework-defined class other than
// those handled specifically by the other kinds of SavedObjects.
//
// This object cannot be evaluated as a tensor, and therefore cannot be bound
// to an input of a function.
message SavedUserObject {
// Corresponds to a registration of the type to use in the loading program.
string identifier = 1;
// Version information from the producer of this SavedUserObject.
VersionDef version = 2;
// Metadata for deserializing this object.
//
// Deprecated! At the time of deprecation, Keras was the only user of this
// field, and its saving and loading code will be updated shortly.
// Please save your application-specific metadata to a separate file.
string metadata = 3 [deprecated = true];
}
// A SavedAsset points to an asset in the MetaGraph.
//
// When bound to a function this object evaluates to a tensor with the absolute
// filename. Users should not depend on a particular part of the filename to
// remain stable (e.g. basename could be changed).
message SavedAsset {
// Index into `MetaGraphDef.asset_file_def[]` that describes the Asset.
//
// Only the field `AssetFileDef.filename` is used. Other fields, such as
// `AssetFileDef.tensor_info`, MUST be ignored.
int32 asset_file_def_index = 1;
}
// A function with multiple signatures, possibly with non-Tensor arguments.
message SavedFunction {
repeated string concrete_functions = 1;
FunctionSpec function_spec = 2;
}
message CapturedTensor {
// Name of captured tensor
string name = 1;
// Name of concrete function which contains the computed graph tensor.
string concrete_function = 2;
}
// Stores low-level information about a concrete function. Referenced in either
// a SavedFunction or a SavedBareConcreteFunction.
message SavedConcreteFunction {
repeated int32 bound_inputs = 2;
// Input in canonicalized form that was received to create this concrete
// function.
StructuredValue canonicalized_input_signature = 3;
// Output that was the return value of this function after replacing all
// Tensors with TensorSpecs. This can be an arbitrary nested function and will
// be used to reconstruct the full structure from pure tensors.
StructuredValue output_signature = 4;
}
message SavedBareConcreteFunction {
// Identifies a SavedConcreteFunction.
string concrete_function_name = 1;
// A sequence of unique strings, one per Tensor argument.
repeated string argument_keywords = 2;
// The prefix of `argument_keywords` which may be identified by position.
int64 allowed_positional_arguments = 3;
// The spec of the function that this ConcreteFunction is traced from. This
// allows the ConcreteFunction to be called with nest structure inputs. This
// field may not be populated. If this field is absent, the concrete function
// can only be called with flat inputs.
// TODO(b/169361281): support calling saved ConcreteFunction with structured
// inputs in C++ SavedModel API.
FunctionSpec function_spec = 4;
}
message SavedConstant {
// An Operation name for a ConstantOp in this SavedObjectGraph's MetaGraph.
string operation = 1;
}
// Represents a Variable that is initialized by loading the contents from the
// checkpoint.
message SavedVariable {
DataType dtype = 1;
TensorShapeProto shape = 2;
bool trainable = 3;
VariableSynchronization synchronization = 4;
VariableAggregation aggregation = 5;
string name = 6;
string device = 7;
// List of component variables for a distributed variable.
//
// When this field is non-empty, the SavedVariable will be assumed
// to be a distributed variable defined by the components listed here.
//
// This is only supported by experimental loaders at the moment.
repeated SavedVariable experimental_distributed_variable_components = 8;
}
// Represents `FunctionSpec` used in `Function`. This represents a
// function that has been wrapped as a TensorFlow `Function`.
message FunctionSpec {
// Full arg spec from inspect.getfullargspec().
StructuredValue fullargspec = 1;
// Whether this represents a class method.
bool is_method = 2;
// The input signature, if specified.
StructuredValue input_signature = 5;
// Whether the function should be compiled by XLA.
//
// The public interface to `tf.function` uses an optional boolean to
// represent three distinct states for this field. Unfortunately, proto3
// removes the ability to explicitly check for the presence or absence of a
// field, so we instead map to an enum.
//
// See `tf.function` for details.
enum JitCompile {
DEFAULT = 0;
ON = 1;
OFF = 2;
}
JitCompile jit_compile = 6;
reserved 3, 4;
}
// A SavedResource represents a TF object that holds state during its lifetime.
// An object of this type can have a reference to a:
// create_resource() and an initialize() function.
message SavedResource {
// A device specification indicating a required placement for the resource
// creation function, e.g. "CPU". An empty string allows the user to select a
// device.
string device = 1;
}
message SaveableObject {
// Node ids of concrete functions for saving and loading from a checkpoint.
// These functions save and restore directly from tensors.
int32 save_function = 2;
int32 restore_function = 3;
}
tensorflow/core/protobuf/saver.proto
+63
View File
@@ -0,0 +1,63 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
option java_outer_classname = "SaverProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.util";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Protocol buffer representing the configuration of a Saver.
message SaverDef {
// The name of the tensor in which to specify the filename when saving or
// restoring a model checkpoint.
string filename_tensor_name = 1;
// The operation to run when saving a model checkpoint.
string save_tensor_name = 2;
// The operation to run when restoring a model checkpoint.
string restore_op_name = 3;
// Maximum number of checkpoints to keep. If 0, no checkpoints are deleted.
int32 max_to_keep = 4;
// Shard the save files, one per device that has Variable nodes.
bool sharded = 5;
// How often to keep an additional checkpoint. If not specified, only the last
// "max_to_keep" checkpoints are kept; if specified, in addition to keeping
// the last "max_to_keep" checkpoints, an additional checkpoint will be kept
// for every n hours of training.
float keep_checkpoint_every_n_hours = 6;
// A version number that identifies a different on-disk checkpoint format.
// Usually, each subclass of BaseSaverBuilder works with a particular
// version/format. However, it is possible that the same builder may be
// upgraded to support a newer checkpoint format in the future.
enum CheckpointFormatVersion {
// Internal legacy format.
LEGACY = 0;
// Deprecated format: tf.Saver() which works with tensorflow::table::Table.
V1 = 1;
// Current format: more efficient.
V2 = 2;
}
CheckpointFormatVersion version = 7;
}
tensorflow/core/protobuf/service_config.proto
+126
View File
@@ -0,0 +1,126 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.data.experimental;
import "tensorflow/core/protobuf/data_service.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Configuration for a tf.data service DispatchServer.
// Next id: 13
message DispatcherConfig {
// The port for the dispatcher to bind to. A value of 0 indicates that the
// dispatcher may bind to any available port.
int64 port = 1;
// The protocol for the dispatcher to use when connecting to workers.
string protocol = 2;
// A work directory to use for storing dispatcher state, and for recovering
// during restarts. The empty string indicates not to use any work directory.
string work_dir = 3;
// Whether to run in fault tolerant mode, where dispatcher state is saved
// across restarts. Requires that `work_dir` is nonempty.
bool fault_tolerant_mode = 4;
// (Optional.) If the job uses auto-sharding, it needs to specify a fixed list
// of worker addresses that will register with the dispatcher. The worker
// addresses should be in the format "host" or "host:port", where "port" is an
// integer, named port, or %port% to match any port.
repeated string worker_addresses = 7;
// (Optional.) tf.data service deployment mode. Supported values are "REMOTE",
// "COLOCATED", and "HYBRID". If unspecified, it is assumed to be "REMOTE".
DeploymentMode deployment_mode = 9;
// How often the dispatcher should scan through to delete old and unused
// jobs. A value of 0 indicates that the decision should be left up to the
// runtime.
int64 job_gc_check_interval_ms = 5;
// How long a job needs to be unused before it becomes a candidate for garbage
// collection. A value of -1 indicates that jobs should never be garbage
// collected. A value of 0 indicates that the decision should be left up to
// the runtime. Note: This does not apply to dynamic sharding unless users
// explicitly opt-in by enabling `gc_dynamic_sharding_jobs` below.
int64 job_gc_timeout_ms = 6;
// Whether dynamically sharded jobs should be eligible for garbage collection.
// These jobs are not garbage collected by default, since if a job is garbage
// collected and then re-created, it will revisit all data from the start. If
// revisiting data is acceptible and you want automatic reclamation of
// iterator memory, set `gc_dynamic_sharding_jobs` to `true`.
bool gc_dynamic_sharding_jobs = 11;
// How long to wait before garbage-collecting a client that hasn't
// heartbeated to the dispatcher. A value of 0 indicates that the timeout
// should be left to the runtime.
int64 client_timeout_ms = 8;
// How long to wait for a worker to heartbeat before considering it missing.
// A value of 0 indicates that the timeout should be left to the runtime.
int64 worker_timeout_ms = 10;
// The maximum number of snapshots that a worker can concurrently process at a
// given point in time. This is a tradeoff between worker resource usage and
// snapshot wall time. A value of 0 indicates that the decision should be left
// up to the runtime.
int64 worker_max_concurrent_snapshots = 12;
}
// Configuration for a tf.data service WorkerServer.
// Next id: 14
message WorkerConfig {
// The port for the worker to bind to. A value of 0 indicates that the
// worker may bind to any available port.
int64 port = 1;
// The protocol for the worker to use when connecting to the dispatcher.
string protocol = 2;
// The address of the dispatcher to register with.
string dispatcher_address = 3;
// The address of the worker server. The substring "%port%", if specified,
// will be replaced with the worker's bound port. This is useful when the port
// is set to `0`.
string worker_address = 4;
// Tags attached to the worker. This allows reading from selected workers.
// For example, by applying a "COLOCATED" tag, tf.data service is able to read
// from the local tf.data worker if one exists, then from off-TF-host workers,
// to avoid cross-TF-host reads.
repeated string worker_tags = 10;
// How often the worker should heartbeat to the master. A value of 0 indicates
// that the decision should be left up to the runtime.
int64 heartbeat_interval_ms = 5;
// How long to retry requests to the dispatcher before giving up and reporting
// an error. A value of 0 indicates that the decision should be left up to the
// runtime.
int64 dispatcher_timeout_ms = 6;
// If set, the name of an alternative data transfer protocol for which the
// worker starts an additional server ("data transfer server"); the trainer
// can then get data from this server. If not set, no such server is started,
// and the trainer can only get data from the regular worker server over
// `protocol`.
string data_transfer_protocol = 7;
// If `data_transfer_protocol` is set, the port to which the data transfer
// server binds. If set to `0`, the server binds to any available port.
int64 data_transfer_port = 13;
// If `data_transfer_protocol` is set, the address of the data transfer
// server. The substring "%dts_port%" can be used to represent -- and is
// replaced with -- the bound port of the data transfer server; this is useful
// when `data_transfer_port` is set to `0`.
string data_transfer_address = 8;
// Maximum size of the cross-trainer cache in bytes. If enabled, make sure
// your training job provides sufficient memory resources.
int64 cross_trainer_cache_size_bytes = 11;
// The maximum size of a distributed snapshot chunk file. A value of 0
// indicates that the decision should be left up to the runtime.
int64 snapshot_max_chunk_size_bytes = 12;
// When shutting down a worker, how long to wait for the gRPC server to
// process the final requests. This is used to achieve clean shutdown in unit
// tests.
int64 shutdown_quiet_period_ms = 9;
}
tensorflow/core/protobuf/snapshot.proto
+73
View File
@@ -0,0 +1,73 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.data.experimental;
import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Each SnapshotRecord represents one batch of pre-processed input data. A batch
// consists of a list of tensors that we encode as TensorProtos. This message
// doesn't store the structure of the batch.
message SnapshotRecord {
repeated .tensorflow.TensorProto tensor = 1;
}
// This stores the metadata information present in each snapshot record.
message SnapshotMetadataRecord {
// Stores the fingerprint of the graph that describes the dataset that is
// snapshotted.
string graph_hash = 1;
// Run ID that this snapshot corresponds to.
string run_id = 2;
// Time when we started creating this snapshot.
int64 creation_timestamp = 3;
// Version of the snapshot data file format.
int64 version = 4;
// A list of tensor dtype corresponding to each element of the snapshot.
repeated .tensorflow.DataType dtype = 5;
// The number of elements in the snapshot.
int64 num_elements = 6;
bool finalized = 1000;
}
// Metadata for a single tensor in the Snapshot Record.
message TensorMetadata {
.tensorflow.TensorShapeProto tensor_shape = 2;
// Number of uncompressed bytes used to store the tensor representation.
int64 tensor_size_bytes = 3;
}
// Metadata for all the tensors in a Snapshot Record.
message SnapshotTensorMetadata {
repeated TensorMetadata tensor_metadata = 1;
}
// Metadata for a `tf.data.Dataset` distributed snapshot.
message DistributedSnapshotMetadata {
// The element spec of the snapshotted dataset.
bytes element_spec = 1;
// Whether and how to compress the snapshot. Supported values are defined in
// `tsl::io::compression`. In particular, an empty string specifies not to
// compress.
string compression = 2;
}
tensorflow/core/protobuf/status.proto
+26
View File
@@ -0,0 +1,26 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
// Add a dummy package name. Having no package like
// core/lib/core/error_codes.proto, or having tensorflow like
// tsl/protobuf/error_codes.proto, results in name collision errors in generated
// code for some users that use JS through J2CL.
package tensorflow.dummy;
import public "xla/tsl/protobuf/status.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
tensorflow/core/protobuf/struct.proto
+179
View File
@@ -0,0 +1,179 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// `StructuredValue` represents a dynamically typed value representing various
// data structures that are inspired by Python data structures typically used in
// TensorFlow functions as inputs and outputs.
//
// For example when saving a Layer there may be a `training` argument. If the
// user passes a boolean True/False, that switches between two concrete
// TensorFlow functions. In order to switch between them in the same way after
// loading the SavedModel, we need to represent "True" and "False".
//
// A more advanced example might be a function which takes a list of
// dictionaries mapping from strings to Tensors. In order to map from
// user-specified arguments `[{"a": tf.constant(1.)}, {"q": tf.constant(3.)}]`
// after load to the right saved TensorFlow function, we need to represent the
// nested structure and the strings, recording that we have a trace for anything
// matching `[{"a": tf.TensorSpec(None, tf.float32)}, {"q": tf.TensorSpec([],
// tf.float64)}]` as an example.
//
// Likewise functions may return nested structures of Tensors, for example
// returning a dictionary mapping from strings to Tensors. In order for the
// loaded function to return the same structure we need to serialize it.
//
// This is an ergonomic aid for working with loaded SavedModels, not a promise
// to serialize all possible function signatures. For example we do not expect
// to pickle generic Python objects, and ideally we'd stay language-agnostic.
message StructuredValue {
// The kind of value.
oneof kind {
// Represents None.
NoneValue none_value = 1;
// Represents a double-precision floating-point value (a Python `float`).
double float64_value = 11;
// Represents a signed integer value, limited to 64 bits.
// Larger values from Python's arbitrary-precision integers are unsupported.
sint64 int64_value = 12;
// Represents a string of Unicode characters stored in a Python `str`.
// In Python 3, this is exactly what type `str` is.
// In Python 2, this is the UTF-8 encoding of the characters.
// For strings with ASCII characters only (as often used in TensorFlow code)
// there is effectively no difference between the language versions.
// The obsolescent `unicode` type of Python 2 is not supported here.
string string_value = 13;
// Represents a boolean value.
bool bool_value = 14;
// Represents a TensorShape.
tensorflow.TensorShapeProto tensor_shape_value = 31;
// Represents an enum value for dtype.
tensorflow.DataType tensor_dtype_value = 32;
// Represents a value for tf.TensorSpec.
TensorSpecProto tensor_spec_value = 33;
// Represents a value for tf.TypeSpec.
TypeSpecProto type_spec_value = 34;
// Represents a value for tf.BoundedTensorSpec.
BoundedTensorSpecProto bounded_tensor_spec_value = 35;
// Represents a list of `Value`.
ListValue list_value = 51;
// Represents a tuple of `Value`.
TupleValue tuple_value = 52;
// Represents a dict `Value`.
DictValue dict_value = 53;
// Represents Python's namedtuple.
NamedTupleValue named_tuple_value = 54;
// Represents a value for tf.Tensor.
tensorflow.TensorProto tensor_value = 55;
// Represents a value for np.ndarray.
tensorflow.TensorProto numpy_value = 56;
}
}
// Represents None.
message NoneValue {}
// Represents a Python list.
message ListValue {
repeated StructuredValue values = 1;
}
// Represents a Python tuple.
message TupleValue {
repeated StructuredValue values = 1;
}
// Represents a Python dict keyed by `str`.
// The comment on Unicode from Value.string_value applies analogously.
message DictValue {
map<string, StructuredValue> fields = 1;
}
// Represents a (key, value) pair.
message PairValue {
string key = 1;
StructuredValue value = 2;
}
// Represents Python's namedtuple.
message NamedTupleValue {
string name = 1;
repeated PairValue values = 2;
}
// A protobuf to represent tf.TensorSpec.
message TensorSpecProto {
string name = 1;
tensorflow.TensorShapeProto shape = 2;
tensorflow.DataType dtype = 3;
}
// A protobuf to represent tf.BoundedTensorSpec.
message BoundedTensorSpecProto {
string name = 1;
tensorflow.TensorShapeProto shape = 2;
tensorflow.DataType dtype = 3;
tensorflow.TensorProto minimum = 4;
tensorflow.TensorProto maximum = 5;
}
// Represents a tf.TypeSpec
message TypeSpecProto {
enum TypeSpecClass {
UNKNOWN = 0;
SPARSE_TENSOR_SPEC = 1; // tf.SparseTensorSpec
INDEXED_SLICES_SPEC = 2; // tf.IndexedSlicesSpec
RAGGED_TENSOR_SPEC = 3; // tf.RaggedTensorSpec
TENSOR_ARRAY_SPEC = 4; // tf.TensorArraySpec
DATA_DATASET_SPEC = 5; // tf.data.DatasetSpec
DATA_ITERATOR_SPEC = 6; // IteratorSpec from data/ops/iterator_ops.py
OPTIONAL_SPEC = 7; // tf.OptionalSpec
PER_REPLICA_SPEC = 8; // PerReplicaSpec from distribute/values.py
VARIABLE_SPEC = 9; // tf.VariableSpec
ROW_PARTITION_SPEC = 10; // RowPartitionSpec from ragged/row_partition.py
reserved 11;
REGISTERED_TYPE_SPEC = 12; // The type registered as type_spec_class_name.
EXTENSION_TYPE_SPEC = 13; // Subclasses of tf.ExtensionType
}
TypeSpecClass type_spec_class = 1;
// The value returned by TypeSpec._serialize().
StructuredValue type_state = 2;
// The name of the TypeSpec class.
// * If type_spec_class == REGISTERED_TYPE_SPEC, the TypeSpec class is
// the one registered under this name. For types registered outside
// core TensorFlow by an add-on library, that library must be loaded
// before this value can be deserialized by nested_structure_coder.
// * If type_spec_class specifies a particular TypeSpec class, this field is
// redundant with the type_spec_class enum, and is only used for error
// reporting in older binaries that do not know the tupe_spec_class enum.
string type_spec_class_name = 3;
// The number of flat tensor components required by this TypeSpec.
int32 num_flat_components = 4;
}
tensorflow/core/protobuf/tensor_bundle.proto
+81
View File
@@ -0,0 +1,81 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/tensor_slice.proto";
import "tensorflow/core/framework/types.proto";
import "tensorflow/core/framework/versions.proto";
option cc_enable_arenas = true;
option java_outer_classname = "TensorBundleProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.util";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Protos used in the tensor bundle module (tf/core/util/tensor_bundle/).
// Special header that is associated with a bundle.
//
// TODO(zongheng,zhifengc): maybe in the future, we can add information about
// which binary produced this checkpoint, timestamp, etc. Sometime, these can be
// valuable debugging information. And if needed, these can be used as defensive
// information ensuring reader (binary version) of the checkpoint and the writer
// (binary version) must match within certain range, etc.
message BundleHeaderProto {
// Number of data files in the bundle.
int32 num_shards = 1;
// An enum indicating the endianness of the platform that produced this
// bundle. A bundle can only be read by a platform with matching endianness.
// Defaults to LITTLE, as most modern platforms are little-endian.
//
// Affects the binary tensor data bytes only, not the metadata in protobufs.
enum Endianness {
LITTLE = 0;
BIG = 1;
}
Endianness endianness = 2;
// Versioning of the tensor bundle format.
VersionDef version = 3;
}
// Describes the metadata related to a checkpointed tensor.
message BundleEntryProto {
// The tensor dtype and shape.
DataType dtype = 1;
TensorShapeProto shape = 2;
// The binary content of the tensor lies in:
// File "shard_id": bytes [offset, offset + size).
int32 shard_id = 3;
int64 offset = 4;
int64 size = 5;
// The CRC32C checksum of the tensor bytes.
fixed32 crc32c = 6;
// Iff present, this entry represents a partitioned tensor. The previous
// fields are interpreted as follows:
//
// "dtype", "shape": describe the full tensor.
// "shard_id", "offset", "size", "crc32c": all IGNORED.
// These information for each slice can be looked up in their own
// BundleEntryProto, keyed by each "slice_name".
repeated TensorSliceProto slices = 7;
}
tensorflow/core/protobuf/tensorflow_server.proto
+64
View File
@@ -0,0 +1,64 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
syntax = "proto3";
package tensorflow;
import "tensorflow/core/protobuf/cluster.proto";
import "tensorflow/core/protobuf/config.proto";
import "tensorflow/core/protobuf/device_filters.proto";
option cc_enable_arenas = true;
option java_outer_classname = "ServerProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.distruntime";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Defines the configuration of a single TensorFlow server.
message ServerDef {
// The cluster of which this server is a member.
ClusterDef cluster = 1;
// The name of the job of which this server is a member.
//
// NOTE(mrry): The `cluster` field must contain a `JobDef` with a `name` field
// that matches this name.
string job_name = 2;
// Replica this server manages.
int32 replica = 8;
// The task index of this server in its job.
//
// NOTE: The `cluster` field must contain a `JobDef` with a matching `name`
// and a mapping in its `tasks` field for this index.
int32 task_index = 3;
// The default configuration for sessions that run on this server.
ConfigProto default_session_config = 4;
// The protocol to be used by this server.
//
// Acceptable values include: "grpc", "grpc+verbs".
string protocol = 5;
// The server port. If not set, then we identify the port from the job_name.
int32 port = 6;
// Device filters for remote tasks in the cluster.
// NOTE: This is an experimental feature and only effective in TensorFlow 2.x.
ClusterDeviceFilters cluster_device_filters = 7;
}
tensorflow/core/protobuf/tpu/BUILD
+119
View File
@@ -0,0 +1,119 @@
load("@rules_python//python:proto.bzl", "py_proto_library")
load(
"//tensorflow/core/platform:build_config.bzl",
"tf_additional_all_protos",
"tf_proto_library",
"tf_pyclif_proto_library",
)
package(
# copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
licenses = ["notice"],
)
tf_proto_library(
name = "tpu_embedding_configuration_proto",
srcs = [
"tpu_embedding_configuration.proto",
],
protodeps = [
":optimization_parameters_proto",
],
visibility = ["//visibility:public"],
)
tf_proto_library(
name = "optimization_parameters_proto",
srcs = [
"optimization_parameters.proto",
],
protodeps = ["@xla//xla/service:hlo_proto"],
visibility = ["//visibility:public"],
)
tf_proto_library(
name = "topology_proto",
srcs = [
"topology.proto",
],
visibility = ["//visibility:public"],
)
tf_proto_library(
name = "dynamic_padding_proto",
srcs = [
"dynamic_padding.proto",
],
visibility = ["//visibility:public"],
)
tf_proto_library(
name = "compilation_result_proto",
srcs = [
"compilation_result.proto",
],
protodeps = tf_additional_all_protos() + [
"@xla//xla:xla_data_proto",
"@xla//xla/service:hlo_proto",
],
visibility = ["//visibility:public"],
)
tf_proto_library(
name = "compile_metadata_proto",
srcs = [
"compile_metadata.proto",
],
protodeps = [
":dynamic_padding_proto",
"@xla//xla:xla_data_proto",
"@xla//xla:xla_proto",
"//tensorflow/core:protos_all",
],
visibility = ["//visibility:public"],
)
tf_pyclif_proto_library(
name = "topology_pyclif",
proto_lib = ":topology_proto",
proto_srcfile = "topology.proto",
visibility = ["//visibility:public"],
)
# copybara:uncomment_begin(google-only)
# py_proto_library(
# name = "tpu_embedding_configuration_py_pb2",
# visibility = ["//visibility:public"],
# deps = [":tpu_embedding_configuration_proto"],
# )
#
# py_proto_library(
# name = "optimization_parameters_py_pb2",
# visibility = ["//visibility:public"],
# deps = [":optimization_parameters_proto"],
# )
#
# py_proto_library(
# name = "topology_py_pb2",
# visibility = ["//visibility:public"],
# deps = [":topology_proto"],
# )
#
# py_proto_library(
# name = "dynamic_padding_py_pb2",
# visibility = ["//visibility:public"],
# deps = [":dynamic_padding_proto"],
# )
#
# py_proto_library(
# name = "compilation_result_py_pb2",
# visibility = ["//visibility:public"],
# deps = [":compilation_result_proto"],
# )
#
# py_proto_library(
# name = "compile_metadata_py_pb2",
# visibility = ["//visibility:public"],
# deps = [":compile_metadata_proto"],
# )
# copybara:uncomment_end
tensorflow/core/protobuf/tpu/compilation_result.proto
+41
View File
@@ -0,0 +1,41 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.tpu;
import "xla/service/hlo.proto";
import "tensorflow/core/protobuf/error_codes.proto";
option cc_enable_arenas = true;
// Describes the result of a TPU compilation. This is also used as TPU
// compilation result status payload.
// URI: "type.googleapis.com/tensorflow.tpu.CompilationResultProto"
message CompilationResultProto {
// The error message, if any, returned during compilation.
error.Code status_code = 1;
string status_error_message = 2;
// HLO proto.
repeated xla.HloProto hlo_protos = 3;
enum ErrorCode {
UNKNOWN = 0;
OUT_OF_MEMORY = 1;
}
ErrorCode error_code = 4;
}
tensorflow/core/protobuf/tpu/compile_metadata.proto
+178
View File
@@ -0,0 +1,178 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.tpu;
import "xla/xla.proto";
import "xla/xla_data.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";
import "tensorflow/core/protobuf/tpu/dynamic_padding.proto";
option cc_enable_arenas = true;
// This is an experimental proto used in the TF/XLA bridge to store metadata to
// a compile op (e.g. _TPUCompileMlir).
// TODO(lyandy): Deprecate proto once generic metadata proto is created.
message TPUCompileMetadataProto {
// Description of the types and shapes of the arguments to a computation.
message Arg {
enum Kind {
INVALID = 0;
PARAMETER = 1;
VARIABLE = 2;
// These are args which have been guaranteed to be constants during the
// session lifetime by the use of the GuaranteeConstOp (or ConstantOp).
GUARANTEED_CONSTANT = 3;
}
DataType dtype = 1;
TensorShapeProto shape = 2;
Kind kind = 3;
// The cross-core sharding of this input within each replica, e.g.,
// assigning to one core, or replicate across all cores.
xla.OpSharding sharding = 4;
// Whether this argument will receive the same data across all replicas.
bool is_same_data_across_replicas = 5;
enum EnableXlaSharding {
DISALLOWED = 0;
// Sharding is allowed if host training loop exists.
TENTATIVE = 1;
ALLOWED = 2;
}
// Whether to allow XLA to produce separate programs to shard/unshard this
// argument. Requires this arg to be an on-device Kind::VARIABLE, or a
// Kind::PARAMETER. For Kind::PARAMETER, it represents the initial value of
// a variable, and retval_index_for_sharding must be specified for the
// corresponding updated value.
EnableXlaSharding enable_xla_sharding = 6;
// If XLA sharding is allowed on a Kind::PARAMETER, this field is used to
// specify the corresponding updated value in the return values. Use -1 for
// variables that are not updated.
int32 retval_index_for_sharding = 8;
// Whether this argument is placed on fast memory or not.
bool fast_mem = 7;
// Whether to let XLA to decide the layout during compilation, as opposed to
// using a fixed layout determined by the shape.
bool unrestricted_layout = 9;
// Name of the node that the arg comes from.
string name = 10;
// Whether to use XLA collectives to broadcast this parameter to all
// replicas, instead of using TensorFlow Send/Recv among the tasks.
bool requires_xla_broadcast = 11;
// Which dimension of this arg is bounded dynamic.
repeated bool is_bounded_dynamic_dim = 12;
}
repeated Arg args = 1;
// Description of the return values from a computation.
message Retval {
// The cross-core sharding of this return value within each replica, e.g.,
// assigning to one core, or replicate across all cores.
xla.OpSharding sharding = 1;
}
repeated Retval retvals = 2;
// Number of replicas of the computation and number of cores in each replica.
// TODO(b/140721404): it may not be necessary to state the number of cores per
// replica here. Reconsider when replicated model-parallelism is implemented
// in XLA.
int32 num_replicas = 3;
int32 num_cores_per_replica = 4;
reserved 5; // was device_names
reserved 7; // was replica_device_assignment
xla.DeviceAssignmentProto device_assignment = 8;
// A fingerprint of the function library. Ensures that any functions called
// by the computation have matching definitions.
uint64 function_library_fingerprint = 6;
// Unique session identifier. Can be empty.
string session_handle = 9;
// Fingerprint of guaranteed_const value. The fingerprint computation inside
// tpu_compile_op may be slow. The computation can be avoided by setting the
// fingerprint value here.
string guaranteed_const_fingerprint = 10;
repeated tpu.PaddingMap padding_maps = 11;
// The location of step markers that XLA compile will instrument.
xla.DebugOptions.StepMarkerLocation step_marker_location = 12;
// Minimum number of batches run through the XLA graph before XLA fusion
// autotuner is enabled. Default value of zero disables the autotuner.
// The XLA fusion autotuner can improve performance by executing a heuristic
// search on the compiler parameters.
int64 xla_fusion_autotuner_thresh = 13;
// Enables TPU compiler to add partitioning policies for inputs/outputs to
// the XLA computation for model parallelism.
bool enable_automatic_model_parallelism = 14;
// Whether to use XLA's SPMD or MPMD partitioner when compiler partitioning is
// requested.
bool use_spmd_for_xla_partitioning = 15;
// Whether to use Shardy partitioner when compiler partitioning is requested.
bool use_shardy_partitioner = 23;
// Whether to automatically generate XLA shardings for SPMD partitioner.
bool use_auto_spmd_for_xla_partitioning = 18;
// Device mesh shape used to create the sharding search space when
// use_auto_spmd_partitioning=true.
repeated int64 auto_spmd_mesh_shape = 19;
// Device mesh ids compatible with the above mesh_shape used when
// use_auto_spmd_partitioning=true.
repeated int64 auto_spmd_mesh_ids = 20;
reserved 16; // Was broadcast_replicated_parameters_via_collectives
// A fingerprint generated by hashing the MLIR module content.
uint64 mlir_fingerprint = 17;
TPUCompileOptions compile_options = 21;
// The name of the MLIR module.
string module_name = 22;
}
// Stable protobuf for TPU compilation options, suitable for persistent storage.
// This proto needs to be backward compatible under maintenance.
// TODO(timshen): investigate and migrate other options from
// TPUCompileMetadataProto.
message TPUCompileOptions {
enum Precision {
DEFAULT = 0;
BFLOAT16 = 1;
FLOAT32 = 2;
TENSOR_FLOAT32 = 3;
}
Precision matrix_unit_operand_precision = 1;
}
tensorflow/core/protobuf/tpu/dynamic_padding.proto
+34
View File
@@ -0,0 +1,34 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.tpu;
option cc_enable_arenas = true;
// A mapping between the dynamic shape dimension of an input and the arg that
// represents the real shape.
message PaddingMap {
// Input arg index with dynamic shapes.
int32 arg_index = 1;
// The dynamic shape dimension index.
int32 shape_index = 2;
// The arg index that dynamic dimension maps to, which represents the value
// of the real shape.
int32 padding_arg_index = 3;
}
tensorflow/core/protobuf/tpu/optimization_parameters.proto
+657
View File
@@ -0,0 +1,657 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.tpu;
import "google/protobuf/wrappers.proto";
import "xla/service/hlo.proto";
message ClippingLimits {
google.protobuf.FloatValue lower = 1; // -inf if not set
google.protobuf.FloatValue upper = 2; // +inf if not set
}
// Configuration for simulated quantization; simulated quantization is used to
// reduce training/serving skew when the serving variables are quantized. The
// same quantization operations are executed during training to minimize
// differences with serving.
//
// Simulated quantization inserts the following operations on the forward pass
// after gathering the embedding vector from HBM. The backward pass operations
// are unchanged.
//
// clipped_val = clip(input, clipping_limits)
// quantum = clipping_limits.range() / (num_buckets - 1)
// quantized_val = floor((clipped_val - clipping_limits.lower()) / quantum + .5)
// return quantized_val * quantum + clipping_limits.lower().
message SimulatedQuantization {
// Whether simulated quantization is enabled.
bool enabled = 1;
// Minimum and maximum values of the range used for quantization.
ClippingLimits clipping_limits = 2;
// Number of possible quantized values.
int32 num_buckets = 3;
}
// Dynamic input specification for optimizers in the TPUEmbeddingConfiguration.
// The actual dynamic inputs are provided as a scalar input list to the
// SendTPUEmbeddingGradients Op indexed by their tag specified through the
// following proto.
message OptimizerDynamicInput {
// For tables where dynamic inputs are needed (e.g., learning rates or other
// dynamic hyperparameters used in optimizers), a tag must be specified for
// the input.
//
// The tag must be a non-negative integer. All tags in the range
// [0, number_of_unique_tags) must be present in the TPU embedding
// configuration, i.e. a tag cannot be skipped if a different tag numerically
// greater than it is used in the configuration.
//
// If multiple tables specify the same tag, they *MUST* have
// the same dynamic input, for example, their dynamic learning rate could be
// computed by the same TensorFlow sub-graph. The partitioning of the
// embedding layer would be more optimal if the number_of_unique_tags is as
// *LOW* as possible, i.e., if many tables share the same tag.
//
// The hyper_parameters input of the SendTPUEmbeddingGradients op is used to
// communicate dynamic hyper-parameters to the TPU embedding program.
// The hyper_parameters input is a list of scalars where the size of the list
// is equal to the number of unique tags. The hyper-parameter associated with
// a particular tag is specified by populating its corresponding index in the
// list of scalars.
int32 tag = 1;
}
// Source of learning rate to use.
message LearningRate {
oneof learning_rate {
float constant = 1;
OptimizerDynamicInput dynamic = 2;
}
}
// Each optimizer's parameter proto has a link to its documentation and CPU
// implementation (if available) for user reference.
// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adagrad
// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1634
message AdagradParameters {
// Old initial accumulator parameter.
reserved "initial_accumulator";
reserved 1;
}
// This optimizer combines the Adagrad and Momentum update rules.
// accum(new) = beta2 == 1.0 ?
// accum(old) + grad^2 :
// beta2 * accum(old) + (1 - beta2) * grad^2
// accum_with_exponent = (accum(new) + epsilon)^(-1.0 / exponent)
// mom_accum(new) = momentum * mom_accum(old) + accum_with_exponent
// update = use_nesterov ?
// momentum * mom_accum(new) + accum_with_exponent :
// mom_accum(new)
// var(new) = var(old) - lr * grad * update
// Algorithm described in https://arxiv.org/abs/2002.11803.
message AdagradMomentumParameters {
// Moving average parameter for the momentum accumulator.
float momentum = 1;
// Whether to use the Nesterov variant of momentum.
bool use_nesterov = 2;
// Exponent for the gradient^2 accumulator.
float exponent = 3;
// Moving average parameter for the gradient^2 accumulator.
float beta2 = 4;
// Offset added to the Adagrad accumulator.
float epsilon = 5;
}
// Algorithm in http://www.jmlr.org/papers/volume12/duchi11a/duchi11a.pdf.
message BoundedAdagradParameters {
// Whether to use the updated or the old value of the accumulator when
// computing the effective learning rate. When update_accumulator_first is set
// to True, the updated value of the accumulator is used.
bool update_accumulator_first = 1;
// The max_var_update value to use. Set value to 0 (default) to disable using
// max_var_update to clip the gradient.
float max_var_update = 2;
// The maximum value of the accumulator. Set max_accumulator to 0 (default)
// to disable using max_accumulator to clip the accumulator.
float max_accumulator = 3;
}
// Frequency Aware Adagrad optimizer. This optimizer implements the AdaGrad
// algorithm and further allows to:
// * Scale the learning rate based on frequency of the update. Sparsely updated
// rows are updated with a higher effective learning rate, and frequently
// updated rows are updated with a lower effective learning rate.
// * Decay the growth of the accumulator values.
// * Use L1 / L2 regularization for the weight updates.
//
// The optimization algorithm is shown below.
// counter(new) = counter(old) + 1
// accum(new) = max(accumulator_decay * accum(old) + grad^2,
// initial_accumulator_value)
// lr_scale = min((step_counter / accum(new)) ^ probability_exponent,
// max_lr_multiplier) update = grad * lr_scale / sqrt(accum(new)) if
// (l1_regularization_strength > 0.0):
// update = update + l1_regularization_strength * sign(var(old))
// if (l2_regularization_strength > 0.0):
// update = update + l2_regularization_strength * var(old)
// var(new) = var(old) - lr_scale * grad * update
message FrequencyAwareAdagradParameters {
// The L1 regularization parameter for adjusting the update based on the sign
// of the variable.
float l1_regularization_strength = 1;
// The L2 regularization parameter for adjusting the update based on the
// variable.
float l2_regularization_strength = 2;
// The exponent used for scaling the learning rate based on the sparsity of
// updates.
float probability_exponent = 4;
// The maximum value of the learning rate scale.
float max_lr_multiplier = 3;
// The decay for the Adagrad accumulator.
float accumulator_decay = 5;
// The initial and minimum value for the Adagrad accumulator.
float initial_accumulator_value = 6;
// The tag for identifying the step counter used for the frequency aware
// Adagrad optimizer.
OptimizerDynamicInput step_counter = 7;
}
// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L629
message StochasticGradientDescentParameters {}
// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Ftrl
// https://static.googleusercontent.com/media/research.google.com/en//pubs/archive/41159.pdf
// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L2646
//
// The hyperparameters for FTRL are the same as for the Keras implementation,
// with some additions. The "beta" parameter matches the behavior described in
// the second link above; "beta" / (2 * learning rate) should be added to "l2"
// to get equivalent behavior in the other TensorFlow implementations of this
// optimizer. When the multiply_linear_by_lr field is set to true, a modified
// formula is used for FTRL that treats the "linear" accumulator as being
// pre-multiplied by the learning rate (i.e., the accumulator named "linear"
// actually stores "linear * learning_rate"). Other than checkpoint
// compatibility, this is mathematically equivalent for a static learning rate;
// for a dynamic learning rate, it is nearly the same as long as the learning
// rate does not change quickly. The benefit of setting multiply_linear_by_lr to
// true is that the modified formula handles zero and near-zero learning rates
// without producing NaNs, improving flexibility for learning rate ramp-up.
message FtrlParameters {
float l1 = 1;
float l2 = 2;
float lr_power = 3;
float beta = 7;
bool multiply_linear_by_lr = 6;
// Previously, allow_zero_accumulator parameter changed some internal formulas
// to allow zero and near-zero accumulator values at the cost of some
// performance. The current implementation ignores this parameter; zero or
// near-zero accumulator values are now always supported.
bool allow_zero_accumulator = 8 [deprecated = true];
// Old initial accumulator parameters.
reserved "initial_accum", "initial_linear";
reserved 4, 5;
}
// The Adam optimizer does not implement hyper-parameter update due to hardware
// limitations; use the dynamic learning rate feature instead, setting the
// learning rate to: user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
// Here, t is the current timestep.
//
// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adam
// https://github.com/tensorflow/tensorflow/blob/ab51450c817674c8ff08a7ae4f8ac50cdc4bed8b/tensorflow/python/training/adam.py#L32
//
// Note that the code by default implements the lazy version of Adam
// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/LazyAdamOptimizer)
// unless the use_non_lazy_adam parameter is set, in which case it implements
// the normal version of Adam that updates all parameters in the embedding
// table, even for entries that are not used in the current minibatch
// (https://www.tensorflow.org/api_docs/python/tf/contrib/opt/AdamOptimizer). If
// use_non_lazy_adam is enabled, gradient accumulation is also required to be
// enabled in order to get correct results; a warning will be printed otherwise
// (which may change to an error in the future). If use_sum_inside_sqrt is set,
// the Adam variable update formula will be changed from m / (sqrt(v) + epsilon)
// to m / sqrt(v + epsilon**2); this option improves the performance of TPU
// training and is not expected to harm model quality.
message AdamParameters {
float beta1 = 3;
float beta2 = 4;
float epsilon = 5;
bool use_non_lazy_adam = 8;
bool use_sum_inside_sqrt = 10;
// Old initial accumulator parameters.
reserved "initial_m", "initial_v";
reserved 6, 7;
}
// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/SGD
// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L3068
message MomentumParameters {
float momentum = 1;
bool use_nesterov = 2;
// Old initial accumulator parameter.
reserved "initial_accum";
reserved 3;
}
// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Lion
// momenta(new) = beta2 * momenta(old) + (1 - beta2) * grad
// momenta_t = beta1 * momenta(old) + (1 - beta1) * grad
// var(new) = var(old) - lr * sign(momenta_t)
// Algorithm described in https://arxiv.org/abs/2302.06675.
message LionParameters {
float beta1 = 1;
float beta2 = 2;
bool use_non_lazy_lion = 3;
}
// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4229
message RmsPropParameters {
float rho = 1;
float momentum = 2;
float epsilon = 3;
// Old initial accumulator parameters.
reserved "initial_ms", "initial_mom";
reserved 4, 5;
}
// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/RMSprop
// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L4358
message CenteredRmsPropParameters {
float rho = 1;
float momentum = 2;
float epsilon = 3;
// Old initial accumulator parameters.
reserved "initial_ms", "initial_mom", "initial_mg";
reserved 4, 5, 6;
}
// Variant of algorithm in http://proceedings.mlr.press/v44/shamir15.pdf
message MdlAdagradLightParameters {
float l2 = 1;
float lr_power = 2;
float min_servable_mdl_benefit = 3;
float mdl_mix_in_margin = 4;
float mdl_benefit_rampup_coeff = 5;
float mdl_min_weight = 6;
float benefit_revisit_scale = 7;
float max_event_benefit = 8;
float max_total_benefit = 9;
float mdl_hard_limit = 10;
bool hard_limit_min_benefit = 11;
bool mdl_regularize = 12;
// Old initial accumulator parameters.
reserved "initial_accumulator", "initial_weight", "initial_benefit";
reserved 13, 14, 15;
}
// https://www.tensorflow.org/api_docs/python/tf/keras/optimizers/Adadelta
// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L933
message AdadeltaParameters {
float rho = 1;
float epsilon = 2;
// Old initial accumulator parameters.
reserved "initial_accumulator", "initial_update";
reserved 3, 4;
}
// https://www.tensorflow.org/api_docs/python/tf/compat/v1/train/ProximalAdagradOptimizer
// https://github.com/tensorflow/tensorflow/blob/6b6471f3ffb7f1fefe42d814aa5fb9ab7a535b58/tensorflow/core/kernels/training_ops.cc#L1961
message ProximalAdagradParameters {
float l1 = 1;
float l2 = 2;
// Old initial accumulator parameter.
reserved "initial_accumulator";
reserved 3;
}
// The online Yogi optimizer does not implement hyper-parameter update; use the
// dynamic learning rate feature instead, setting the learning rate to:
// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
// Here, t is the current timestep.
//
// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
// plus some extensions based on FTRL.
//
// Note that the code by default implements the lazy version of online Yogi.
message OnlineYogiParameters {
// The L1 regularization parameter (used analogously to the one in FTRL).
float l1 = 1;
// The L2 regularization parameter (used analogously to the one in FTRL).
float l2 = 2;
// \beta_2 from Algorithm 2 in the paper.
float beta2 = 3;
// Reserved ids corresponding to removed tanh activation.
reserved 6; // sign
reserved 7; // tanh
}
// The online Yogi optimizer does not implement hyper-parameter update; use the
// dynamic learning rate feature instead, setting the learning rate to:
// user learning_rate * sqrt(1 - beta2^t) / (1 - beta1^t)
// Here, t is the current timestep.
//
// https://papers.nips.cc/paper/8186-adaptive-methods-for-nonconvex-optimization.pdf
// plus some extensions based on FTRL.
//
// Note that the code by default implements the lazy version of proximal Yogi.
message ProximalYogiParameters {
// The L1 regularization parameter.
float l1 = 1;
// The L2 regularization parameter.
float l2 = 2;
// The exponential decay rate for the 1st moment estimates.
float beta1 = 3;
// The exponential decay rate for the 2nd moment estimates.
float beta2 = 4;
// A constant trading off adaptivity and noise.
float epsilon = 5;
// Reserved ids corresponding to removed tanh activation.
reserved 8; // sign
reserved 9; // tanh
}
// Estimator for the frequency of updates to a lookup table. It maintains an
// array (tf.Variable) D, where each element records the average number of
// global steps between two consecutive batches that hit the corresponding
// bucket. Once an item with bucket id i is sampled, D[i] is updated by:
// D[i] <- D[i] * (1 - tau) + delta[i] * tau,
//
// where tau is a learning rate between 0 and 1 (exclusive), and
// delta[i] = current global step - last step i is sampled.
//
// The estimated frequency (sampling rate in a batch) is thus 1 / D[i].
//
// Elements in D are initialized with a large value max_delta. delta[i] will
// also be capped by this value.
//
// The exact sequence of operations used in the optimizer is shown below.
// last_hit_step[i] is a tf.Variable that holds the last global step at which i
// was sampled.
//
// delta = global_step - last_hit_step[i]
// clipped_delta = min(delta, params.max_delta)
// is_outlier = (delta >= params.outlier_threshold * D[i])
// D[i] <- is_outlier ? clipped_delta
// : D[i] * (1 - params.tau) + clipped_delta * params.tau
// last_hit_step[i] <- global_step
message FrequencyEstimatorParameters {
// Learning rate between (0, 1) that is used to update the array D.
float tau = 1;
// Maximum value of delta: difference between the current global step and the
// last global step at which the row was sampled.
float max_delta = 2;
// Threshold used to determine whether the current update is an outlier.
float outlier_threshold = 3;
// The weight exponent used to transform the estimated delta into weights.
// The transformation function is: (delta / max_delta) ^ (weight_exponent)
float weight_exponent = 4;
}
// A user-defined optimizer.
// The contained HLO program must take the following arguments in the following
// order:
// 1. gradients
// 2. table weights
// 3. slot variables
// 4. an optional scalar input that is passed in via the dynamic learning
// rate mechanism.
//
// It must return/end in a tuple op that contains the following values in the
// following order:
// 1. new table values
// 2. new slot variable value
//
// The program must have shape (1,1) with dtype float32 throughout and only use
// HLO that operate elementwise (e.g., no reduce, no variables, no control flow
// and no broadcasting outside of the single scalar input).
// The HLO program should be written as if it were a dense update. It will be
// called on each row that needs an update and will applied elementwise.
message UserDefinedProgramParameters {
xla.HloModuleProto program = 1;
reserved 2; // Was padding_values
}
// Optimizer that just sets the variable to the value of the gradient. To be
// correct, this requires either gradient accumulation (to sum the values of a
// computed expression across the samples) or to deduplicate IDs within a single
// host (to assign the value from an arbitrary sample).
message AssignParameters {}
// Status of using gradient accumulation (doing two passes over the input
// gradients: one to accumulate them into a temporary array and another to apply
// them using the actual optimization algorithm). The extra message is to wrap
// the enum for scoping.
message GradientAccumulationStatus {
// if UNSPECIFIED (default), gradient accumulation is ENABLED.
enum Status {
UNSPECIFIED = 0;
ENABLED = 1;
DISABLED = 2;
}
}
// Whether to optimize the packing of low-dimensional embedding tables in HBM
// (high bandwidth memory). TPUs access HBM at 32-byte (8-float) granularity.
// For functional correctness, the TPU software internally pads the embedding
// dimension to a multiple of 8. This can sometimes lead to significant memory
// wastage due to padding. For 1-dimensional, 2-dimensional, and 4-dimensional,
// the TPU software can remove this padding by packing multiple rows into the
// same 8-float HBM chunk. For example, 8 rows could be packed into the same
// 8-float chunk for a 1-dimensional embedding table.
// There is one important limitation for this HBM packing though. When only a
// subset of rows in an 8-float chunk are accessed on a particular step, the
// adjoining rows in the same chunk are updated with zero gradients on the
// backward pass even if they are not touched. This is an artifact of the
// packing implementation. This operation is NOT functionally correct for
// optimizers where zero gradients change the embeddings/slot-variable values,
// e.g., momentum-based optimizers. Hence, this HBM packing cannot be enabled
// for embedding tables with such optimizers. The TPU software automatically
// recognizes that a zero gradient can modify state and turns off the low
// dimensional embedding packing in that scenario.
//
// However, for optimizers where a zero gradient is a NoOp, such as SGD,
// Adagrad, and FTRL, this packing optimization can be used. However, there are
// some important considerations:
// * Clipping limits: The initial values for such embeddings should fall within
// the clipping limits specified in the optimization parameters. Otherwise, a
// zero gradient will cause the embeddings to be clipped. This changes state
// and hence, is not a NoOp.
// * FTRL: The embedding vector is computed directly from the values of the
// accumulator and linear slot variables. Hence, the initial embedding values
// should match that computed from the initial values of the accumulator and
// linear slot variables. Note that in nearly all cases, the linear value is
// initialized to zero; this corresponds to an embedding value of zero.
//
// Performance: The TPU has to perform additional work when low dimensional
// packing is enabled. In certain situations when the vocabulary size is small,
// it may not make sense to turn on this packing since the total memory usage
// due to padding is extremely low. Hence, the TPU software automatically turns
// off the packing optimization in such scenarios.
message LowDimensionalPackingStatus {
// if UNSPECIFIED (default), the low dimension packing status is DISABLED.
// This can change in future.
//
// if ENABLED, the low dimension packing is enabled only if the following
// three additional conditions are true:
// * The optimizer treats the zero gradient as a NoOp.
// * The embedding dimension is 1, 2, or 4.
// * The vocabulary size is large enough to avoid performance issues.
//
// if DISABLED, the low dimension packing is always disabled.
enum Status {
UNSPECIFIED = 0;
ENABLED = 1;
DISABLED = 2;
}
}
// Configuration proto for hot ID optimization. This is an experimental feature
// that is currently disabled (by default).
message HotIdReplicationConfiguration {
// Whether to enable or disable hot ID optimization.
// If set to UNSPECIFIED (default), hot ID optimization is DISABLED.
// If set to ENABLED, hot ID replication is turned ON.
// If set to MIGRATION_ONLY, hot ID migration is turned ON.
enum Status {
UNSPECIFIED = 0;
ENABLED = 1;
DISABLED = 2;
MIGRATION_ONLY = 3;
}
Status status = 1;
}
message OptimizationParameters {
// Learning rate used for updating the embedding layer parameters.
LearningRate learning_rate = 13;
// Limits to which to clip the weight values after the backward pass; not
// present means no limits are applied.
ClippingLimits clipping_limits = 2;
// Limits to which to clip the backward pass gradient before using it for
// updates; not present means no limits are applied.
ClippingLimits gradient_clipping_limits = 7;
// Amount of weight decay to apply; see weight_decay_optimizers.py for
// details. All optimizers except MDL Adagrad Light are supported with this
// option. Although there is no check, users who want weight decay will also
// want to ensure that gradient accumulation is enabled so that the decay will
// happen once per global batch.
float weight_decay_factor = 16;
// If true, the weight decay factor is multiplied by the current learning rate
// before use; this is to match the note in DecoupledWeightDecayExtension in
// weight_decay_optimizers.py.
bool multiply_weight_decay_factor_by_learning_rate = 22;
// Configuration for simulated quantization which is used to reduce
// training/serving skew when the serving variables are quantized. The same
// quantization operations are executed during training to minimize
// differences with serving.
SimulatedQuantization simulated_quantization = 27;
// Status of using gradient accumulation (doing two passes over the input
// gradients: one to accumulate them into a temporary array and another to
// apply them using the actual optimization algorithm).
GradientAccumulationStatus.Status gradient_accumulation_status = 17;
// Status of the low-dimensional embedding packing optimization. This controls
// whether to optimize the packing of 1-dimensional, 2-dimensional, and
// 4-dimensional embedding tables in memory.
LowDimensionalPackingStatus.Status low_dimensional_packing_status = 28;
// Configuration proto for hot ID replication. This is an experimental
// feature that is currently disabled (by default).
HotIdReplicationConfiguration hot_id_replication_configuration = 18;
// Optimization algorithm parameters; which field is selected determines which
// algorithm to use.
oneof parameters {
AdagradParameters adagrad = 3;
AdagradMomentumParameters adagrad_momentum = 26;
BoundedAdagradParameters bounded_adagrad = 19;
FrequencyAwareAdagradParameters frequency_aware_adagrad = 30;
StochasticGradientDescentParameters stochastic_gradient_descent = 4;
FtrlParameters ftrl = 5;
AdamParameters adam = 6;
MomentumParameters momentum = 8;
LionParameters lion = 29;
RmsPropParameters rms_prop = 9;
CenteredRmsPropParameters centered_rms_prop = 10;
MdlAdagradLightParameters mdl_adagrad_light = 11;
AdadeltaParameters adadelta = 12;
ProximalAdagradParameters proximal_adagrad = 14;
OnlineYogiParameters online_yogi = 20;
ProximalYogiParameters proximal_yogi = 21;
FrequencyEstimatorParameters frequency_estimator = 23;
UserDefinedProgramParameters user_defined_program = 24;
AssignParameters assign = 25;
}
reserved 1, 15;
// NEXT_ID: 31
}
// Specification of an optimization algorithm's state variables (both the main
// value vector and any extra accumulators, etc.). This proto is only used
// internally by the TPU software and is not exposed directly to the TF model.
message StateVariableSpecification {
// Parameter name for the state variable.
string name = 1;
// A normal state variable that should be saved and restored in checkpoints
// and used as an input or output to non-debug TensorFlow ops.
message UserDefined {
reserved 1; // Was padding_initial_value.
}
// A state variable that should be filled with a constant and normally hidden
// from users (used for intermediate gradients being accumulated, for
// example).
message FillWithConstant {
double initial_value = 1;
}
// Usage type of this state variable.
oneof usage {
UserDefined user_defined = 2;
FillWithConstant fill_with_constant = 3;
}
}
tensorflow/core/protobuf/tpu/topology.proto
+66
View File
@@ -0,0 +1,66 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.tpu;
option cc_enable_arenas = true;
// Describes features of a tpu.
message TPUHardwareFeature {
// Embedding feature of a tpu.
enum EmbeddingFeature {
// No embedding lookup accelerator available on the tpu.
UNSUPPORTED = 0;
// Embedding lookup accelerator V1. The embedding lookup operation can only
// be placed at the beginning of computation. Only one instance of embedding
// lookup layer is allowed.
V1 = 1;
// Embedding lookup accelerator V2. The embedding lookup operation can be
// placed anywhere of the computation. Multiple instances of embedding
// lookup layer is allowed.
V2 = 2;
}
EmbeddingFeature embedding_feature = 1;
// Number of embedding accelerator devices per chip.
int32 num_embedding_devices_per_chip = 2;
}
// Describes the geometry of a TPU mesh.
message TopologyProto {
// The dimensions of the TPU topology, in cores. Typically, this is a 4D
// topology [x, y, z, core], where the major dimensions correspond to TPU
// chips, and the minor dimension describes the number of cores on a multicore
// chip.
repeated int32 mesh_shape = 1;
// Number of TensorFlow tasks in the cluster.
int32 num_tasks = 2;
// Number of TPU devices per task.
int32 num_tpu_devices_per_task = 3;
// A flattened rank 3 int32 array with shape
// [num_tasks, num_tpu_devices_per_task, len(mesh_shape)].
// `tasks` is the number of tasks in the TPU cluster, `devices` is the number
// of TPU devices per task, and the minor dimension corresponds to a position
// in the TPU mesh topology. Each entry [task, device, axis] gives the
// `axis`-th coordinate in the topology of a task/device pair.
repeated int32 device_coordinates = 4;
// TPU supported features.
TPUHardwareFeature tpu_hardware_feature = 5;
}
tensorflow/core/protobuf/tpu/tpu_embedding_configuration.proto
+176
View File
@@ -0,0 +1,176 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow.tpu;
import "tensorflow/core/protobuf/tpu/optimization_parameters.proto";
message TPUEmbeddingConfiguration {
// Description of the various embedding tables.
message TableDescriptor {
// Name of the table.
string name = 1;
// Size of the vocabulary (i.e., number of rows) in the table.
int64 vocabulary_size = 2;
// The embedding dimension (i.e., the width of the embedding table).
int32 dimension = 3;
// Number of features mapped to this table.
int32 num_features = 4;
// Details of the learning algorithm used to update the embedding
// parameters.
OptimizationParameters optimization_parameters = 5;
}
repeated TableDescriptor table_descriptor = 1;
// Mode. Should the embedding layer program be run for inference (just forward
// pass), training (both forward and backward pass) or just the backward_pass.
enum Mode {
UNSPECIFIED = 0;
INFERENCE = 1;
TRAINING = 2;
BACKWARD_PASS_ONLY = 3;
}
Mode mode = 2;
// Number of samples in each batch of embedding layer activations sent to
// the TensorCore.
int32 batch_size_per_tensor_core = 3;
// Number of TPU hosts used for inference/training.
int32 num_hosts = 4;
// Number of TensorCore used for inference/training.
int32 num_tensor_cores = 5;
// Sharding strategy of the embedding tables among the hosts.
// If the sharding_strategy is "mod", each id is assigned to host
// "id % num_hosts". For instance, 13 ids are split across 5 hosts as:
// [[0, 5, 10], [1, 6, 11], [2, 7, 12], [3, 8], [4, 9]].
// If the sharding_strategy is "div", ids are assigned to hosts in a
// contiguous manner. In this case, 13 ids are split across 5 hosts as:
// [[0, 1, 2], [3, 4, 5], [6, 7, 8], [9, 10], [11, 12]].
// In both the strategies, if the id space does not evenly divide the number
// of hosts, each of the first "table_descriptor.vocabulary_size % num_hosts"
// hosts will be assigned one more id.
// This partitioning strategy exactly follows that in the embedding_lookup
// TensorFlow function at tensorflow/python/ops/embedding_ops.py.
enum ShardingStrategy {
DIV_DEFAULT = 0;
MOD = 1;
}
ShardingStrategy sharding_strategy = 6;
// This parameter determines if the execution of the sparse core will be
// pipelined with that of the TensorCore. This parameter only affects results
// when mode=TRAINING. If mode=INFERENCE or BACKWARD_PASS_ONLY, this parameter
// does not affect execution and hence, is a don't care value.
//
// false: The execution of the sparse core is not pipelined with that of the
// TensorCore. The forward pass of every step on the sparse core is executed
// only after the backward pass of the previous step is complete. And the
// backward pass on the sparse core is executed only after the embedding
// gradients have been computed on the TensorCore on every step. This ensures
// that the activations on every step observe the gradient updates from the
// previous step on both the sparse core and the TensorCore.
//
// true: The execution of the sparse core is pipelined with that of the
// TensorCore. The forward pass of every step on the sparse core can be
// executed after the forward pass of the previous step is complete without
// waiting for the backward pass. This improves the utilization of the sparse
// core allowing it to process step N+1 while the embedding gradients for step
// N are computed on the TensorCore. The backward pass of every step on the
// sparse core is executed directly after the forward pass for the next step
// is complete. The drawback is that embedding activations for step N+1 do not
// observe the embedding gradient updates from step N. This could affect model
// quality if step N and N+1 involve the same set of embedding IDs. However,
// since the embedding updates are sparse, this is generally not considered a
// problem.
bool pipeline_execution_with_tensor_core = 7;
// Directory where embedding lookup statistics are stored. These statistics
// summarize information about the inputs to the embedding lookup
// operation, in particular, the average number of embedding IDs per example
// and how well the embedding IDs are load balanced across the system. The
// lookup statistics are used during TPU initialization for embedding table
// partitioning. Collection of lookup statistics is done at runtime by
// profiling the embedding inputs: only 3% of input samples are profiled to
// minimize host CPU overhead. Once a suitable number of samples are
// profiled, the lookup statistics are saved to table-specific files in the
// profile data directory generally at the end of a TPU training loop. The
// filename corresponding to each table is obtained by hashing table specific
// parameters (e.g., table name and number of features) and global
// configuration parameters (e.g., sharding strategy and TPU worker task
// count). The same profile data directory can be shared amongst several
// models to reuse embedding lookup statistics.
string profile_data_directory = 9;
// Description of different input features.
message FeatureDescriptor {
// Name of the input feature.
string name = 1;
// Index of the corresponding table in the TableDescriptor list.
int32 table_id = 2;
// Static shape of the inputs (excluding the reduction axis). Note that
// the shape of the actual inputs provided using the infeed op must be
// strictly smaller than input_shape. The outputs received at the TensorCore
// will have rank = input_shape.size() + 1. The innermost axis corresponds
// to the embedding dimension. If the input has shape [m, n, k] (excluding
// the reduction axis) and the embedding dimension is d, the output received
// at the TensorCore will have shape [m, n, k, d].
repeated int32 input_shape = 3;
}
// If the feature_descriptor field is populated, the model should NOT populate
// TableDescriptor.num_features and batch_size_per_tensor_core. These two
// fields will be auto-populated by the TPUEmbedding rewrite passes.
repeated FeatureDescriptor feature_descriptor = 10;
// SPMD (Single Program Multiple Data) sharding configuration for
// TPUEmbedding. When model parallelism is used on the TensorCore, the number
// of cores per replica must be passed to TPUEmbedding so that the right
// shapes can be computed in the TF/XLA bridge.
message SpmdSharding {
// Whether SPMD sharding is enabled.
bool enabled = 1;
// Number of cores per replica.
int32 num_cores_per_replica = 2;
// If true, the tensors are manually partitioned. Otherwise, use the
// automatic SPMD partitioning. This should be true when users use
// `shard_map`.
bool use_manual_partitioning = 3;
// If true, lowers to Shardy custom-call for manual partitioning.
// Default is false.
bool use_shardy_partitioner = 4;
}
SpmdSharding spmd_sharding = 11;
// Old TPU embedding output layout.
reserved "output_layout";
reserved 8;
}
// A placeholder message that is used to define a unique Status payload
// URL for TPU embedding errors.
message TPUEmbeddingError {}
tensorflow/core/protobuf/trackable_object_graph.proto
+95
View File
@@ -0,0 +1,95 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
import "google/protobuf/wrappers.proto";
option cc_enable_arenas = true;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// A TensorBundle addition which saves extra information about the objects which
// own variables, allowing for more robust checkpoint loading into modified
// programs.
message TrackableObjectGraph {
message TrackableObject {
message ObjectReference {
// An index into `TrackableObjectGraph.nodes`, indicating the object
// being referenced.
int32 node_id = 1;
// A user-provided name for the edge.
string local_name = 2;
}
message SerializedTensor {
// A name for the Tensor. Simple variables have only one
// `SerializedTensor` named "VARIABLE_VALUE" by convention. This value may
// be restored on object creation as an optimization.
string name = 1;
// The full name of the variable/tensor, if applicable. Used to allow
// name-based loading of checkpoints which were saved using an
// object-based API. Should match the checkpoint key which would have been
// assigned by tf.train.Saver.
string full_name = 2;
// The generated name of the Tensor in the checkpoint.
string checkpoint_key = 3;
// Deprecated bool field for optional restore. This field has never been
// set to True.
reserved "optional_restore";
reserved 4;
}
message SlotVariableReference {
// An index into `TrackableObjectGraph.nodes`, indicating the
// variable object this slot was created for.
int32 original_variable_node_id = 1;
// The name of the slot (e.g. "m"/"v").
string slot_name = 2;
// An index into `TrackableObjectGraph.nodes`, indicating the
// `Object` with the value of the slot variable.
int32 slot_variable_node_id = 3;
}
// Objects which this object depends on.
repeated ObjectReference children = 1;
// Serialized data specific to this object.
repeated SerializedTensor attributes = 2;
// Slot variables owned by this object.
repeated SlotVariableReference slot_variables = 3;
// The registered saver used to save this object. If this saver is not
// present when loading the checkpoint, then loading will fail.
RegisteredSaver registered_saver = 4;
// Whether this object has checkpoint values or descendants with checkpoint
// values. This is computed at save time to avoid traversing the entire
// object graph proto when restoring (which also has to traverse the live
// object graph).
google.protobuf.BoolValue has_checkpoint_values = 5;
}
repeated TrackableObject nodes = 1;
}
message RegisteredSaver {
// The name of the registered saver/restore function.
string name = 1;
// Unique auto-generated name of the object.
string object_name = 2;
}
tensorflow/core/protobuf/transport_options.proto
+25
View File
@@ -0,0 +1,25 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
// Extra data needed on a non-RDMA RecvBufResponse.
message RecvBufRespExtra {
repeated bytes tensor_content = 1;
}
tensorflow/core/protobuf/verifier_config.proto
+42
View File
@@ -0,0 +1,42 @@
// Copyright 2026 The TensorFlow Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
// ==============================================================================
syntax = "proto3";
package tensorflow;
option cc_enable_arenas = true;
option java_outer_classname = "VerifierConfigProtos";
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";
// The config for graph verifiers.
message VerifierConfig {
enum Toggle {
DEFAULT = 0;
ON = 1;
OFF = 2;
}
// Deadline for completion of all verification i.e. all the Toggle ON
// verifiers must complete execution within this time.
int64 verification_timeout_in_ms = 1;
// Perform structural validation on a tensorflow graph. Default is OFF.
Toggle structure_verifier = 2;
// Next tag: 3
}
tensorflow/core/protobuf/worker.proto
+610
View File
@@ -0,0 +1,610 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
syntax = "proto3";
package tensorflow;
import "google/protobuf/any.proto";
import "tensorflow/core/framework/cost_graph.proto";
import "tensorflow/core/framework/device_attributes.proto";
import "tensorflow/core/framework/graph.proto";
import "tensorflow/core/framework/step_stats.proto";
import "tensorflow/core/framework/tensor.proto";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";
import "tensorflow/core/protobuf/config.proto";
import "tensorflow/core/protobuf/debug.proto";
import "tensorflow/core/protobuf/error_codes.proto";
import "tensorflow/core/protobuf/named_tensor.proto";
import "tensorflow/core/protobuf/tensorflow_server.proto";
option cc_enable_arenas = true;
option java_outer_classname = "WorkerProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.distruntime";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
////////////////////////////////////////////////////////////////////////////////
//
// GetStatus method request/response messages
//
////////////////////////////////////////////////////////////////////////////////
message GetStatusRequest {}
message GetStatusResponse {
repeated DeviceAttributes device_attributes = 1;
}
////////////////////////////////////////////////////////////////////////////////
//
// CreateSession method request/response messages
//
// For each session,
//
////////////////////////////////////////////////////////////////////////////////
message CreateWorkerSessionRequest {
// Sessions are identified by a given handle.
string session_handle = 1;
// Defines the configuration of a TensorFlow worker.
ServerDef server_def = 2;
// If true, any resources such as Variables used in the session will not be
// shared with other sessions.
bool isolate_session_state = 3;
// The device attributes of all the devices in the cluster.
repeated DeviceAttributes cluster_device_attributes = 4;
// The master task name from which the request is sent.
string master_task = 5;
// The incarnation ID of the master task local CPU device.
// If the target worker already has a WorkerSession created previously with
// the same master task name but a different incarnation, it usually indicates
// that the previous master failed before deleting the WorkerSession on the
// worker. To prevent memory leaks, the worker should garbage collect the old
// WorkerSessions.
int64 master_incarnation = 6;
reserved 7; // Deprecated config that is embedded within server_def now.
}
message CreateWorkerSessionResponse {}
////////////////////////////////////////////////////////////////////////////////
//
// DeleteSession method request/response messages
//
// Deletes all worker-side state associated with the given session handle.
//
////////////////////////////////////////////////////////////////////////////////
message DeleteWorkerSessionRequest {
// Sessions are identified by a given handle.
string session_handle = 1;
}
message DeleteWorkerSessionResponse {}
////////////////////////////////////////////////////////////////////////////////
//
// RegisterGraph method request/response messages
//
// For each session, after the master placed every node on a device,
// it partitions the whole graph into many subgraphs. All the nodes in
// a subgraph were in the same worker, but potentially on many devices
// owned by that worker (e.g. cpu0, plus gpu0, gpu1, ..., gpu7). The
// master registers subgraphs for a worker before running any steps. A
// successful registration returns a graph handle to be used in latter
// RunGraph requests.
//
////////////////////////////////////////////////////////////////////////////////
message RegisterGraphRequest {
// Subgraphs are scoped within one session.
string session_handle = 1;
// Set to true if `CreateWorkerSession` was called for `session_handle`.
bool create_worker_session_called = 6;
// "graph_def" has the subgraph of nodes for this worker, with each node
// having its device_name filled in.
GraphDef graph_def = 2;
// True iff the graph (before partitioning) contains control flow nodes.
//
// As of 01/11/2015, this is no longer set by clients.
bool has_control_flow = 3 [deprecated = true];
// Configuration options for the session in which this graph was created.
GraphOptions graph_options = 4;
// Field(s) used by TensorFlow Debugger (tfdbg).
DebugOptions debug_options = 5;
// If graph_def contains any collective ops this must be a positive
// integer used to coordinate execution with other graphs. All
// graphs in a distributed execution with the same
// collective_graph_key will coordinate to use the same step_id
// concurrently so that BufRendezvous entries will make the correct
// values accessible.
int64 collective_graph_key = 7;
// ConfigProto from the session in which this graph was created.
// Contains additional parameters beyond graph_options, including
// the name of the requested executor.
ConfigProto config_proto = 8;
}
message RegisterGraphResponse {
// If the registration succeeds, returns an opaque graph_handle to
// the master. The master calls RunGraph with graph_handle to
// compute different steps.
string graph_handle = 1;
}
////////////////////////////////////////////////////////////////////////////////
//
// DeregisterGraph method request/response messages
//
// The master deregisters the given graph_handle when the graph is no
// longer needed (e.g., the overall graph is re-scheduled and nodes
// are re-placed).
//
// The worker deregisters a graph_handle automatically according to on
// a TTL-base policy in case of master restarts.
//
////////////////////////////////////////////////////////////////////////////////
message DeregisterGraphRequest {
// The session_handle used when registering the graph. If session_handle is
// empty, a single global namespace is used.
string session_handle = 2;
// Set to true if `CreateWorkerSession` was called for `session_handle`.
bool create_worker_session_called = 3;
// REQUIRED: graph_handle must be returned by a RegisterGraph call
// to the same WorkerService.
string graph_handle = 1;
}
message DeregisterGraphResponse {
// TODO(mrry): Optionally add summary stats for the graph.
}
////////////////////////////////////////////////////////////////////////////////
//
// CleanupAll method request/response messages
//
////////////////////////////////////////////////////////////////////////////////
message CleanupAllRequest {
// A list of container names.
//
// If 'container' is not empty, releases resources in the given
// containers in all devices.
//
// If 'container' is empty, releases resources in the default
// container in all devices.
repeated string container = 1;
}
message CleanupAllResponse {}
////////////////////////////////////////////////////////////////////////////////
//
// RunGraph request / response messages
//
// The worker executes all subgraphs registered under graph_handle.
// RunGraph returns after the execution finishes or an error is
// encountered.
// A sequence of RunGraphRequests with is_partial may be sent to RunGraph for
// partial graph execution.
//
////////////////////////////////////////////////////////////////////////////////
// Options specific to the execution of a single step.
message ExecutorOpts {
bool record_costs = 1;
bool record_timeline = 3;
bool record_partition_graphs = 4;
bool report_tensor_allocations_upon_oom = 5;
}
message RunGraphRequest {
// session_handle is the master-generated unique id for this session.
// If session_handle is non-empty, it must be the same as used when
// registering the graph. If it is empty, a single global namespace is used to
// search for the graph_handle.
string session_handle = 8;
// Set to true if `CreateWorkerSession` was called for `session_handle`.
bool create_worker_session_called = 10;
// REQUIRED: graph_handle must be returned by a RegisterGraph call
// to the same WorkerService.
string graph_handle = 1;
// A unique ID to distinguish different runs of the same graph.
//
// The master generates a global unique `step_id` to distinguish
// different runs of the graph computation. Subgraphs communicate
// (e.g., send/recv ops) with each other using `step_id` to
// distinguish tensors generated by different runs.
int64 step_id = 2;
// Options for this step.
ExecutorOpts exec_opts = 5;
// Runs the graph.
//
// Sends the tensors in "send" into the graph before the run and
// fetches the keys into `RunGraphResponse.recv` after the run.
repeated NamedTensorProto send = 3;
repeated string recv_key = 4;
// True if the RunGraphRequest is a partial run request.
bool is_partial = 6;
// True if this is the last partial run request in a sequence of requests.
bool is_last_partial_run = 7;
// If true then some errors, e.g., execution errors that have long
// error messages, may return an OK RunGraphResponse with the actual
// error saved in the status_code/status_error_message fields of the
// response body. This is a workaround since the RPC subsystem may
// truncate long metadata messages.
bool store_errors_in_response_body = 9;
// Unique identifier for this request. Every RunGraphRequest must have a
// unique request_id, and retried RunGraphRequests must have the same
// request_id. If request_id is zero, retry detection is disabled.
//
// Retried RunGraphRequests are problematic because they may issue a
// RecvTensor that will have no corresponding sender and will wait forever.
// Workers use request_ids to reject retried RunGraph requests instead of
// waiting forever.
int64 request_id = 11;
// Next: 12
}
message RunGraphResponse {
// A list of tensors corresponding to those requested by
// `RunGraphRequest.recv_key`.
repeated NamedTensorProto recv = 1;
// If the request asked for execution stats, the cost graph, or the partition
// graphs, these are returned here.
// TODO(suharshs): Package these in a RunMetadata instead.
StepStats step_stats = 2;
CostGraphDef cost_graph = 3;
repeated GraphDef partition_graph = 4;
// If store_errors_in_response_body is true in the request, then
// optionally the server may return an OK status for the RPC and
// fill the true status into the fields below, to allow for messages
// that are too long to fit in metadata.
error.Code status_code = 5;
string status_error_message = 6;
}
////////////////////////////////////////////////////////////////////////////////
//
// CleanupGraph method request/response messages
//
// After the master receives RunGraph responses from all workers, the
// master instructs every worker to cleanup any remaining state of a
// step (e.g. tensors buffered by a `Send` op but not picked up by
// other workers). The master does not necessarily need to wait for
// completion of CleanupGraph calls.
//
// Workers should cleanup step states automatically according to a
// TTL-based policy in case of master restarts.
//
////////////////////////////////////////////////////////////////////////////////
message CleanupGraphRequest {
int64 step_id = 1;
}
message CleanupGraphResponse {}
////////////////////////////////////////////////////////////////////////////////
//
// RecvTensor method request/response messages
//
////////////////////////////////////////////////////////////////////////////////
message RecvTensorRequest {
// The step in which the tensor will be produced.
//
// REQUIRED: This must eventually correspond to the `step_id` passed
// into a RunGraph call on the same WorkerService.
int64 step_id = 1;
// A key identifying the channel to receive tensors from. A RecvTensor request
// retrieves one tensor from the channel, but multiple tensors can be sent and
// received over the same channel with multiple RecvTensor requests. See
// rendezvous.h for details.
string rendezvous_key = 2;
// If true, use an out-of-band DMA mechanism to transfer the
// received tensor.
bool dma_ok = 3;
// Optional information on client-side device locality.
DeviceLocality client_locality = 4;
// Optional information on server-side device locality.
DeviceLocality server_locality = 5;
// Optional information needed by the RPC subsystem.
google.protobuf.Any transport_options = 6;
// Unique identifier for this request. Every RecvTensorRequest must have a
// unique request_id, and retried RecvTensorRequests must have the same
// request_id. If request_id is zero, retry detection and response cache
// are disabled.
//
// Retried RecvTensorRequests are problematic because a RecvTensor with no
// corresponding sender will wait forever, and the tensor may have been
// delivered to a previous retry. Workers use request_ids to reject retried
// RecvTensor requests instead of waiting forever.
int64 request_id = 7;
}
message RecvTensorResponse {
// The tensor as a proto.
TensorProto tensor = 1;
// If true, this tensor was the output of a dead node, and the
// content is invalid.
bool is_dead = 2;
// The time at which tensor was available and started to be returned.
int64 send_start_micros = 3;
// Optional additional information about how to receive the tensor,
// e.g. in the event that `RecvTensorRequest.dma_ok` was true.
google.protobuf.Any transport_options = 4;
// Whether the receiver should send a MarkRecvFinishedRequest to the sender
// to ack the message.
bool require_ack = 5;
}
// Message for managing the response cache maintained on the sender side.
// Currently only used by the gRPC worker service.
message MarkRecvFinishedRequest {
int64 request_id = 1;
}
message MarkRecvFinishedResponse {}
////////////////////////////////////////////////////////////////////////////////
//
// Logging method request/response messages
//
// NOTE(mrry): This feature is not supported in the open-source
// version, and these messages are expected to change.
//
////////////////////////////////////////////////////////////////////////////////
// Out-of-band request to begin or end logging, or
// to retrieve logs for particular steps.
message LoggingRequest {
// If true, RPC logging will be enabled.
bool enable_rpc_logging = 1;
// If true, RPC logging will be disabled.
bool disable_rpc_logging = 4;
// If true, discard any saved logging data (for all steps).
bool clear = 2;
// When set, requests all saved log data pertaining to the step.
// Any log data retrieved is eliminated from the store and cannot be
// retrieved again.
repeated int64 fetch_step_id = 3;
}
message LabeledStepStats {
int64 step_id = 1;
StepStats step_stats = 2;
}
message LoggingResponse {
repeated LabeledStepStats step = 1;
}
////////////////////////////////////////////////////////////////////////////////
//
// Tracing method request/response messages
//
// NOTE(mrry): This feature is not supported in the open-source
// version, and these messages are expected to change.
//
////////////////////////////////////////////////////////////////////////////////
message TraceOpts {
// Length of the trace to be taken, in seconds.
double duration = 1;
// If true, capture step profile locally in each worker. Currently
// unimplemented.
bool use_step_profiler = 2;
// If true, capture kernel events from each worker.
bool use_kernel_profiler = 3;
// If true, capture extended profiling events from TensorFlow process.
bool use_extended_profiler = 4;
// If true, capture GPU profiling events locally on each
// machine. Currently unimplemented.
bool use_gpu_profiler = 5;
// If true, collect sampled profile events. Currently unimplemented.
bool use_sample_profiler = 6;
}
// Out-of-band request to configure distributed tracing.
message TracingRequest {
TraceOpts options = 1;
}
message TracingResponse {}
////////////////////////////////////////////////////////////////////////////////
//
// Raw data transfers in support of Collective Ops.
// These methods are experimental and subject to change.
//
// The intention is to allow collectives to take advantage of the most
// efficient methods available on a platform, e.g. RDMA, and not be
// constrained to use the RPC system in use by other methods.
//
////////////////////////////////////////////////////////////////////////////////
message RecvBufRequest {
// Use of the fields below may vary by implementation. For example
// the buf_ptr and num_bytes may be set only for local operations and
// not sent on the wire, or only sent on the wire in one direction.
// Used at server side to find the correct BufRendezvous.
int64 step_id = 1;
// Arbitrary string identifying a BufRendezvous entry.
string buf_rendezvous_key = 2;
// Size of value expected, must agree with BufRendezvous entry.
int64 num_bytes = 3;
// When RDMA is in use, address of destination field on client.
fixed64 buf_ptr = 4;
// Optional information on client-side device locality.
DeviceLocality client_locality = 5;
// Optional information on server-side device locality.
DeviceLocality server_locality = 6;
// Optional, implementation-specific data.
google.protobuf.Any transport_options = 7;
// For annotating timeline and device incarnation check.
string src_device = 8;
// Optional, for annotating the timeline.
string dst_device = 9;
// Depending on the RPC system in use, it may be necessary to set this
// id to detect resends of RPCs where the server is not aware that
// the prior RPC failed.
int64 request_id = 10;
// Incarnation number of the source device, used to detect worker failures.
uint64 src_incarnation = 11;
}
message RecvBufResponse {
// Use of the fields below may vary by implementation. Comments give
// intended use.
fixed64 buf_ptr = 1; // Address of source field on server.
int64 num_bytes = 2; // Byte length of buf_ptr field, if set.
bool is_dead = 3; // True if value is 'dead' like a tensor.
// Optional, implementation-specific data.
google.protobuf.Any transport_options = 4;
// Optional, for timeline.
int64 send_start_micros = 5;
// Whether the receiver should send a MarkRecvFinishedRequest to the sender
// to ack the message.
bool require_ack = 6;
}
////////////////////////////////////////////////////////////////////////////////
//
// Collective Op dynamic group resolution messages.
//
////////////////////////////////////////////////////////////////////////////////
// Supplies one or more device names as members of the group identified by
// group_key. Service will respond when all group_size devices become known.
// All devices in group must have same type.
message CompleteGroupRequest {
int32 group_key = 1;
int32 group_size = 2;
string device_type = 3;
int32 collective_type = 5;
DeviceAttributes device_attributes = 6;
reserved 4;
}
// Gives the complete membership of the group identified by group_key.
message CompleteGroupResponse {
int32 group_key = 1;
int32 group_size = 2;
string device_type = 3;
int32 num_tasks = 4; // number of distinct tasks hosting the devices
bytes communicator_key = 7;
repeated DeviceAttributes device_attributes = 8;
reserved 5, 6;
}
// Supplies data about one collective op belonging to the instance identified
// by instance_key and step_id. Service will respond when all group_size ops
// have become known. Most of the data being sent is for correctness checking,
// to ensure that all ops in the instance share common attributes.
message CompleteInstanceRequest {
string name = 1;
int32 type = 2;
DataType data_type = 3;
TensorShapeProto shape = 4;
int32 group_key = 5;
int32 group_size = 6;
int32 instance_key = 7;
string device_type = 8;
repeated int32 subdiv_offset = 9;
string device = 10;
bool is_source = 11;
int64 step_id = 12;
}
// Confirms that every op in the instance has consistently declared itself.
// Also gives the source_rank in case of broadcast.
message CompleteInstanceResponse {
int32 instance_key = 1;
int32 source_rank = 2;
reserved 3;
}
// Request for next agreed-upon step_id for the specified graph_keys.
// This is used to enable multiple graphs containing nodes from
// a common collective instance to coordinate using the same step_ids.
message GetStepSequenceRequest {
repeated int64 graph_key = 1;
}
message StepSequence {
int64 graph_key = 1;
int64 next_step_id = 2;
}
// Next valid step_ids for one or more graph_keys.
message GetStepSequenceResponse {
repeated StepSequence step_sequence = 1;
}
tensorflow/core/protobuf/worker_service.proto
+126
View File
@@ -0,0 +1,126 @@
/* Copyright 2016 The TensorFlow Authors. All Rights Reserved.
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/
syntax = "proto3";
package tensorflow.grpc;
import "tensorflow/core/protobuf/worker.proto";
option java_outer_classname = "WorkerServiceProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.distruntime";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf/for_core_protos_go_proto";
////////////////////////////////////////////////////////////////////////////////
//
// WorkerService defines a TensorFlow service that executes dataflow
// graphs on a set of local devices, on behalf of a MasterService.
//
// A worker service keeps track of multiple "registered graphs". Each
// registered graph is a subgraph of a client's graph, corresponding to
// only the nodes that should execute on this worker (and any
// additional nodes necessary for inter-process communication using
// the `RecvTensor` method).
//
////////////////////////////////////////////////////////////////////////////////
service WorkerService {
// See worker.proto for details.
rpc GetStatus(GetStatusRequest) returns (GetStatusResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc CreateWorkerSession(CreateWorkerSessionRequest)
returns (CreateWorkerSessionResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc DeleteWorkerSession(DeleteWorkerSessionRequest)
returns (DeleteWorkerSessionResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc RegisterGraph(RegisterGraphRequest) returns (RegisterGraphResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc DeregisterGraph(DeregisterGraphRequest)
returns (DeregisterGraphResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc RunGraph(RunGraphRequest) returns (RunGraphResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc CleanupGraph(CleanupGraphRequest) returns (CleanupGraphResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc CleanupAll(CleanupAllRequest) returns (CleanupAllResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc RecvTensor(RecvTensorRequest) returns (RecvTensorResponse) {
// RecvTensor Method
}
// See worker.proto for details.
rpc MarkRecvFinished(MarkRecvFinishedRequest)
returns (MarkRecvFinishedResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc Logging(LoggingRequest) returns (LoggingResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc Tracing(TracingRequest) returns (TracingResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc RecvBuf(RecvBufRequest) returns (RecvBufResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc GetStepSequence(GetStepSequenceRequest)
returns (GetStepSequenceResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc CompleteGroup(CompleteGroupRequest) returns (CompleteGroupResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
// See worker.proto for details.
rpc CompleteInstance(CompleteInstanceRequest)
returns (CompleteInstanceResponse) {
// [AUTOMATION]: Internal rpc option goes here.
}
}