Fix funcdef lookup in auto_mixed_precision grappler pass

tensorflow/core/grappler/optimizers/auto_mixed_precision.cc

@@ -221,6 +221,8 @@ class NodeTypeAttrMap {
       return errors::InvalidArgument("NodeTypeAttrMap is already initialized.");
     }
     graph_ = &graph;
+    function_library_.reset(
+        new FunctionLibraryDefinition(OpRegistry::Global(), graph.library()));
     for (const NodeDef& node : graph.node()) {
       TF_RETURN_IF_ERROR(AddNode(node));
     }
@@ -272,10 +274,7 @@ class NodeTypeAttrMap {
  private:
   Status AddNode(const NodeDef& node) {
     const OpDef* op_def_ptr = nullptr;
-    // TODO(benbarsdell): This may fail if node.op() is a function. It will
-    // need to be addressed when we add support for functions.
-    TF_RETURN_IF_ERROR(
-        OpRegistry::Global()->LookUpOpDef(node.op(), &op_def_ptr));
+    TF_RETURN_IF_ERROR(function_library_->LookUpOpDef(node.op(), &op_def_ptr));
     const OpDef& op_def = *op_def_ptr;
     auto& type2io_entry = type2io_[&node];
     auto& io2type_entry = io2type_[&node];
@@ -343,6 +342,7 @@ class NodeTypeAttrMap {
 
   // WARN: `graph_` must outlive this object (node pointers must remain valid).
   const GraphDef* graph_ = nullptr;  // do not own
+  std::unique_ptr<FunctionLibraryDefinition> function_library_;
 
   typedef absl::flat_hash_set<int> IntSet;
   // Maps a type attr id -> (input port set, output port set)

tensorflow/python/grappler/auto_mixed_precision_test.py

@@ -28,6 +28,7 @@
 from tensorflow.python.compat import compat
 from tensorflow.python.framework import constant_op
 from tensorflow.python.framework import dtypes
+from tensorflow.python.framework import function
 from tensorflow.python.framework import random_seed
 from tensorflow.python.framework import test_util
 from tensorflow.python.ops import array_ops
@@ -251,6 +252,30 @@ def _build_node_map(nodes):
   return node_map
 
 
+def _example_noninlined_funcdef_shape(op):
+  return [op.inputs[0].shape]
+
+
+@function.Defun(shape_func=_example_noninlined_funcdef_shape,
+                func_name="example_noninlined_funcdef_grad", noinline=True)
+def _example_noninlined_funcdef_grad(features, grad):
+  """Gradient of Swish function defined below."""
+  sigmoid_features = math_ops.sigmoid(features)
+  activation_grad = (
+      sigmoid_features * (1.0 + features * (1.0 - sigmoid_features)))
+  return grad * activation_grad
+
+
+@function.Defun(
+    grad_func=_example_noninlined_funcdef_grad,
+    shape_func=_example_noninlined_funcdef_shape,
+    func_name="example_noninlined_funcdef",
+    noinline=True)
+def _example_noninlined_funcdef(features):
+  """Computes the Swish activation function: `x * sigmoid(x)`."""
+  return features * math_ops.sigmoid(features)
+
+
 class AutoMixedPrecisionTest(test.TestCase):
   """Tests the Grappler auto mixed precision optimizer."""
   IGNORE_PERF_VAR = 'TF_AUTO_MIXED_PRECISION_GRAPH_REWRITE_IGNORE_PERFORMANCE'
@@ -567,6 +592,28 @@ def test_propagation_through_simple_loop_7(self):
   def test_propagation_through_simple_loop_8(self):
     self._run_simple_loop_test('C', 'CgbgWC', 'g')
 
+  @test_util.run_deprecated_v1
+  def test_noninlined_funcdef(self):
+    """Test graph with non-inlined function subgraph.
+
+    This requires the grappler pass to handle an OpDef that only appears in the
+    graph's function registry instead of the global op registry.
+    """
+    if test.is_gpu_available(cuda_only=True):
+      random_seed.set_random_seed(0)
+      x = _input([8, 8])
+      y = _matmul_act(x)
+      y = _example_noninlined_funcdef(y)
+      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=0.01)
+      g = optimizer.compute_gradients(y, [x])
+      output = (g, y)
+
+      output_val_ref, output_val, cost_graph = self._run(output)
+      node_map = _build_node_map(cost_graph.node)
+
+      self._assert_output_fp16(node_map, 'MatMul')
+      self.assertAllClose(output_val_ref, output_val, atol=1e-3, rtol=1e-3)
+
 
 if __name__ == '__main__':
   test.main()