Adds VariableV2 with a sane shape_fn.

Change: 141071094

1 files changed, 4 insertions, 3 deletions

diff --git a/tensorflow/core/common_runtime/function_test.cc b/tensorflow/core/common_runtime/function_test.cc
index 0f63e41957..03879f9ce6 100644
--- a/tensorflow/core/common_runtime/function_test.cc
+++ b/tensorflow/core/common_runtime/function_test.cc
@@ -229,6 +229,7 @@ class FunctionLibraryRuntimeTest : public ::testing::Test {
 TEST_F(FunctionLibraryRuntimeTest, IsStateful) {
   Init({});
   EXPECT_TRUE(lib_->IsStateful("Variable"));
+  EXPECT_TRUE(lib_->IsStateful("VariableV2"));
   EXPECT_FALSE(lib_->IsStateful("Matmul"));
 }
 
@@ -764,14 +765,14 @@ TEST(OptimizationTest, RemoveIdentityNodes_Ref) {
       {},
       // Nodes
       {// variable
-       {{"v"}, "Variable", {}, {{"dtype", T}, {"shape", TensorShape({})}}},
+       {{"v"}, "VariableV2", {}, {{"dtype", T}, {"shape", TensorShape({})}}},
        // read the variable. Shouldn't be removed.
        {{"v_read"}, "Identity", {"v"}, {{"T", T}}},
        // returns v + v
        {{"ret"}, "Add", {"v_read", "v_read"}, {{"T", T}}}});
   const char* e0 = R"S(
 () -> (n2:float) {
-  n0 = Variable[container="", dtype=float, shape=[], shared_name=""]()
+  n0 = VariableV2[container="", dtype=float, shape=[], shared_name=""]()
   n1 = Identity[T=float](n0)
   n2 = Add[T=float](n1, n1)
 }
@@ -780,7 +781,7 @@ TEST(OptimizationTest, RemoveIdentityNodes_Ref) {
 
   const char* e1 = R"S(
 () -> (n2:float) {
-  n0 = Variable[container="", dtype=float, shape=[], shared_name=""]()
+  n0 = VariableV2[container="", dtype=float, shape=[], shared_name=""]()
   n1 = Identity[T=float](n0)
   n2 = Add[T=float](n1, n1)
 }