Add shape inference functions for _Arg and _Retval.

The shape inference function for _Arg just returns UnknownShape, but the presence of a shape inference function makes ShapeRefiner happy.
Change: 145688141

1 files changed, 8 insertions, 0 deletions

diff --git a/tensorflow/core/ops/function_ops.cc b/tensorflow/core/ops/function_ops.cc
index 8de6123943..9fbebdb088 100644
--- a/tensorflow/core/ops/function_ops.cc
+++ b/tensorflow/core/ops/function_ops.cc
@@ -14,6 +14,7 @@ limitations under the License.
 ==============================================================================*/
 
 #include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
 
 namespace tensorflow {
 
@@ -22,6 +23,10 @@ REGISTER_OP("_Arg")
     .Attr("T: type")
     .Attr("index: int >= 0")
     .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext* context) {
+      context->set_output(0, context->UnknownShape());
+      return Status::OK();
+    })
     .Doc(R"doc(
 A graph node which represents an argument to a function.
 
@@ -34,6 +39,9 @@ REGISTER_OP("_Retval")
     .Attr("T: type")
     .Attr("index: int >= 0")
     .SetIsStateful()
+    .SetShapeFn([](shape_inference::InferenceContext* context) {
+      return Status::OK();
+    })
     .Doc(R"doc(
 A graph node which represents a return value of a function.