diff options
author | 2017-06-29 11:44:13 -0700 | |
---|---|---|
committer | 2017-06-29 11:48:10 -0700 | |
commit | e85d3df92deb9d717befdf173966a2913ac2aea0 (patch) | |
tree | 1c9e80f85c965e039121d8591daab95c62fe7041 /tensorflow/core/ops/state_ops.cc | |
parent | 9b11f458196f6f0528c9974238497a6c8b6da547 (diff) |
Prepare to remove a bunch of proto.h includes from tensorflow/core headers
The goal is to make kernels mostly independent of proto headers, which will let
us lock down our .so imports. This CL does not remove any actual headers, but
changes a bunch of files so that header removal is possible in a followup CL.
It also marks the headers that will be removed with
// TODO(b/62899350): Remove
RELNOTES: n/a
PiperOrigin-RevId: 160552878
Diffstat (limited to 'tensorflow/core/ops/state_ops.cc')
-rw-r--r-- | tensorflow/core/ops/state_ops.cc | 16 |
1 files changed, 7 insertions, 9 deletions
diff --git a/tensorflow/core/ops/state_ops.cc b/tensorflow/core/ops/state_ops.cc index 35f965b6a9..1c9ae90a26 100644 --- a/tensorflow/core/ops/state_ops.cc +++ b/tensorflow/core/ops/state_ops.cc @@ -30,11 +30,11 @@ REGISTER_OP("VariableV2") .Attr("shared_name: string = ''") .SetIsStateful() .SetShapeFn([](InferenceContext* c) { - TensorShapeProto shape_proto; - TF_RETURN_IF_ERROR(c->GetAttr("shape", &shape_proto)); + PartialTensorShape shape; + TF_RETURN_IF_ERROR(c->GetAttr("shape", &shape)); ShapeHandle output_shape; TF_RETURN_IF_ERROR( - c->MakeShapeFromShapeProto(shape_proto, &output_shape)); + c->MakeShapeFromPartialTensorShape(shape, &output_shape)); c->set_output(0, output_shape); return Status::OK(); }) @@ -72,10 +72,8 @@ REGISTER_OP("Variable") return shape_inference::UnknownShape(c); } - TensorShapeProto shape_proto; - shape.AsProto(&shape_proto); ShapeHandle out; - TF_RETURN_IF_ERROR(c->MakeShapeFromShapeProto(shape_proto, &out)); + TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(shape, &out)); c->set_output(0, out); return Status::OK(); }) @@ -103,10 +101,10 @@ REGISTER_OP("TemporaryVariable") .Attr("var_name: string = ''") .SetIsStateful() .SetShapeFn([](InferenceContext* c) { - TensorShapeProto shape_proto; - TF_RETURN_IF_ERROR(c->GetAttr("shape", &shape_proto)); + PartialTensorShape shape; + TF_RETURN_IF_ERROR(c->GetAttr("shape", &shape)); ShapeHandle output; - TF_RETURN_IF_ERROR(c->MakeShapeFromShapeProto(shape_proto, &output)); + TF_RETURN_IF_ERROR(c->MakeShapeFromPartialTensorShape(shape, &output)); c->set_output(0, output); return Status::OK(); }) |