Expose autograph related API as tf.__internal__.

These methods are used by Keras for creating autographed functions.

PiperOrigin-RevId: 355498359
Change-Id: I189d6005905ab1478935d8fd137ac92fe04296d1

tensorflow/python/autograph/core/ag_ctx.py

@@ -18,9 +18,10 @@ from __future__ import absolute_import
 from __future__ import division
 from __future__ import print_function
 
+import enum
 import threading
 
-import enum
+from tensorflow.python.util.tf_export import tf_export
 
 
 stacks = threading.local()
@@ -32,7 +33,23 @@ def _control_ctx():
   return stacks.control_status
 
 
+@tf_export('__internal__.autograph.control_status_ctx', v1=[])
 def control_status_ctx():
+  """Returns the current control context for autograph.
+
+  This method is useful when calling `tf.__internal__.autograph.tf_convert`,
+  The context will be used by tf_convert to determine whether it should convert
+  the input function. See the sample usage like below:
+
+  ```
+  def foo(func):
+    return tf.__internal__.autograph.tf_convert(
+       input_fn, ctx=tf.__internal__.autograph.control_status_ctx())()
+  ```
+
+  Returns:
+    The current control context of autograph.
+  """
   ret = _control_ctx()[-1]
   return ret
 

tensorflow/python/autograph/impl/api.py

@@ -511,17 +511,56 @@ def _fall_back_unconverted(f, args, kwargs, options, exc):
 #
 
 
+@tf_export('__internal__.autograph.tf_convert', v1=[])
 def tf_convert(f, ctx, convert_by_default=True, user_requested=False):
   """Decorator that applies AutoGraph to a function.
 
   Use in internal APIs.
 
   This API is suitable for high order functions internal to the TensorFlow API,
-  and more generally any function to which Autograph is not applied.
+  and more generally any function to which AutoGraph is not applied.
 
-  Guidance: convert was a decorator meant for use directly by developers, and
-  will be soon deprecated in favor of tf.function. tf_convert is to be called
-  from high order functions internal to TF.
+  Guidance: `convert` was a decorator meant for use directly by developers, but
+  most of today's uses go through `tf.function`. `tf_convert` is to be called
+  from high order functions internal to TF. By default, all the internal
+  TensorFlow functions are skipped when AutoGraph processes the code. This may
+  lead to user-supplied functions to be incorrectly skipped as well.
+  `tf_convert` helps avoid that. See the following example for more details.
+
+  ```
+  =====tf_internal_module.py=====
+
+  def unconverted(input_fn):
+    return input_fn()
+
+  def converted(input_fn):
+    return tf.__internal__.autograph.tf_convert(
+       input_fn, ctx=tf.__internal__.autograph.control_status_ctx())()
+
+  ======user_module.py======
+
+  @tf.function
+  def foo(input_fn)
+    return unconverted(input_fn)
+
+  @tf.function
+  def bar(input_fn)
+    return converted(input_fn)
+
+  @tf.function(autograph=False)
+  def baz(input_fn)
+    return converted(input_fn)
+  ```
+
+  The `foo` method above will execute the `input_fn` without autograph
+  conversion, while the `bar` method will run an autographed `input_fn`. The
+  `baz` method will run an unconverted `input_fn`, since `tf_convert` respect
+  the control status context.
+
+  Note that both methods in `tf_internal_module` are skipped by autograph when
+  tracing the `tf.function`. The configuration of whether a module/package
+  should be skipped by autograph is controlled in
+  tensorflow/python/autograph/core/config.py.
 
   Args:
     f: Callable.

tensorflow/python/tools/api/generator/api_init_files.bzl

@@ -5,6 +5,7 @@ TENSORFLOW_API_INIT_FILES = [
     # BEGIN GENERATED FILES
     "__init__.py",
     "__internal__/__init__.py",
+    "__internal__/autograph/__init__.py",
     "__internal__/decorator/__init__.py",
     "__internal__/dispatch/__init__.py",
     "__internal__/distribute/__init__.py",

tensorflow/tools/api/golden/v2/tensorflow.__internal__.autograph.pbtxt

@@ -0,0 +1,11 @@
+path: "tensorflow.__internal__.autograph"
+tf_module {
+  member_method {
+    name: "control_status_ctx"
+    argspec: "args=[], varargs=None, keywords=None, defaults=None"
+  }
+  member_method {
+    name: "tf_convert"
+    argspec: "args=[\'f\', \'ctx\', \'convert_by_default\', \'user_requested\'], varargs=None, keywords=None, defaults=[\'True\', \'False\'], "
+  }
+}

tensorflow/tools/api/golden/v2/tensorflow.__internal__.pbtxt

@@ -12,6 +12,10 @@ tf_module {
     name: "FuncGraph"
     mtype: "<type \'type\'>"
   }
+  member {
+    name: "autograph"
+    mtype: "<type \'module\'>"
+  }
   member {
     name: "decorator"
     mtype: "<type \'module\'>"