chore: import upstream snapshot with attribution

tensorflow/python/autograph/impl/BUILD

@@ -0,0 +1,116 @@
+load("@xla//third_party/rules_python/python:py_library.bzl", "py_library")
+load("//tensorflow:tensorflow.default.bzl", "tf_py_strict_test")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    licenses = ["notice"],
+)
+
+py_library(
+    name = "api",
+    srcs = ["api.py"],
+    strict_deps = True,
+    visibility = ["//tensorflow:__subpackages__"],
+    deps = [
+        ":conversion",
+        "//tensorflow/python/autograph/converters:asserts",
+        "//tensorflow/python/autograph/converters:break_statements",
+        "//tensorflow/python/autograph/converters:call_trees",
+        "//tensorflow/python/autograph/converters:conditional_expressions",
+        "//tensorflow/python/autograph/converters:continue_statements",
+        "//tensorflow/python/autograph/converters:control_flow",
+        "//tensorflow/python/autograph/converters:directives",
+        "//tensorflow/python/autograph/converters:functions",
+        "//tensorflow/python/autograph/converters:lists",
+        "//tensorflow/python/autograph/converters:logical_expressions",
+        "//tensorflow/python/autograph/converters:return_statements",
+        "//tensorflow/python/autograph/converters:slices",
+        "//tensorflow/python/autograph/converters:variables",
+        "//tensorflow/python/autograph/core:ag_ctx",
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/core:function_wrappers",
+        "//tensorflow/python/autograph/core:unsupported_features_checker",
+        "//tensorflow/python/autograph/lang:special_functions",
+        "//tensorflow/python/autograph/operators:__init__",
+        "//tensorflow/python/autograph/operators:py_builtins",
+        "//tensorflow/python/autograph/pyct:anno",
+        "//tensorflow/python/autograph/pyct:cfg",
+        "//tensorflow/python/autograph/pyct:error_utils",
+        "//tensorflow/python/autograph/pyct:errors",
+        "//tensorflow/python/autograph/pyct:inspect_utils",
+        "//tensorflow/python/autograph/pyct:origin_info",
+        "//tensorflow/python/autograph/pyct:qual_names",
+        "//tensorflow/python/autograph/pyct:transpiler",
+        "//tensorflow/python/autograph/pyct/static_analysis:activity",
+        "//tensorflow/python/autograph/pyct/static_analysis:reaching_definitions",
+        "//tensorflow/python/autograph/utils:__init__",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/eager/polymorphic_function:tf_method_target",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/util:tf_decorator_py",
+        "//tensorflow/python/util:tf_export",
+        "//tensorflow/python/util:tf_inspect",
+        "//tensorflow/python/util:tf_stack",
+    ],
+)
+
+py_library(
+    name = "conversion",
+    srcs = ["conversion.py"],
+    strict_deps = True,
+    visibility = ["//visibility:private"],
+    deps = [
+        "//tensorflow/python/autograph/core:config",
+        "//tensorflow/python/autograph/pyct:cache",
+        "//tensorflow/python/autograph/pyct:inspect_utils",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/eager/polymorphic_function:tf_method_target",
+        "//tensorflow/python/util:tf_inspect",
+    ],
+)
+
+tf_py_strict_test(
+    name = "api_test",
+    srcs = ["api_test.py"],
+    deps = [
+        ":api",
+        ":conversion",
+        "//tensorflow/python/autograph/core:ag_ctx",
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/core:test_lib",
+        "//tensorflow/python/autograph/pyct:errors",
+        "//tensorflow/python/autograph/pyct:inspect_utils",
+        "//tensorflow/python/autograph/pyct:parser",
+        "//tensorflow/python/autograph/utils:ag_logging",
+        "//tensorflow/python/data/ops:dataset_ops",
+        "//tensorflow/python/eager:def_function",
+        "//tensorflow/python/framework:_errors_test_helper",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/framework:errors",
+        "//tensorflow/python/framework:ops",
+        "//tensorflow/python/framework:test_lib",
+        "//tensorflow/python/ops:math_ops",
+        "//tensorflow/python/ops:variables",
+        "//tensorflow/python/platform:client_testlib",
+        "//tensorflow/python/util:function_utils",
+        "//tensorflow/python/util:tf_decorator_py",
+        "//tensorflow/python/util:tf_inspect",
+        "//third_party/py/numpy",
+    ],
+)
+
+tf_py_strict_test(
+    name = "conversion_test",
+    srcs = ["conversion_test.py"],
+    deps = [
+        ":api",
+        ":conversion",
+        "//tensorflow/python/autograph/core:config",
+        "//tensorflow/python/autograph/core:converter",
+        "//tensorflow/python/autograph/impl/testing:pybind_for_testing",
+        "//tensorflow/python/autograph/utils:__init__",
+        "//tensorflow/python/eager:function",
+        "//tensorflow/python/framework:constant_op",
+        "//tensorflow/python/platform:client_testlib",
+    ],
+)

tensorflow/python/autograph/impl/api.py

