In C++ shape inference, support registering a shape inference function in OP

registration.  Change op registration signature to return Status and return the registration data as an out parameter.

Add a shape inference function and test for AddN.

Support in function library needs to be expanded in a future change.
Change: 124871850

1 files changed, 15 insertions, 20 deletions

diff --git a/tensorflow/core/framework/function_test.cc b/tensorflow/core/framework/function_test.cc
index 5849051062..d7e967ec3d 100644
--- a/tensorflow/core/framework/function_test.cc
+++ b/tensorflow/core/framework/function_test.cc
@@ -33,9 +33,7 @@ namespace tensorflow {
 typedef FunctionDefHelper FDH;
 
 Status GetOpSig(const string& op, const OpDef** sig) {
-  Status s;
-  *sig = OpRegistry::Global()->LookUp(op, &s);
-  return s;
+  return OpRegistry::Global()->LookUpOpDef(op, sig);
 }
 
 REGISTER_OP("One")
@@ -643,12 +641,12 @@ TEST(FunctionLibraryDefinitionTest, LookUp) {
   *proto.add_function() = test::function::XTimesTwo();
   FunctionLibraryDefinition lib_def(proto);
 
-  Status s;
-  EXPECT_EQ(lib_def.LookUp("XTimes16", &s), nullptr);
+  const OpDef* op_def;
+  EXPECT_TRUE(!lib_def.LookUpOpDef("XTimes16", &op_def).ok());
 
-  auto found = lib_def.LookUp("XTimesTwo", &s);
-  ASSERT_NE(found, nullptr);
-  EXPECT_EQ(found->DebugString(),
+  TF_EXPECT_OK(lib_def.LookUpOpDef("XTimesTwo", &op_def));
+  ASSERT_NE(op_def, nullptr);
+  EXPECT_EQ(op_def->DebugString(),
             test::function::XTimesTwo().signature().DebugString());
 }
 
@@ -662,14 +660,15 @@ TEST(FunctionLibraryDefinitionTest, AddFunctionDef) {
   TF_EXPECT_OK(lib_def.AddFunctionDef(test::function::WXPlusB()));
 
   // Test lookup of first function.
-  Status s;
-  auto first = lib_def.LookUp("XTimesTwo", &s);
+  const OpDef* first;
+  TF_EXPECT_OK(lib_def.LookUpOpDef("XTimesTwo", &first));
   ASSERT_NE(first, nullptr);
   EXPECT_EQ(first->DebugString(),
             test::function::XTimesTwo().signature().DebugString());
 
   // Test lookup of second function.
-  auto second = lib_def.LookUp("WXPlusB", &s);
+  const OpDef* second;
+  TF_EXPECT_OK(lib_def.LookUpOpDef("WXPlusB", &second));
   ASSERT_NE(second, nullptr);
   EXPECT_EQ(second->DebugString(),
             test::function::WXPlusB().signature().DebugString());
@@ -689,18 +688,14 @@ TEST(FunctionLibraryDefinitionTest, ToProto) {
   FunctionLibraryDefinition lib_def2(proto2);
 
   // Test that the first function exists in both libraries.
-  Status s;
-  auto f1 = lib_def1.LookUp("XTimesTwo", &s);
-  TF_EXPECT_OK(s);
-  auto f2 = lib_def1.LookUp("XTimesTwo", &s);
-  TF_EXPECT_OK(s);
+  const OpDef *f1, *f2, *f3, *f4;
+  TF_EXPECT_OK(lib_def1.LookUpOpDef("XTimesTwo", &f1));
+  TF_EXPECT_OK(lib_def2.LookUpOpDef("XTimesTwo", &f2));
   EXPECT_EQ(f1->DebugString(), f2->DebugString());
 
   // Test that the second function exists in both libraries.
-  auto f3 = lib_def1.LookUp("WXPlusB", &s);
-  TF_EXPECT_OK(s);
-  auto f4 = lib_def1.LookUp("WXPlusB", &s);
-  TF_EXPECT_OK(s);
+  TF_EXPECT_OK(lib_def1.LookUpOpDef("WXPlusB", &f3));
+  TF_EXPECT_OK(lib_def2.LookUpOpDef("WXPlusB", &f4));
   EXPECT_EQ(f3->DebugString(), f4->DebugString());
 }