Improve shape inference in TF 2.0, by propagating input shapes to functions.

This significantly improves the ability of Grappler to optimize tf.function bodies.

Fix a bug in ShapeOrHandleShape in training_ops.cc: If a resource argument does not have the handle_shapes available, we should return UnknownShape, rather than the shape of the argument itself, which is just a scalar resource handle.

PiperOrigin-RevId: 291991651
Change-Id: Iae8f3c5c5cf75742f0f2ce8f99cef2ae1ece5648

tensorflow/core/framework/graph_to_functiondef.cc

@@ -441,6 +441,12 @@ Status GraphToFunctionDef(const Graph& fn_body, const string& fn_name,
       // _Arg/Placeholder nodes.
       if (absl::StartsWith(attr.first, "_")) {
         arg_attrs.mutable_attr()->insert(attr);
+      } else if (attr.first == "shape") {
+        // Preserve known shapes by moving them to the _output_shapes list.
+        // The _Arg shape function knows how to extract them from there.
+        AttrValue value;
+        *(value.mutable_list()->add_shape()) = attr.second.shape();
+        arg_attrs.mutable_attr()->insert({"_output_shapes", value});
       }
       if (attr.first == "_resource_arg_unique_id") {
         resource_arg_unique_id = attr.second.i();

tensorflow/core/grappler/optimizers/meta_optimizer.cc

@@ -800,7 +800,7 @@ Status MetaOptimizer::Optimize(Cluster* cluster, const GrapplerItem& item,
 
   VLOG(1) << "Optimized " << optimized_funcs.size()
           << " functions: " << absl::StrJoin(optimized_funcs, ", ");
-
+  VLOG(3) << "Optimized graph =\n" << optimized_graph->DebugString();
   if (VLOG_IS_ON(1)) {
     DumpGraphDefToFile(
         strings::StrCat("after_MetaOptimizer_",

tensorflow/python/BUILD

@@ -7440,6 +7440,26 @@ cuda_py_test(
     ],
 )
 
+cuda_py_test(
+    name = "arithmetic_optimizer_test",
+    size = "small",
+    srcs = [
+        "grappler/arithmetic_optimizer_test.py",
+    ],
+    python_version = "PY3",
+    tags = [
+        "grappler",
+    ],
+    deps = [
+        ":array_ops",
+        ":client_testlib",
+        ":framework_for_generated_wrappers",
+        ":math_ops",
+        "//tensorflow/core:protos_all_py",
+        "//third_party/py/numpy",
+    ],
+)
+
 # TODO(b/131764887) Remove once LayoutOptimizer is swapped out with GenericLayoutOptimizer.
 #
 # cuda_py_test(

tensorflow/python/framework/function_def_to_graph.py

@@ -175,7 +175,9 @@ def function_def_to_graph_def(fdef, input_shapes=None):
     for k in arg_attrs:
       # Only copy internal attributes. Normal attributes for nodes cannot be
       # applied to these Placeholder nodes.
-      if k.startswith("_"):
+      if k == "_output_shapes":
+        node_def.attr["shape"].shape.CopyFrom(arg_attrs[k].list.shape[0])
+      elif k.startswith("_"):
         node_def.attr[k].CopyFrom(arg_attrs[k])
 
   # 2. Copy all body NodeDefs to the GraphDef.

tensorflow/python/grappler/arithmetic_optimizer_test.py

@@ -0,0 +1,50 @@
+# Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Tests for Grappler Arithmetic Optimizer."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from tensorflow.python.eager import context
+from tensorflow.python.eager import def_function
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import math_ops
+from tensorflow.python.platform import test
+
+
+class ArithmeticOptimizerTest(test.TestCase):
+
+  # See b/146524878.
+  def testFunctionArgShapeInference(self):
+
+    @def_function.function
+    def f(x, y):
+      return math_ops.matmul(
+          x, array_ops.reshape(array_ops.transpose(y), [384, 1536]))
+
+    with context.eager_mode():
+      x = array_ops.ones((1, 384))
+      y = array_ops.ones((1536, 384))
+      with context.collect_graphs(optimized=True) as graphs:
+        f(x, y).numpy()
+      self.assertLen(graphs, 1)
+      self.assertLen(graphs[0].node, 4)
+      self.assertEqual(graphs[0].node[2].name,
+                       'ArithmeticOptimizer/FoldTransposeIntoMatMul_MatMul')
+
+
+if __name__ == '__main__':
+  test.main()