@@ -0,0 +1,949 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""This module contains the user- and codegen-facing API for AutoGraph."""
+
+import functools
+import importlib
+import inspect
+import os
+import sys
+import textwrap
+import traceback
+
+from tensorflow.python.autograph import operators
+from tensorflow.python.autograph import utils
+from tensorflow.python.autograph.converters import asserts
+from tensorflow.python.autograph.converters import break_statements
+from tensorflow.python.autograph.converters import call_trees
+from tensorflow.python.autograph.converters import conditional_expressions
+from tensorflow.python.autograph.converters import continue_statements
+from tensorflow.python.autograph.converters import control_flow
+from tensorflow.python.autograph.converters import directives
+from tensorflow.python.autograph.converters import functions
+from tensorflow.python.autograph.converters import lists
+from tensorflow.python.autograph.converters import logical_expressions
+from tensorflow.python.autograph.converters import return_statements
+from tensorflow.python.autograph.converters import slices
+from tensorflow.python.autograph.converters import variables
+from tensorflow.python.autograph.core import ag_ctx
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.core import function_wrappers
+from tensorflow.python.autograph.core import unsupported_features_checker
+from tensorflow.python.autograph.impl import conversion
+from tensorflow.python.autograph.lang import special_functions
+from tensorflow.python.autograph.operators import py_builtins
+from tensorflow.python.autograph.pyct import anno
+from tensorflow.python.autograph.pyct import cfg
+from tensorflow.python.autograph.pyct import error_utils
+from tensorflow.python.autograph.pyct import errors
+from tensorflow.python.autograph.pyct import inspect_utils
+from tensorflow.python.autograph.pyct import origin_info
+from tensorflow.python.autograph.pyct import qual_names
+from tensorflow.python.autograph.pyct import transpiler
+from tensorflow.python.autograph.pyct.static_analysis import activity
+from tensorflow.python.autograph.pyct.static_analysis import reaching_definitions
+from tensorflow.python.autograph.utils import ag_logging as logging
+from tensorflow.python.eager.polymorphic_function import tf_method_target
+from tensorflow.python.framework import errors_impl
+from tensorflow.python.util import tf_decorator
+from tensorflow.python.util import tf_inspect
+from tensorflow.python.util import tf_stack
+from tensorflow.python.util.tf_export import tf_export
+
+
+def is_autograph_strict_conversion_mode():
+  return int(os.environ.get('AUTOGRAPH_STRICT_CONVERSION', '0')) > 0
+
+
+#
+# Error handling
+#
+
+
+# TODO(mdan): Export this symbol.
+class AutoGraphError(errors.PyCTError):
+  """Base class for all AutoGraph exceptions."""
+  pass
+
+
+class ConversionError(AutoGraphError):
+  """Raised during the conversion process."""
+  pass
+
+
+class StagingError(AutoGraphError):
+  """Raised during the staging (i.e. Python execution) of converted code."""
+  pass
+
+
+class _ErrorMetadata(error_utils.ErrorMetadataBase):
+  """AutoGraph-specific error metadata. See base class."""
+
+  def create_exception(self, source_error):
+    preferred_type = type(source_error)
+    if issubclass(preferred_type, errors_impl.OpError):
+      # Best-effort unpacking of OpError exceptions.
+      # TODO(mdan): Use a mechanism that is more future-proof.
+      init_argspec = tf_inspect.getfullargspec(preferred_type.__init__)
+      message = self.get_message()
+      init_args = tuple(init_argspec.args)
+      # At the time of this writing, TF errors either take 3 or 4 arguments,
+      # the argument '*args' may or may not be used.
+      if init_args == ('self', 'node_def', 'op', 'message'):
+        return preferred_type(source_error.node_def, source_error.op, message,
+                              source_error.experimental_payloads)
+
+    elif preferred_type in (errors.PyCTError, AutoGraphError, ConversionError,
+                            StagingError, errors_impl.InaccessibleTensorError,
+                            errors_impl.OperatorNotAllowedInGraphError):
+      return preferred_type(self.get_message())
+
+    exc = super(_ErrorMetadata, self).create_exception(source_error)
+    if exc is not None:
+      return exc
+
+    # Note: While changing an error's message property to change the message it
+    # displays will probably work a lot of times, there is no standard way in
+    # Python to do that. The safest way is therefore to create a new exception.
+    # For user defined exceptions, we could define an interface that allowed
+    # them to work under this mechanism.
+    return StagingError(self.get_message())
+
+
+def _attach_error_metadata(e, f):
+  """Augments an error with the metadata necessary for rewrite."""
+  if hasattr(e, 'ag_pass_through'):
+    return
+
+  metadata = getattr(e, 'ag_error_metadata', None)
+  source_map = f.ag_source_map
+
+  if metadata is None:
+    logging.log(1, 'Caught error in user callable %s', f, exc_info=True)
+    message = '{}: {}'.format(e.__class__.__name__, e)
+  else:
+    message = None
+
+  cause_tb = traceback.extract_tb(sys.exc_info()[2])[1:]
+
+  e.ag_error_metadata = _ErrorMetadata(cause_tb, metadata, message, source_map,
+                                       __file__)
+
+
+class StackTraceMapper(tf_stack.StackTraceMapper):
+  """Remaps generated code to code it originated from."""
+
+  def __init__(self, converted_fn):
+    super().__init__()
+    self._source_map = converted_fn.ag_source_map
+    # This may be called repeatedly: once on entry, by the superclass, then by
+    # each child context manager.
+    self._cached_map = None
+
+  def get_effective_source_map(self):
+    if self._cached_map is not None:
+      return self._cached_map
+
+    parent_map = self.parent.get_effective_source_map()
+
+    effective_source_map = {}
+    for loc, origin in self._source_map.items():
+      effective_source_map[(loc.filename, loc.lineno)] = (origin.loc.filename,
+                                                          origin.loc.lineno,
+                                                          origin.function_name)
+
+    for key, value in parent_map.items():
+      filename, lineno, _ = value
+      value_loc = origin_info.LineLocation(filename=filename, lineno=lineno)
+      if value_loc in self._source_map:
+        origin = self._source_map[value_loc]
+        effective_source_map[key] = (origin.loc.filename, origin.loc.lineno,
+                                     origin.function_name)
+      else:
+        effective_source_map[key] = value
+
+    self._cached_map = effective_source_map
+    return effective_source_map
+
+
+#
+# Actual source code transformation
+#
+
+
+class PyToTF(transpiler.PyToPy):
+  """The TensorFlow AutoGraph transformer."""
+
+  def __init__(self):
+    super(PyToTF, self).__init__()
+    self._extra_locals = None
+
+  def get_transformed_name(self, node):
+    return 'tf__' + super(PyToTF, self).get_transformed_name(node)
+
+  def get_extra_locals(self):
+    if self._extra_locals is None:
+      # TODO(mdan): Move into core or replace with an actual importable module.
+      # Craft a module that exposes the external API as well as certain
+      # internal modules.
+      module_spec = importlib.machinery.ModuleSpec('autograph', None)
+      ag_internal = importlib.util.module_from_spec(module_spec)
+      ag_internal.__dict__.update(inspect.getmodule(PyToTF).__dict__)
+      ag_internal.ConversionOptions = converter.ConversionOptions
+      ag_internal.STD = converter.STANDARD_OPTIONS
+      ag_internal.Feature = converter.Feature
+      ag_internal.utils = utils
+      ag_internal.FunctionScope = function_wrappers.FunctionScope
+      ag_internal.with_function_scope = function_wrappers.with_function_scope
+      # TODO(mdan): Add safeguards against name clashes.
+      # We don't want to create a submodule because we want the operators to be
+      # accessible as ag__.<operator>
+      ag_internal.__dict__.update(special_functions.__dict__)
+      ag_internal.__dict__.update(operators.__dict__)
+
+      self._extra_locals = {'ag__': ag_internal}
+    return self._extra_locals
+
+  def get_caching_key(self, ctx):
+    return ctx.options
+
+  def initial_analysis(self, node, ctx):
+    graphs = cfg.build(node)
+    node = qual_names.resolve(node)
+    node = activity.resolve(node, ctx, None)
+    node = reaching_definitions.resolve(node, ctx, graphs)
+    anno.dup(
+        node,
+        {
+            anno.Static.DEFINITIONS: anno.Static.ORIG_DEFINITIONS,
+        },
+    )
+    return node
+
+  def transform_ast(self, node, ctx):
+    unsupported_features_checker.verify(node)
+    node = self.initial_analysis(node, ctx)
+
+    node = functions.transform(node, ctx)
+    node = directives.transform(node, ctx)
+    node = break_statements.transform(node, ctx)
+    if ctx.user.options.uses(converter.Feature.ASSERT_STATEMENTS):
+      node = asserts.transform(node, ctx)
+    # Note: sequencing continue canonicalization before for loop one avoids
+    # dealing with the extra loop increment operation that the for
+    # canonicalization creates.
+    node = continue_statements.transform(node, ctx)
+    node = return_statements.transform(node, ctx)
+    if ctx.user.options.uses(converter.Feature.LISTS):
+      node = lists.transform(node, ctx)
+      node = slices.transform(node, ctx)
+    node = call_trees.transform(node, ctx)
+    node = control_flow.transform(node, ctx)
+    node = conditional_expressions.transform(node, ctx)
+    node = logical_expressions.transform(node, ctx)
+    node = variables.transform(node, ctx)
+    return node
+
+
+def _convert_actual(entity, program_ctx):
+  """Applies AutoGraph to entity."""
+
+  # TODO(mdan): Put these extra fields inside __autograph_info__.
+  if not hasattr(entity, '__code__'):
+    raise ValueError('Cannot apply autograph to a function that doesn\'t '
+                     'expose a __code__ object. If this is a @tf.function,'
+                     ' try passing f.python_function instead.')
+
+  transformed, module, source_map = _TRANSPILER.transform(entity, program_ctx)
+
+  assert not hasattr(transformed, 'ag_module')
+  assert not hasattr(transformed, 'ag_source_map')
+  transformed.ag_module = module
+  transformed.ag_source_map = source_map
+  return transformed
+
+
+#
+# Generated code support
+#
+
+
+def autograph_artifact(entity, extras=None):
+  if inspect.ismethod(entity):
+    setattr(entity.__func__, 'autograph_info__', extras)
+  else:
+    setattr(entity, 'autograph_info__', extras)
+  return entity
+
+
+def is_autograph_artifact(entity):
+  return hasattr(entity, 'autograph_info__')
+
+
+def converted_call(f, args, kwargs, caller_fn_scope=None, options=None):
+  """Converts a function call inline.
+
+  For internal use only.
+
+  Note: The argument list is optimized for readability of generated code, which
+  may look like this:
+
+    ag__.converted_call(f, (arg1, arg2), None, fscope)
+    ag__.converted_call(f, (), dict(arg1=val1, **kwargs), fscope)
+    ag__.converted_call(f, (arg1, arg2) + varargs, dict(**kwargs), lscope)
+
+  Args:
+    f: The function to convert.
+    args: Tuple, the original positional arguments of f
+    kwargs: Optional[Dict], the original keyword arguments of f
+    caller_fn_scope: Optional[function_wrappers.FunctionScope], the function
+      scope of the converted function in which this call was originally made.
+    options: Optional[converter.ConversionOptions], conversion options. If not
+      specified, the value of caller_fn_scope.callopts is used. Either options
+      or caller_fn_scope must be present.
+
+  Returns:
+    Any, the result of executing a possibly-converted `f` with the given
+      arguments.
+  """
+  logging.log(1, 'Converted call: %s\n    args: %s\n    kwargs: %s\n', f, args,
+              kwargs)
+
+  if options is None:
+    if caller_fn_scope is None:
+      raise ValueError('either caller_fn_scope or options must have a value')
+    options = caller_fn_scope.callopts
+
+  if conversion.is_in_allowlist_cache(f, options):
+    logging.log(2, 'Allowlisted %s: from cache', f)
+    return _call_unconverted(f, args, kwargs, options, False)
+
+  if ag_ctx.control_status_ctx().status == ag_ctx.Status.DISABLED:
+    logging.log(2, 'Allowlisted: %s: AutoGraph is disabled in context', f)
+    return _call_unconverted(f, args, kwargs, options, False)
+
+  if is_autograph_artifact(f):
+    logging.log(2, 'Permanently allowed: %s: AutoGraph artifact', f)
+    return _call_unconverted(f, args, kwargs, options)
+
+  # If this is a partial, unwrap it and redo all the checks.
+  if isinstance(f, functools.partial):
+    new_kwargs = {}
+    if f.keywords is not None:
+      # Use copy to avoid mutating the underlying keywords.
+      new_kwargs = f.keywords.copy()
+    if kwargs is not None:
+      new_kwargs.update(kwargs)
+    new_args = f.args + args
+    logging.log(3, 'Forwarding call of partial %s with\n%s\n%s\n', f, new_args,
+                new_kwargs)
+    return converted_call(
+        f.func,
+        new_args,
+        new_kwargs,
+        caller_fn_scope=caller_fn_scope,
+        options=options)
+
+  if inspect_utils.isbuiltin(f):
+    if f is eval:
+      return py_builtins.eval_in_original_context(f, args, caller_fn_scope)
+    if f is super:
+      return py_builtins.super_in_original_context(f, args, caller_fn_scope)
+    if f is globals:
+      return py_builtins.globals_in_original_context(caller_fn_scope)
+    if f is locals:
+      return py_builtins.locals_in_original_context(caller_fn_scope)
+    if kwargs:
+      return py_builtins.overload_of(f)(*args, **kwargs)
+    else:
+      return py_builtins.overload_of(f)(*args)
+
+  if conversion.is_unsupported(f):
+    return _call_unconverted(f, args, kwargs, options)
+
+  if not options.user_requested and conversion.is_allowlisted(f):
+    return _call_unconverted(f, args, kwargs, options)
+
+  # internal_convert_user_code is for example turned off when issuing a dynamic
+  # call conversion from generated code while in nonrecursive mode. In that
+  # case we evidently don't want to recurse, but we still have to convert
+  # things like builtins.
+  if not options.internal_convert_user_code:
+    return _call_unconverted(f, args, kwargs, options)
+
+  try:
+    if inspect.ismethod(f) or inspect.isfunction(f):
+      target_entity = f
+      effective_args = args
+
+      f_self = getattr(f, '__self__', None)
+      if f_self is not None:
+        if isinstance(f_self, tf_method_target.TfMethodTarget):
+          f_self = f_self.target
+        effective_args = (f_self,) + effective_args
+
+    elif hasattr(f, '__class__') and hasattr(f.__class__, '__call__'):
+      # Callable objects. Dunder methods have special lookup rules, see:
+      # https://docs.python.org/3/reference/datamodel.html#specialnames
+      # TODO(mdan): Recurse into converted_call to simplify other verifications.
+      # This should be handled in the same way as partials.
+      target_entity = f.__class__.__call__
+      effective_args = (f,) + args
+
+    else:
+      target_entity = f
+      raise NotImplementedError('unknown callable type "%s"' % type(f))
+
+  except Exception as e:  # pylint:disable=broad-except
+    logging.log(1, 'Error transforming entity %s', target_entity, exc_info=True)
+    if is_autograph_strict_conversion_mode():
+      raise
+    return _fall_back_unconverted(f, args, kwargs, options, e)
+
+  if not hasattr(target_entity, '__code__'):
+    logging.log(2, 'Permanently allowed: %s: native binding', target_entity)
+    return _call_unconverted(f, args, kwargs, options)
+  elif (hasattr(target_entity.__code__, 'co_filename') and
+        target_entity.__code__.co_filename == '<string>'):
+    # TODO(mdan): __globals__['txt'] might work in Py3.
+    logging.log(2, 'Permanently allowed: %s: dynamic code (exec?)',
+                target_entity)
+    return _call_unconverted(f, args, kwargs, options)
+
+  try:
+    program_ctx = converter.ProgramContext(options=options)
+    converted_f = _convert_actual(target_entity, program_ctx)
+    if logging.has_verbosity(2):
+      _log_callargs(converted_f, effective_args, kwargs)
+  except Exception as e:  # pylint:disable=broad-except
+    logging.log(1, 'Error transforming entity %s', target_entity, exc_info=True)
+    if is_autograph_strict_conversion_mode():
+      raise
+    return _fall_back_unconverted(f, args, kwargs, options, e)
+
+  with StackTraceMapper(converted_f), tf_stack.CurrentModuleFilter():
+    try:
+      if kwargs is not None:
+        result = converted_f(*effective_args, **kwargs)
+      else:
+        result = converted_f(*effective_args)
+    except Exception as e:
+      _attach_error_metadata(e, converted_f)
+      raise
+
+  return result
+
+
+def _call_unconverted(f, args, kwargs, options, update_cache=True):
+  """Calls the original function without converting with AutoGraph."""
+  if update_cache:
+    conversion.cache_allowlisted(f, options)
+
+  if (inspect.ismethod(f) and
+      isinstance(f.__self__, tf_method_target.TfMethodTarget)):
+    return f.__self__.call(args, kwargs)
+
+  if kwargs is not None:
+    return f(*args, **kwargs)
+  return f(*args)
+
+
+def _fall_back_unconverted(f, args, kwargs, options, exc):
+  """Falls back to calling the function unconverted, in case of error."""
+  # TODO(mdan): Consider adding an internal metric.
+  warning_template = (
+      'AutoGraph could not transform %s and will run it as-is.\n'
+      '%s'
+      'Cause: %s\n'
+      'To silence this warning, decorate the function with'
+      ' @tf.autograph.experimental.do_not_convert')
+  if isinstance(exc, errors.InaccessibleSourceCodeError):
+    if ag_ctx.INSPECT_SOURCE_SUPPORTED:
+      logging.warning(warning_template, f, '', exc)
+  elif isinstance(exc, errors.UnsupportedLanguageElementError):
+    if not conversion.is_in_allowlist_cache(f, options):
+      logging.warning(warning_template, f, '', exc)
+  else:
+    file_bug_message = (
+        'Please report this to the TensorFlow team. When filing the bug, set'
+        ' the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and'
+        ' attach the full output.\n')
+    logging.warning(warning_template, f, file_bug_message, exc)
+
+  return _call_unconverted(f, args, kwargs, options)
+
+
+#
+# TensorFlow integration
+#
+
+
+@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.
+
+  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.
+    ctx: ag_ctx.ControlStatusCtx, the Autograph context in which `f` is used.
+    convert_by_default: bool, whether to use AutoGraph when the context doesn't
+      specify.
+    user_requested: bool, whether to ignore the conversion allowlist. See
+      ConversionOptions.user_requested.
+
+  Returns:
+    Either `f or the converted version of `f`.
+  """
+
+  if is_autograph_artifact(f):
+    return f
+  f_wrapper = f
+  decorators, f = tf_decorator.unwrap(f)
+
+  # TODO(mdan): Grab features from context.
+  # Note: we pass the original context through to convert to properly handle the
+  # following scenario, which can be used inside TF implementations:
+  #
+  #   ctx = ag_ctx.control_status_ctx()
+  #   @function(autograph=False)  # Low-level graph code
+  #   def inner_fn():
+  #     # The context is disabled here, but should be enabled in user user_fn
+  #     tf_convert(user_fn, ctx=ctx)
+  if ctx.status == ag_ctx.Status.ENABLED:
+    wrapper_factory = convert(
+        recursive=True, user_requested=user_requested, conversion_ctx=ctx)
+  elif ctx.status == ag_ctx.Status.DISABLED:
+    wrapper_factory = do_not_convert
+  elif ctx.status == ag_ctx.Status.UNSPECIFIED:
+    if convert_by_default:
+      wrapper_factory = convert(
+          recursive=True, user_requested=user_requested, conversion_ctx=ctx)
+    else:
+      wrapper_factory = call_with_unspecified_conversion_status
+  else:
+    assert False, 'This switch contains all possible cases!'
+  wrapper = wrapper_factory(f)
+
+  if decorators:
+    wrapper = tf_decorator.rewrap(f_wrapper, f, wrapper)
+
+  return autograph_artifact(wrapper)
+
+
+def call_with_unspecified_conversion_status(func):
+  """Decorator that resets the conversion context to the unspecified status."""
+
+  def wrapper(*args, **kwargs):
+    with ag_ctx.ControlStatusCtx(status=ag_ctx.Status.UNSPECIFIED):
+      return func(*args, **kwargs)
+
+  if inspect.isfunction(func) or inspect.ismethod(func):
+    wrapper = functools.update_wrapper(wrapper, func)
+
+  return autograph_artifact(wrapper)
+
+
+def _log_callargs(f, args, kwargs):
+  """Logging helper."""
+  logging.log(2, 'Defaults of %s : %s', f, f.__defaults__)
+  logging.log(2, 'KW defaults of %s : %s', f, f.__kwdefaults__)
+
+  if kwargs is not None:
+    callargs = tf_inspect.getcallargs(f, *args, **kwargs)
+  else:
+    callargs = tf_inspect.getcallargs(f, *args)
+
+  formatted_callargs = '\n'.join(
+      '    {}: {}'.format(k, v) for k, v in callargs.items())
+  logging.log(2, 'Calling %s with\n%s\n', f, formatted_callargs)
+
+
+#
+# Public API
+#
+
+
+@tf_export('autograph.experimental.do_not_convert')
+def do_not_convert(func=None):
+  """Decorator that suppresses the conversion of a function.
+
+  Args:
+    func: function to decorate.
+
+  Returns:
+    If `func` is not None, returns a `Callable` which is equivalent to
+    `func`, but is not converted by AutoGraph.
+    If `func` is None, returns a decorator that, when invoked with a
+    single `func` argument, returns a `Callable` equivalent to the
+    above case.
+  """
+  if func is None:
+    return do_not_convert
+
+  def wrapper(*args, **kwargs):
+    with ag_ctx.ControlStatusCtx(status=ag_ctx.Status.DISABLED):
+      return func(*args, **kwargs)
+
+  if inspect.isfunction(func) or inspect.ismethod(func):
+    wrapper = functools.update_wrapper(wrapper, func)
+
+  return autograph_artifact(wrapper)
+
+
+# TODO(mdan): Make private.
+def convert(recursive=False,
+            optional_features=None,
+            user_requested=True,
+            conversion_ctx=ag_ctx.NullCtx()):
+  """Decorator that compiles a function to use TensorFlow ops.
+
+  The decorator is dynamic - it recompiles the target whenever the decorated
+  function is called. This means the parameter values are known at conversion.
+  It also means that repeated calls with different types of parameters will be
+  correctly processed.
+
+  Args:
+    recursive: bool, whether to recursively convert any functions or classes
+      that the converted function may use.
+    optional_features: converted.Feature, allows toggling optional or
+      experimental features. When set to None, only the core features are
+      enabled.
+    user_requested: bool, whether this is a function that the user explicitly
+      asked to be converted. See ConversionOptions.user_requested.
+    conversion_ctx: Optional ag_ctx.ControlStatusCtx, the Autograph context in
+      which `f` is used.
+
+  Returns:
+    Callable, a decorator that converts the given function into an equivalent
+    function that uses TensorFlow ops.
+  """
+
+  def decorator(f):
+    """Decorator implementation."""
+
+    def wrapper(*args, **kwargs):
+      """Wrapper that calls the converted version of f."""
+      options = converter.ConversionOptions(
+          recursive=recursive,
+          user_requested=user_requested,
+          optional_features=optional_features)
+      try:
+        with conversion_ctx:
+          return converted_call(f, args, kwargs, options=options)
+      except Exception as e:  # pylint:disable=broad-except
+        if hasattr(e, 'ag_error_metadata'):
+          raise e.ag_error_metadata.to_exception(e)
+        else:
+          raise
+
+    if inspect.isfunction(f) or inspect.ismethod(f):
+      wrapper = functools.update_wrapper(wrapper, f)
+
+    decorated_wrapper = tf_decorator.make_decorator(f, wrapper)
+    return autograph_artifact(decorated_wrapper)
+
+  return decorator
+
+
+# pylint:disable=line-too-long
+@tf_export('autograph.to_graph', v1=[])
+def to_graph(entity, recursive=True, experimental_optional_features=None):
+  """Converts a Python entity into a TensorFlow graph.
+
+  Also see: `tf.autograph.to_code`, `tf.function`.
+
+  Unlike `tf.function`, `to_graph` is a low-level transpiler that converts
+  Python code to TensorFlow graph code. It does not implement any caching,
+  variable management or create any actual ops, and is best used where greater
+  control over the generated TensorFlow graph is desired. Another difference
+  from `tf.function` is that `to_graph` will not wrap the graph into a
+  TensorFlow function or a Python callable. Internally, `tf.function` uses
+  `to_graph`.
+
+  Example usage:
+
+  >>> def f(x):
+  ...   if x > 0:
+  ...     y = x * x
+  ...   else:
+  ...     y = -x
+  ...   return y
+  ...
+  >>> converted_f = to_graph(f)
+  >>> x = tf.constant(2)
+  >>> converted_f(x)  # converted_foo is like a TensorFlow Op.
+  <tf.Tensor: shape=(), dtype=int32, numpy=4>
+
+  Supported Python entities include:
+    * functions
+    * classes
+    * object methods
+
+  Functions are converted into new functions with converted code.
+
+  Classes are converted by generating a new class whose methods use converted
+  code.
+
+  Methods are converted into unbound function that have an additional first
+  argument called `self`.
+
+  For a tutorial, see the
+  [tf.function and AutoGraph guide](https://www.tensorflow.org/guide/function).
+  For more detailed information, see the
+  [AutoGraph reference documentation](https://github.com/tensorflow/tensorflow/blob/master/tensorflow/python/autograph/g3doc/reference/index.md).
+
+  Args:
+    entity: Python callable or class to convert.
+    recursive: Whether to recursively convert any functions that the converted
+      function may call.
+    experimental_optional_features: `None`, a tuple of, or a single
+      `tf.autograph.experimental.Feature` value.
+
+  Returns:
+    Same as `entity`, the converted Python function or class.
+
+  Raises:
+    ValueError: If the entity could not be converted.
+  """
+  try:
+    program_ctx = converter.ProgramContext(
+        options=converter.ConversionOptions(
+            recursive=recursive,
+            user_requested=True,
+            optional_features=experimental_optional_features))
+    return autograph_artifact(_convert_actual(entity, program_ctx))
+  except (ValueError, AttributeError, KeyError, NameError, AssertionError) as e:
+    logging.error(1, 'Error converting %s', entity, exc_info=True)
+    raise ConversionError('converting {}: {}: {}'.format(
+        entity, e.__class__.__name__, str(e)))
+
+
+@tf_export(v1=['autograph.to_graph'])
+def to_graph_v1(entity,
+                recursive=True,
+                arg_values=None,
+                arg_types=None,
+                experimental_optional_features=None):
+  """Converts a Python entity into a TensorFlow graph.
+
+  Also see: `tf.autograph.to_code`, `tf.function`.
+
+  Unlike `tf.function`, `to_graph` is a low-level transpiler that converts
+  Python code to TensorFlow graph code. It does not implement any caching,
+  variable management or create any actual ops, and is best used where greater
+  control over the generated TensorFlow graph is desired. Another difference
+  from `tf.function` is that `to_graph` will not wrap the graph into a
+  TensorFlow function or a Python callable. Internally, `tf.function` uses
+  `to_graph`.
+
+  _Example Usage_
+
+  ```python
+    def foo(x):
+      if x > 0:
+        y = x * x
+      else:
+        y = -x
+      return y
+
+    converted_foo = to_graph(foo)
+
+    x = tf.constant(1)
+    y = converted_foo(x)  # converted_foo is a TensorFlow Op-like.
+    assert is_tensor(y)
+  ```
+
+  Supported Python entities include:
+    * functions
+    * classes
+    * object methods
+
+  Functions are converted into new functions with converted code.
+
+  Classes are converted by generating a new class whose methods use converted
+  code.
+
+  Methods are converted into unbound function that have an additional first
+  argument called `self`.
+
+  Args:
+    entity: Python callable or class to convert.
+    recursive: Whether to recursively convert any functions that the converted
+      function may call.
+    arg_values: Deprecated.
+    arg_types: Deprecated.
+    experimental_optional_features: `None`, a tuple of, or a single
+      `tf.autograph.experimental.Feature` value.
+
+  Returns:
+    Same as `entity`, the converted Python function or class.
+
+  Raises:
+    ValueError: If the entity could not be converted.
+  """
+  del arg_types
+  del arg_values
+  return to_graph(
+      entity,
+      recursive=recursive,
+      experimental_optional_features=experimental_optional_features)
+
+
+@tf_export(v1=['autograph.to_code'])
+def to_code_v1(entity,
+               recursive=True,
+               arg_values=None,
+               arg_types=None,
+               indentation='  ',
+               experimental_optional_features=None):
+  """Returns the source code generated by AutoGraph, as a string.
+
+  Example usage:
+
+  >>> def f(x):
+  ...   if x < 0:
+  ...     x = -x
+  ...   return x
+  >>> tf.autograph.to_code(f)
+  "...def tf__f(x):..."
+
+  Also see: `tf.autograph.to_graph`.
+
+  Note: If a function has been decorated with `tf.function`, pass its
+  underlying Python function, rather than the callable that `tf.function
+  creates:
+
+  >>> @tf.function
+  ... def f(x):
+  ...   if x < 0:
+  ...     x = -x
+  ...   return x
+  >>> tf.autograph.to_code(f.python_function)
+  "...def tf__f(x):..."
+
+  Args:
+    entity: Python callable or class.
+    recursive: Whether to recursively convert any functions that the converted
+      function may call.
+    arg_values: Deprecated.
+    arg_types: Deprecated.
+    indentation: Deprecated.
+    experimental_optional_features: `None`, a tuple of, or a single
+      `tf.autograph.experimental.Feature` value.
+
+  Returns:
+    The converted code as string.
+  """
+  del arg_values
+  del arg_types
+  del indentation
+  return to_code(
+      entity,
+      recursive=recursive,
+      experimental_optional_features=experimental_optional_features)
+
+
+@tf_export('autograph.to_code', v1=[])
+def to_code(entity, recursive=True, experimental_optional_features=None):
+  """Returns the source code generated by AutoGraph, as a string.
+
+  Example usage:
+
+  >>> def f(x):
+  ...   if x < 0:
+  ...     x = -x
+  ...   return x
+  >>> tf.autograph.to_code(f)
+  "...def tf__f(x):..."
+
+  Also see: `tf.autograph.to_graph`.
+
+  Note: If a function has been decorated with `tf.function`, pass its
+  underlying Python function, rather than the callable that `tf.function
+  creates:
+
+  >>> @tf.function
+  ... def f(x):
+  ...   if x < 0:
+  ...     x = -x
+  ...   return x
+  >>> tf.autograph.to_code(f.python_function)
+  "...def tf__f(x):..."
+
+  Args:
+    entity: Python callable or class to convert.
+    recursive: Whether to recursively convert any functions that the converted
+      function may call.
+    experimental_optional_features: `None`, a tuple of, or a single
+      `tf.autograph.experimental.Feature` value.
+
+  Returns:
+    The converted code as string.
+  """
+  source = tf_inspect.getsource(
+      to_graph(
+          entity,
+          recursive=recursive,
+          experimental_optional_features=experimental_optional_features))
+  return textwrap.dedent(source)
+
+
+_TRANSPILER = PyToTF()

