Fix go proto handling

This commit is contained in:
Jonathan Hseu 2020-01-31 12:47:04 -08:00
parent 20f9bb5ba4
commit 25251f057c
19 changed files with 58 additions and 25 deletions

View File

@ -10,6 +10,8 @@ package tensorflow;
import "google/protobuf/any.proto";
import "google/protobuf/duration.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
message CudnnVersion {
int32 major = 1;
int32 minor = 2;

View File

@ -2,6 +2,8 @@ syntax = "proto3";
package tensorflow;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
// Some of the data from AllocatorStats
message MemAllocatorStats {
int64 num_allocs = 1;

View File

@ -6,6 +6,8 @@ package tensorflow;
import "tensorflow/stream_executor/dnn.proto";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
// 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 {

View File

@ -10,6 +10,8 @@ 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";
// 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.

View File

@ -22,6 +22,8 @@ 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";
// This file contains protos to be used when defining a TensorFlow
// cluster.
//

View File

@ -11,6 +11,8 @@ 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";
// A proto representation of an eager operation.
message Operation {
// A unique identifier for the operation. Set by the client so that the client

View File

@ -7,6 +7,8 @@ option java_outer_classname = "GraphDebugInfoProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.framework";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
message GraphDebugInfo {
// This represents a file/line location in the source code.
message FileLineCol {

View File

@ -10,6 +10,8 @@ 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";
message ResourceDtypeAndShape {
DataType dtype = 1;
TensorShapeProto shape = 2;

View File

@ -6,6 +6,8 @@ import "tensorflow/core/protobuf/master.proto";
option cc_enable_arenas = true;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
// Records the creation of a new replay session. We record the device listing
// here to capture the state of the cluster.
message NewReplaySession {

View File

@ -11,6 +11,8 @@ option cc_enable_arenas = true;
package tensorflow;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
// 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.

View File

@ -2,6 +2,8 @@ syntax = "proto3";
package tensorflow;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
import "tensorflow/core/framework/tensor_shape.proto";
import "tensorflow/core/framework/types.proto";

View File

@ -7,6 +7,8 @@ option java_outer_classname = "TraceEventsProtos";
option java_multiple_files = true;
option java_package = "org.tensorflow.framework";
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
// A 'Trace' contains metadata for the individual traces of a system.
message Trace {
// The devices that this trace has information about. Maps from device_id to

View File

@ -4,6 +4,8 @@ option cc_enable_arenas = true;
package tensorflow;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
// A TensorBundle addition which saves extra information about the objects which
// own variables, allowing for more robust checkpoint loading into modified
// programs.

View File

@ -2,6 +2,8 @@ syntax = "proto3";
package tensorflow;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf";
// Extra data needed on a non-RDMA RecvBufResponse.
message RecvBufRespExtra {
repeated bytes tensor_content = 1;

View File

@ -51,8 +51,12 @@ fi
# Ensure that protoc-gen-go is available in $PATH
# Since ${PROTOC} will require it.
export PATH=$PATH:${GOPATH}/bin
mkdir -p ./internal/proto
${PROTOC} \
-I ${TF_DIR} \
--go_out=./internal/proto \
${TF_DIR}/tensorflow/core/framework/*.proto
mkdir -p ../vendor
for FILE in ${TF_DIR}/tensorflow/core/framework/*.proto \
${TF_DIR}/tensorflow/core/protobuf/*.proto \
${TF_DIR}/tensorflow/stream_executor/*.proto; do
${PROTOC} \
-I ${TF_DIR} \
--go_out=../vendor \
$FILE
done

View File

@ -22,7 +22,7 @@ import (
"github.com/golang/protobuf/proto"
tfpb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/github.com/tensorflow/tensorflow/tensorflow/go/core"
tfpb "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf"
)
// #include <stdlib.h>

View File

@ -16,7 +16,7 @@ limitations under the License.
package tensorflow
import tfpb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/github.com/tensorflow/tensorflow/tensorflow/go/core"
import tfpb "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf"
// #include "tensorflow/c/c_api.h"
import "C"

View File

@ -20,9 +20,8 @@ import (
"fmt"
"testing"
tspb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/github.com/tensorflow/tensorflow/tensorflow/go/core/framework/tensor_shape_go_proto"
typb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/github.com/tensorflow/tensorflow/tensorflow/go/core/framework/types_go_proto"
tfpb "github.com/tensorflow/tensorflow/tensorflow/go/genop/internal/proto/github.com/tensorflow/tensorflow/tensorflow/go/core"
frameworkpb "github.com/tensorflow/tensorflow/tensorflow/go/core/framework"
tfpb "github.com/tensorflow/tensorflow/tensorflow/go/core/protobuf"
)
func TestSignatureFromProto(t *testing.T) {
@ -32,9 +31,9 @@ func TestSignatureFromProto(t *testing.T) {
Encoding: &tfpb.TensorInfo_Name{
Name: "tensor_1",
},
Dtype: typb.DataType_DT_INT8,
TensorShape: &tspb.TensorShapeProto{
Dim: []*tspb.TensorShapeProto_Dim{
Dtype: frameworkpb.DataType_DT_INT8,
TensorShape: &frameworkpb.TensorShapeProto{
Dim: []*frameworkpb.TensorShapeProto_Dim{
{Size: 1},
{Size: 2},
{Size: 3},
@ -45,9 +44,9 @@ func TestSignatureFromProto(t *testing.T) {
Encoding: &tfpb.TensorInfo_Name{
Name: "tensor_2",
},
Dtype: typb.DataType_DT_FLOAT,
TensorShape: &tspb.TensorShapeProto{
Dim: []*tspb.TensorShapeProto_Dim{
Dtype: frameworkpb.DataType_DT_FLOAT,
TensorShape: &frameworkpb.TensorShapeProto{
Dim: []*frameworkpb.TensorShapeProto_Dim{
{Size: 4},
{Size: 5},
{Size: 6},
@ -60,9 +59,9 @@ func TestSignatureFromProto(t *testing.T) {
Encoding: &tfpb.TensorInfo_Name{
Name: "tensor_3",
},
Dtype: typb.DataType_DT_STRING,
TensorShape: &tspb.TensorShapeProto{
Dim: []*tspb.TensorShapeProto_Dim{
Dtype: frameworkpb.DataType_DT_STRING,
TensorShape: &frameworkpb.TensorShapeProto{
Dim: []*frameworkpb.TensorShapeProto_Dim{
{Size: 1},
{Size: 2},
{Size: 3},
@ -73,9 +72,9 @@ func TestSignatureFromProto(t *testing.T) {
Encoding: &tfpb.TensorInfo_Name{
Name: "tensor_4",
},
Dtype: typb.DataType_DT_BOOL,
TensorShape: &tspb.TensorShapeProto{
Dim: []*tspb.TensorShapeProto_Dim{
Dtype: frameworkpb.DataType_DT_BOOL,
TensorShape: &frameworkpb.TensorShapeProto{
Dim: []*frameworkpb.TensorShapeProto_Dim{
{Size: 4},
{Size: 5},
{Size: 6},
@ -144,9 +143,9 @@ func TestTensorInfoFromProto(t *testing.T) {
Encoding: &tfpb.TensorInfo_Name{
Name: "tensor",
},
Dtype: typb.DataType_DT_INT8,
TensorShape: &tspb.TensorShapeProto{
Dim: []*tspb.TensorShapeProto_Dim{
Dtype: frameworkpb.DataType_DT_INT8,
TensorShape: &frameworkpb.TensorShapeProto{
Dim: []*frameworkpb.TensorShapeProto_Dim{
{Size: 1},
{Size: 2},
{Size: 3},

View File

@ -3,6 +3,8 @@ syntax = "proto3";
package stream_executor.dnn;
option go_package = "github.com/tensorflow/tensorflow/tensorflow/go/stream_executor";
// Specifies the data type used by an operation.
enum DataType {
kFloat = 0;