aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/core/ops/state_ops.cc
diff options
context:
space:
mode:
authorGravatar Geoffrey Irving <geoffreyi@google.com>2017-06-29 11:44:13 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2017-06-29 11:48:10 -0700
commite85d3df92deb9d717befdf173966a2913ac2aea0 (patch)
tree1c9e80f85c965e039121d8591daab95c62fe7041 /tensorflow/core/ops/state_ops.cc
parent9b11f458196f6f0528c9974238497a6c8b6da547 (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.cc16
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();
})