tensorflow/python/autograph/impl/conversion.py

@@ -0,0 +1,227 @@
+# Copyright 2016 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.
+# ==============================================================================
+"""Core conversion logic, serves as main point of access."""
+
+import functools
+import inspect
+import sys
+import unittest
+
+from tensorflow.python.autograph.core import config
+from tensorflow.python.autograph.pyct import cache
+from tensorflow.python.autograph.pyct import inspect_utils
+from tensorflow.python.autograph.utils import ag_logging as logging
+from tensorflow.python.eager.polymorphic_function import tf_method_target
+from tensorflow.python.util import tf_inspect
+
+
+_ALLOWLIST_CACHE = cache.UnboundInstanceCache()
+
+
+def _is_of_known_loaded_module(f, module_name):
+  mod = sys.modules.get(module_name, None)
+  if mod is None:
+    return False
+  if any(v is not None for v in mod.__dict__.values() if f is v):
+    return True
+  return False
+
+
+def _is_known_loaded_type(f, module_name, entity_name):
+  """Tests whether the function or method is an instance of a known type."""
+  if (module_name not in sys.modules or
+      not hasattr(sys.modules[module_name], entity_name)):
+    return False
+  type_entity = getattr(sys.modules[module_name], entity_name)
+  if isinstance(f, type_entity):
+    # The method if of this type. Example:
+    #
+    # o = ClassType()
+    # function(o.method)()
+    return True
+  # Note: inspect is required here, to avoid unpacking tf.function decorators.
+  if inspect.ismethod(f):
+    # The unbound method if of this type. Example:
+    #
+    # class ClassType:
+    #   @function
+    #   def method(self):
+    #     ...
+    # o = ClassType()
+    # o.method()
+    if isinstance(f.__func__, type_entity):
+      return True
+  return False
+
+
+def is_unsupported(o):
+  """Checks whether an entity is supported by AutoGraph at all."""
+
+  # TODO(b/122265385): Remove this bypass.
+  if (_is_known_loaded_type(o, 'wrapt', 'FunctionWrapper') or
+      _is_known_loaded_type(o, 'wrapt', 'BoundFunctionWrapper')):
+    logging.warning(
+        '{} appears to be decorated by wrapt, which is not yet supported'
+        ' by AutoGraph. The function will run as-is.'
+        ' You may still apply AutoGraph before the wrapt decorator.'.format(o))
+    logging.log(2, 'Permanently allowed: %s: wrapt decorated', o)
+    return True
+
+  if _is_known_loaded_type(o, 'functools', '_lru_cache_wrapper'):
+    logging.log(2, 'Permanently allowed: %s: lru_cache', o)
+    return True
+
+  # Constructors are permanently allowed.
+  # TODO(mdan): Toggle as experimental feature instead.
+  # TODO(b/124016764): Remove this limitation.
+  if inspect_utils.isconstructor(o):
+    logging.log(2, 'Permanently allowed: %s: constructor', o)
+    return True
+
+  # Other built-in modules are permanently allowed.
+  # TODO(mdan): Figure out how to do this consistently for all stdlib modules.
+  if any(
+      _is_of_known_loaded_module(o, m)
+      for m in ('collections', 'pdb', 'copy', 'inspect', 're')):
+    logging.log(2, 'Permanently allowed: %s: part of builtin module', o)
+    return True
+
+  # Custom ops and kernels are also permanently allowed.
+  # See tensorflow.framework.load_library.
+  if (hasattr(o, '__module__') and
+      hasattr(o.__module__, '_IS_TENSORFLOW_PLUGIN')):
+    logging.log(2, 'Permanently allowed: %s: TensorFlow plugin', o)
+    return True
+
+  return False
+
+
+# TODO(mdan): allow_namedtuple_subclass should be hardcoded to True.
+def is_allowlisted(
+    o, check_call_override=True, allow_namedtuple_subclass=False):
+  """Checks whether an entity is allowed for use in graph mode.
+
+  Examples of allowed entities include all members of the tensorflow
+  package.
+
+  Args:
+    o: A Python entity.
+    check_call_override: Reserved for internal use. When set to `False`, it
+      disables the rule according to which classes are allowed if their
+      __call__ method is allowed.
+    allow_namedtuple_subclass: Reserved for internal use. When `True`,
+      namedtuple subclasses are not allowed.
+
+  Returns:
+    Boolean
+  """
+  # TODO(b/120224672): Fix this.
+  if isinstance(o, functools.partial):
+    # tf_inspect.getmodule(functools.partial(...)) otherwise returns None since
+    # functools.partial objects do not have a __module__ attribute.
+    m = functools
+  else:
+    m = tf_inspect.getmodule(o)
+
+  # Examples of callables that lack a __module__ property include builtins.
+  if hasattr(m, '__name__'):
+    for rule in config.CONVERSION_RULES:
+      action = rule.get_action(m)
+      if action == config.Action.CONVERT:
+        logging.log(2, 'Not allowed: %s: %s', o, rule)
+        return False
+      elif action == config.Action.DO_NOT_CONVERT:
+        logging.log(2, 'Allowlisted: %s: %s', o, rule)
+        return True
+
+  # The check for __code__ below is because isgeneratorfunction crashes
+  # without one.
+  if hasattr(o, '__code__') and tf_inspect.isgeneratorfunction(o):
+    logging.log(2, 'Allowlisted: %s: generator functions are not converted', o)
+    return True
+
+  if (check_call_override and not tf_inspect.isclass(o) and
+      hasattr(o, '__call__')):
+    # Callable objects: allowed if their __call__ method is.
+    # The type check avoids infinite recursion around the __call__ method
+    # of function objects.
+    if (type(o) != type(o.__call__)) and is_allowlisted(o.__call__):  # pylint: disable=unidiomatic-typecheck
+      logging.log(2, 'Allowlisted: %s: object __call__ allowed', o)
+      return True
+
+  owner_class = None
+  if tf_inspect.ismethod(o):
+    # Methods of allowed classes are also allowed, even if they are
+    # bound via user subclasses.
+    #
+    # For example, suppose `tf.Foo` has a method called `bar`, and `baz` is
+    # defined as below. `tf.Foo` is allowed. Then `baz.bar` is also
+    # allowed.
+    #
+    #   class Custom(tf.Foo):
+    #     pass
+    #
+    #   baz = Custom()
+    #
+    # For the example above, if `Custom` did overload `bar`, then it would no
+    # longer be allowed.
+
+    owner_class = inspect_utils.getmethodclass(o)
+    if owner_class is tf_method_target.TfMethodTarget:
+      owner_class = o.__self__.target_class
+    if owner_class is not None:
+      if issubclass(owner_class, unittest.TestCase):
+        logging.log(2, 'Allowlisted: %s: method of TestCase subclass', o)
+        return True
+
+      owner_class = inspect_utils.getdefiningclass(o, owner_class)
+      if is_allowlisted(
+          owner_class,
+          check_call_override=False,
+          allow_namedtuple_subclass=True):
+        logging.log(2, 'Allowlisted: %s: owner is allowed %s', o,
+                    owner_class)
+        return True
+
+  if inspect_utils.isnamedtuple(o):
+    # Due to the way they're constructed, namedtuple types cannot be converted
+    # because they don't expose source code. But we assume they are safe for
+    # graph mode since they are just containers.
+    if allow_namedtuple_subclass:
+      if not any(inspect_utils.isnamedtuple(base) for base in o.__bases__):
+        logging.log(2, 'Allowlisted: %s: named tuple', o)
+        return True
+    else:
+      logging.log(2, 'Allowlisted: %s: named tuple or subclass', o)
+      return True
+
+  logging.log(2, 'Not allowed: %s: default rule', o)
+  return False
+
+
+def is_in_allowlist_cache(entity, options):
+  try:
+    return _ALLOWLIST_CACHE.has(entity, options)
+  except TypeError:
+    # Catch-all for entities that are unhashable or don't allow weakrefs.
+    return False
+
+
+def cache_allowlisted(entity, options):
+  try:
+    _ALLOWLIST_CACHE[entity][options] = True
+  except TypeError:
+    # Catch-all for entities that are unhashable or don't allow weakrefs.
+    pass

