experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Mark saver/restore ops that depend on filesystem as stateful to disable them

Browse Source 
from being folded into a constant by graph optimizer.

PiperOrigin-RevId: 158182282
This commit is contained in:
A. Unique TensorFlower 2017-06-06 13:34:29 -07:00 committed by TensorFlower Gardener
parent 96cb4d1820
commit df5906fb72
1 changed files with 7 additions and 0 deletions

7
tensorflow/core/ops/io_ops.cc
View File

@ -62,6 +62,7 @@ REGISTER_OP("SaveV2")
.Input("shape_and_slices: string")
.Input("tensors: dtypes")
.Attr("dtypes: list(type)")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle unused;
ShapeHandle s;
@ -101,6 +102,7 @@ REGISTER_OP("RestoreV2")
.Input("shape_and_slices: string")
.Output("tensors: dtypes")
.Attr("dtypes: list(type)")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle shape0, shape1, shape2;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &shape0));
@ -141,6 +143,7 @@ REGISTER_OP("MergeV2Checkpoints")
.Input("checkpoint_prefixes: string")
.Input("destination_prefix: string")
.Attr("delete_old_dirs: bool = true")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle unused;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 1, &unused));
@ -169,6 +172,7 @@ REGISTER_OP("Save")
.Input("tensor_names: string")
.Input("data: T")
.Attr("T: list(type)")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle unused;
ShapeHandle s;
@ -204,6 +208,7 @@ REGISTER_OP("SaveSlices")
.Input("shapes_and_slices: string")
.Input("data: T")
.Attr("T: list(type)")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle unused;
ShapeHandle s;
@ -261,6 +266,7 @@ REGISTER_OP("Restore")
.Output("tensor: dt")
.Attr("dt: type")
.Attr("preferred_shard: int = -1")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle unused;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused));
@ -305,6 +311,7 @@ REGISTER_OP("RestoreSlice")
.Output("tensor: dt")
.Attr("dt: type")
.Attr("preferred_shard: int = -1")
.SetIsStateful()
.SetShapeFn([](InferenceContext* c) {
ShapeHandle unused;
TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 0, &unused));