tensorflow/python/autograph/impl/conversion_test.py

@@ -0,0 +1,123 @@
+# Copyright 2017 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 conversion module."""
+
+import sys
+import types
+import weakref
+
+from tensorflow.python.autograph import utils
+from tensorflow.python.autograph.core import config
+from tensorflow.python.autograph.core import converter
+from tensorflow.python.autograph.impl import api
+from tensorflow.python.autograph.impl import conversion
+from tensorflow.python.autograph.impl.testing import pybind_for_testing
+from tensorflow.python.eager import function
+from tensorflow.python.framework import constant_op
+from tensorflow.python.platform import test
+
+
+class ConversionTest(test.TestCase):
+
+  def _simple_program_ctx(self):
+    return converter.ProgramContext(
+        options=converter.ConversionOptions(recursive=True),
+        autograph_module=api)
+
+  def test_is_allowlisted(self):
+
+    def test_fn():
+      return constant_op.constant(1)
+
+    self.assertFalse(conversion.is_allowlisted(test_fn))
+    self.assertTrue(conversion.is_allowlisted(utils))
+    self.assertTrue(conversion.is_allowlisted(constant_op.constant))
+
+  def test_is_allowlisted_tensorflow_like(self):
+
+    tf_like = types.ModuleType('tensorflow_foo')
+    def test_fn():
+      pass
+    tf_like.test_fn = test_fn
+    test_fn.__module__ = tf_like
+
+    self.assertFalse(conversion.is_allowlisted(tf_like.test_fn))
+
+  def test_is_allowlisted_callable_allowlisted_call(self):
+
+    allowlisted_mod = types.ModuleType('test_allowlisted_call')
+    sys.modules['test_allowlisted_call'] = allowlisted_mod
+    config.CONVERSION_RULES = ((config.DoNotConvert('test_allowlisted_call'),) +
+                               config.CONVERSION_RULES)
+
+    class TestClass:
+
+      def __call__(self):
+        pass
+
+      def allowlisted_method(self):
+        pass
+
+    TestClass.__module__ = 'test_allowlisted_call'
+    TestClass.__call__.__module__ = 'test_allowlisted_call'
+
+    class Subclass(TestClass):
+
+      def converted_method(self):
+        pass
+
+    tc = Subclass()
+
+    self.assertTrue(conversion.is_allowlisted(TestClass.__call__))
+    self.assertTrue(conversion.is_allowlisted(tc))
+    self.assertTrue(conversion.is_allowlisted(tc.__call__))
+    self.assertTrue(conversion.is_allowlisted(tc.allowlisted_method))
+    self.assertFalse(conversion.is_allowlisted(Subclass))
+    self.assertFalse(conversion.is_allowlisted(tc.converted_method))
+
+  def test_is_allowlisted_tfmethodwrapper(self):
+    allowlisted_mod = types.ModuleType('test_allowlisted_call')
+    sys.modules['test_allowlisted_call'] = allowlisted_mod
+    config.CONVERSION_RULES = ((config.DoNotConvert('test_allowlisted_call'),) +
+                               config.CONVERSION_RULES)
+
+    class TestClass:
+
+      def member_function(self):
+        pass
+
+    TestClass.__module__ = 'test_allowlisted_call'
+    test_obj = TestClass()
+
+    def test_fn(self):
+      del self
+
+    bound_method = types.MethodType(
+        test_fn,
+        function.TfMethodTarget(
+            weakref.ref(test_obj), test_obj.member_function))
+
+    self.assertTrue(conversion.is_allowlisted(bound_method))
+
+  def test_is_allowlisted_pybind(self):
+    test_object = pybind_for_testing.TestClassDef()
+    with test.mock.patch.object(config, 'CONVERSION_RULES', ()):
+      # TODO(mdan): This should return True for functions and methods.
+      # Note: currently, native bindings are allowlisted by a separate check.
+      self.assertFalse(conversion.is_allowlisted(test_object.method))
+
+
+if __name__ == '__main__':
+  test.main()

tensorflow/python/autograph/impl/testing/BUILD

@@ -0,0 +1,19 @@
+load("//tensorflow:tensorflow.default.bzl", "pybind_extension")
+
+package(
+    # copybara:uncomment default_applicable_licenses = ["//tensorflow:license"],
+    default_visibility = ["//tensorflow:__subpackages__"],
+    licenses = ["notice"],
+)
+
+pybind_extension(
+    name = "pybind_for_testing",
+    srcs = ["pybind_for_testing.cc"],
+    enable_stub_generation = True,
+    pytype_srcs = [
+        "pybind_for_testing.pyi",
+    ],
+    deps = [
+        "@pybind11",
+    ],
+)

tensorflow/python/autograph/impl/testing/pybind_for_testing.cc

@@ -0,0 +1,39 @@
+// Copyright 2017 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.
+// ==============================================================================
+
+#include "pybind11/pybind11.h"  // from @pybind11
+#include "pybind11/pytypes.h"  // from @pybind11
+#include "pybind11/stl.h"  // from @pybind11
+
+namespace autograph {
+
+namespace py = pybind11;
+
+class TestClassDef {
+ public:
+  TestClassDef() = default;
+
+  py::object Method() const;
+};
+
+py::object TestClassDef::Method() const { return py::none(); }
+
+PYBIND11_MODULE(pybind_for_testing, m) {
+  py::class_<TestClassDef>(m, "TestClassDef")
+      .def(py::init<>())
+      .def("method", &TestClassDef::Method);
+}
+
+}  // namespace autograph

tensorflow/python/autograph/impl/testing/pybind_for_testing.pyi

@@ -0,0 +1,18 @@
+# Copyright 2023 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.
+# ==============================================================================
+
+class TestClassDef:
+    def __init__(self) -> None: ...
+    def method(self) -> object: ...